Update app.py

app.py

@@ -16,7 +16,8 @@ if importlib.util.find_spec("openvino"):
     def predict_openvino(
             image: Union[str, Path, np.ndarray],
             model_path: Union[str, Path],
-            device: str) -> Dict[str, np.ndarray]:
+            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(
@@ -24,6 +25,7 @@ if importlib.util.find_spec("openvino"):
                 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):
@@ -34,27 +36,31 @@ def convert_to_heatmap(heatmap):
     heatmap = cv2.applyColorMap((heatmap*255).astype("uint8"), cv2.COLORMAP_HSV)
     return heatmap
 
-def predict(image, device, model_dir):
-    outputs = predict_openvino(image, model_dir, device)
+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
 
 def launch():
-    input_image = gr.Image(label="Input image")
+    input_image = gr.Image(
+        label="Input image",
+        value="images/171436008_Fail.jpeg"
+        )
+    threshold = gr.Slider(value=5.5, step=0.1, label="Threshold")
     devices = gr.Radio(
         label="Device",
-        choices=["AUTO", "CPU", "CUDA"],
+        choices=["AUTO", "CPU", "GPU"],
         value="CPU",
         interactive=False
     )
     output_image = gr.Image(label="Output image")
     output_heatmap = gr.Image(label="Heatmap")
-    model = gr.Text(label="Model", interactive=False, value="models/default")
+    model = gr.Text(label="Model", interactive=False, value="models/glass-tr-5-ov")
     intf = gr.Interface(
         title="Anomaly Detection",
         fn=predict,
-        inputs=[input_image, devices, model],
+        inputs=[input_image, threshold, devices, model],
         outputs=[output_image, output_heatmap]
     )
     intf.launch()