Update app.py

app.py

@@ -26,9 +26,11 @@ def make_prediction(img, feature_extractor, model):
     )
 
 
-def detect_objects(model_name, image_input, image_path, threshold, display_mask=False):
+def detect_objects(model_name, image_input, threshold, display_mask=False, img_input_mask=None):
     feature_extractor = AutoFeatureExtractor.from_pretrained(MODELS_REPO[model_name])
 
+    print("Set threshold to: ", threshold)
+
     if "DETR" in model_name:
         model = DetrForObjectDetection.from_pretrained(MODELS_REPO[model_name])
         model_details = "DETR details"
@@ -40,15 +42,16 @@ def detect_objects(model_name, image_input, image_path, threshold, display_mask=
         cross_attention_map,
     ) = make_prediction(image_input, feature_extractor, model)
 
-    mask_pil_image = None
     if display_mask:
-        mask_path = image_path.parent / (
-            image_path.stem.replace("_HE", "_mask") + image_path.suffix
-        )
-        mask_pil_image = Image.open(mask_path)
+        mask_pil_image = img_input_mask
 
     viz_img = visualize_prediction(
-        image_input, processed_outputs, threshold, model.config.id2label, display_mask, mask_pil_image
+        image_input,
+        processed_outputs,
+        threshold,
+        model.config.id2label,
+        display_mask,
+        mask_pil_image
     )
     decoder_attention_map_img = visualize_attention_map(
         image_input, decoder_attention_map
@@ -130,7 +133,7 @@ with gr.Blocks(css=css) as app:
 
     detect_button.click(
         detect_objects,
-        inputs=[options, img_input, slider_input, display_mask],
+        inputs=[options, img_input, img_input_mask, slider_input, display_mask],
         outputs=[
             img_output_from_upload,
             decoder_att_map_output,