Update app.py

app.py

@@ -9,6 +9,8 @@ from transformers import (AutoFeatureExtractor, DetrForObjectDetection,)
 from visualization import visualize_attention_map, visualize_prediction
 from style import css, description, title
 
+from PIL import Image
+
 
 def make_prediction(img, feature_extractor, model):
     inputs = feature_extractor(img, return_tensors="pt")
@@ -16,14 +18,6 @@ def make_prediction(img, feature_extractor, model):
     img_size = torch.tensor([tuple(reversed(img.size))])
     processed_outputs = feature_extractor.post_process(outputs, img_size)
     print(outputs.keys())
-    # if model type is YOLOS, then return "attentions"
-    if "attentions" in outputs.keys():
-        return (
-            processed_outputs[0],
-            outputs["attentions"],
-            outputs["attentions"],
-            outputs["attentions"],
-        )
     return (
         processed_outputs[0],
         outputs["decoder_attentions"],
@@ -32,7 +26,7 @@ def make_prediction(img, feature_extractor, model):
     )
 
 
-def detect_objects(model_name, image_input, threshold):
+def detect_objects(model_name, image_input, threshold, display_mask=False):
     feature_extractor = AutoFeatureExtractor.from_pretrained(MODELS_REPO[model_name])
 
     if "DETR" in model_name:
@@ -46,8 +40,17 @@ def detect_objects(model_name, image_input, threshold):
         cross_attention_map,
     ) = make_prediction(image_input, feature_extractor, model)
 
+    mask_pil_image = None
+    if display_mask:
+        # get image path
+        image_path = pathlib.Path(image_input.name)
+        mask_path = image_path.parent / (
+            image_path.stem.replace("_HE", "_mask") + image_path.suffix
+        )
+        mask_pil_image = Image.open(mask_path)
+
     viz_img = visualize_prediction(
-        image_input, processed_outputs, threshold, model.config.id2label
+        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
@@ -86,9 +89,15 @@ with gr.Blocks(css=css) as app:
                         label="Select an object detection model",
                         show_label=True,
                     )
-                    slider_input = gr.Slider(
-                        minimum=0.2, maximum=1, value=0.7, label="Prediction threshold"
-                    )
+                    with gr.Row():
+                        with gr.Column():
+                            slider_input = gr.Slider(
+                                minimum=0.2, maximum=1, value=0.7, label="Prediction threshold"
+                            )
+                        with gr.Column():
+                            display_mask = gr.Checkbox(
+                                label="Display masks", default=False
+                            )
                     detect_button = gr.Button("Detect leukocytes")
             with gr.Row():
                 example_images = gr.Dataset(
@@ -96,7 +105,7 @@ with gr.Blocks(css=css) as app:
                     samples=[
                         [path.as_posix()]
                         for path in sorted(
-                            pathlib.Path("cd45rb_test_imgs").rglob("*.png")
+                            pathlib.Path("cd45rb_test_imgs").rglob("*_HE.png")
                         )
                     ],
                 )
@@ -106,7 +115,7 @@ with gr.Blocks(css=css) as app:
                         """The selected image with detected bounding boxes by the model"""
                     )
                     img_output_from_upload = gr.Image(shape=(850, 850))
-        with gr.TabItem("Attention maps visualization"):
+        with gr.TabItem("Attentions visualization"):
             gr.Markdown("""Encoder attentions""")
             with gr.Row():
                 encoder_att_map_output = gr.Image(shape=(850, 850))
@@ -122,7 +131,7 @@ with gr.Blocks(css=css) as app:
 
     detect_button.click(
         detect_objects,
-        inputs=[options, img_input, slider_input],
+        inputs=[options, img_input, slider_input, display_mask],
         outputs=[
             img_output_from_upload,
             decoder_att_map_output,
@@ -137,3 +146,4 @@ with gr.Blocks(css=css) as app:
     )
 
     app.launch(enable_queue=True)
+