Update app.py

app.py

@@ -10,7 +10,7 @@ from fastai.vision.all import *
 model_multi = load_learner('vit_tiny_patch16.pkl')
 
 def binary_label(path):
-  return 'No-anomaly' if (parent_label(path) == 'No-Anomaly') else 'Anomaly'
+    return 'No-anomaly' if (parent_label(path) == 'No-Anomaly') else 'Anomaly'
 
 model_binary = load_learner('vit_tiny_patch16_binary.pkl')
 
@@ -22,17 +22,17 @@ seg_feature_extractor = SegformerFeatureExtractor.from_pretrained('zklee98/segfo
 seg_model = SegformerForSemanticSegmentation.from_pretrained('zklee98/segformer-b1-solarModuleAnomaly-v0.1')
 
 def get_seg_overlay(image, seg):
-  color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
-  palette = np.array(sidewalk_palette())
-  for label, color in enumerate(palette):
-      color_seg[seg == label, :] = color
+    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
+    palette = np.array(sidewalk_palette())
+    for label, color in enumerate(palette):
+        color_seg[seg == label, :] = color
 
-  # Show image + mask
-  img = np.array(image) * 0.5 + color_seg * 0.5
-  img = img.astype(np.uint8)
-  #img = PIL.Image.open(img)
+    # Show image + mask
+    img = np.array(image) * 0.5 + color_seg * 0.5
+    img = img.astype(np.uint8)
+    #img = PIL.Image.open(img)
 
-  return img
+    return img
 
 #@title `def sidewalk_palette()`
 
@@ -80,46 +80,45 @@ def sidewalk_palette():
 
 def predict(classification_mode, image):
 
-  if (classification_mode == 'Binary Classification'):
-    model = model_binary
-  else:
-    model = model_multi
+    if (classification_mode == 'Binary Classification'):
+        model = model_binary
+    else:
+        model = model_multi
 
-  labels = model.dls.vocab
-  # Classification model prediction
-  image  = PILImage.create(image)
-  pred, pred_idx, probs = model.predict(image)
+    labels = model.dls.vocab
+    # Classification model prediction
+    #image  = PILImage.create(image)
+    pred, pred_idx, probs = model.predict(image)
 
-  seg_img = None
-  percentage_affected = '0%'
-  if (pred.upper() != 'NO-ANOMALY'):
-    addChannel = Grayscale(num_output_channels=3)
-    image = addChannel(image)
+    seg_img = None
+    percentage_affected = '0%'
+    if (pred.upper() != 'NO-ANOMALY'):
+        addChannel = Grayscale(num_output_channels=3)
+        image = addChannel(image)
 
-    inputs = seg_feature_extractor(images=image, return_tensors="pt")
-    outputs = seg_model(**inputs)
-    logits = outputs.logits  # shape (batch_size, num_labels, height/4, width/4)
+        inputs = seg_feature_extractor(images=image, return_tensors="pt")
+        outputs = seg_model(**inputs)
+        logits = outputs.logits  # shape (batch_size, num_labels, height/4, width/4)
 
-    # First, rescale logits to original image size
-    upsampled_logits = nn.functional.interpolate(
-        logits,
-        size=image.size[::-1], # (height, width)
-        mode='bilinear',
-        align_corners=False
-    )
+        # First, rescale logits to original image size
+        upsampled_logits = nn.functional.interpolate(
+            logits,
+            size=image.size[::-1], # (height, width)
+            mode='bilinear',
+            align_corners=False)
 
-    # Second, apply argmax on the class dimension
-    pred_seg = upsampled_logits.argmax(dim=1)[0]
+        # Second, apply argmax on the class dimension
+        pred_seg = upsampled_logits.argmax(dim=1)[0]
 
-    seg_img = get_seg_overlay(image, pred_seg)
+        seg_img = get_seg_overlay(image, pred_seg)
 
-    classified_pixels = np.unique(pred_seg.numpy(), return_counts=True)
-    pixels_count = dict({classified_pixels[0][0]: classified_pixels[1][0],
-                                classified_pixels[0][1]: classified_pixels[1][1]})
-    percentage_affected = round((pixels_count[1]/960)*100, 1)
-    percentage_affected = str(percentage_affected) + '%'
+        classified_pixels = np.unique(pred_seg.numpy(), return_counts=True)
+        pixels_count = dict({classified_pixels[0][0]: classified_pixels[1][0],
+                             classified_pixels[0][1]: classified_pixels[1][1]})
+        percentage_affected = round((pixels_count[1]/960)*100, 1)
+        percentage_affected = str(percentage_affected) + '%'
 
-    seg_img = PIL.Image.fromarray(seg_img)
+        seg_img = PIL.Image.fromarray(seg_img)
 
   return ({labels[i]: float(probs[i]) for i in range(len(labels))}, seg_img, percentage_affected)
 
@@ -133,7 +132,7 @@ description = """
 gr.Interface(fn=predict,
              inputs= [gr.Dropdown(choices=['Binary Classification', 'Multiclass Classification'], label='Classification Mode:', 
                                   info='Choose to classify between anomaly and no-anomaly OR between 12 different types of anomalies.'),
-                      gr.Image(label='Input infrared image: ')],
+                      gr.Image(type='pil', label='Input infrared image: ')],
              outputs=[gr.outputs.Label(num_top_classes=3, label='Detected:').style(container=False),
                       gr.Image(type='pil', label=' ').style(height=240, width=144),
                       gr.Textbox(label='Affected area:').style(container=False)],