Update interface

app.py

@@ -1,67 +1,39 @@
 from typing import *
-import importlib
-from pathlib import Path
 
-import cv2
 import gradio as gr
-import numpy as np
 
-available_modes = []
-inferencer = None
-last_model_path: str = None
+from predict import predict_fn
+from utils import populate_examples
 
-if importlib.util.find_spec("openvino"):
-    available_modes.append("OpenVINO")
-    from inference.openvino import OpenVINOInferencer
-    def predict_openvino(
-            image: Union[str, Path, np.ndarray],
-            model_path: Union[str, Path],
-            device: str,
-            threshold: float) -> Dict[str, np.ndarray]:
-        global inferencer, last_model_path
-        if not isinstance(inferencer, OpenVINOInferencer) or last_model_path != str(model_path):
-            inferencer = OpenVINOInferencer(
-                path = model_path,
-                device=device.upper()
-            )
-            last_model_path = str(model_path)
-        inferencer.metadata["pixel_threshold"] = threshold
-        return inferencer(image)
+description = """
+    Anomaly detection models are trained with only <span style="color:lime;font-weight:bold">normal</span> images,
+    and aimed to segment <span style="color:red;font-weight:bold">anomalies (deviations)</span> in input images.
 
-def draw_contour(image, mask, color=(255,0,0), thickness=3):
-    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
-    return cv2.drawContours(image.copy(), contours, -1, color, thickness)
-
-def convert_to_heatmap(heatmap):
-    heatmap = cv2.applyColorMap((heatmap*255).astype("uint8"), cv2.COLORMAP_HSV)
-    return heatmap
-
-def predict(image, threshold, device, model_dir):
-    outputs = predict_openvino(image, model_dir, device, threshold)
-    out_image = draw_contour(image, outputs["pred_mask"])
-    heatmap = convert_to_heatmap(outputs["anomaly_map"])
-    return out_image, heatmap
+    Scroll to bottom of this demo for a list of pretrained examples.
+"""
 
 def launch():
-    input_image = gr.Image(
-        label="Input image",
-        value="images/171436008_Fail.jpeg"
-        )
-    threshold = gr.Slider(value=5.5, step=0.1, label="Threshold")
+    input_image = gr.Image(label="Input image")
+    threshold = gr.Slider(value=1, step=0.1, label="Threshold")
     devices = gr.Radio(
         label="Device",
         choices=["AUTO", "CPU", "GPU"],
         value="CPU",
         interactive=False
     )
+    model = gr.Text(label="Model", interactive=False)
+
     output_image = gr.Image(label="Output image")
     output_heatmap = gr.Image(label="Heatmap")
-    model = gr.Text(label="Model", interactive=False, value="models/glass-tr-5-ov")
+
     intf = gr.Interface(
         title="Anomaly Detection",
-        fn=predict,
+        description=description,
+        fn=predict_fn,
         inputs=[input_image, threshold, devices, model],
-        outputs=[output_image, output_heatmap]
+        outputs=[output_image, output_heatmap],
+        examples=populate_examples(),
+        allow_flagging="never"
     )
     intf.launch()
 

examples.py

@@ -0,0 +1,10 @@
+"""List of example inputs, a dictionary where
+key is model_name, and value is a list of (image_name, threshold).
+"""
+
+examples = {
+    "glass-tr-5-ov":
+        [
+            ("171436008_Fail.jpeg", 5.5),
+        ],
+}

predict.py

@@ -0,0 +1,38 @@
+from typing import *
+from pathlib import Path
+
+import cv2
+from inference.openvino import OpenVINOInferencer
+import numpy as np
+
+inferencer = None
+last_model_path: str = None
+
+def predict_openvino(
+        image: Union[str, Path, np.ndarray],
+        model_path: Union[str, Path],
+        device: str,
+        threshold: float) -> Dict[str, np.ndarray]:
+    global inferencer, last_model_path
+    if not isinstance(inferencer, OpenVINOInferencer) or last_model_path != str(model_path):
+        inferencer = OpenVINOInferencer(
+            path = model_path,
+            device=device.upper()
+        )
+        last_model_path = str(model_path)
+    inferencer.metadata["pixel_threshold"] = threshold
+    return inferencer(image)
+
+def draw_contour(image, mask, color=(255,0,0), thickness=3):
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+    return cv2.drawContours(image.copy(), contours, -1, color, thickness)
+
+def convert_to_heatmap(heatmap):
+    heatmap = cv2.applyColorMap((heatmap*255).astype("uint8"), cv2.COLORMAP_HSV)
+    return heatmap
+
+def predict_fn(image, threshold, device, model_dir):
+    outputs = predict_openvino(image, model_dir, device, threshold)
+    out_image = draw_contour(image, outputs["pred_mask"])
+    heatmap = convert_to_heatmap(outputs["anomaly_map"])
+    return out_image, heatmap

utils.py

@@ -0,0 +1,29 @@
+from typing import *
+import os
+
+from examples import examples as exp_dict
+
+MODEL_DIR = "models"
+IMAGE_DIR = "images"
+
+def get_models() -> List[str]:
+    """List all model names in MODEL_DIR."""
+    model_names = []
+    for p in os.listdir(MODEL_DIR):
+        if os.path.isdir(os.path.join(MODEL_DIR, p)):
+            model_names.append(p)
+    return model_names
+
+def populate_examples() -> List[Any]:
+    """Populate example inputs into a list."""
+    examples = []
+    for m in get_models():
+        if m in exp_dict:
+            for img, thres in exp_dict[m]:
+                examples.append([
+                    f"{IMAGE_DIR}/{m}/{img}",
+                    thres,
+                    "CPU",
+                    f"{MODEL_DIR}/{m}"
+                ])
+    return examples