Update app.py

app.py

@@ -4,7 +4,6 @@ from constants import MODELS_REPO, MODELS_NAMES
 import gradio as gr
 import torch
 
-
 from transformers import (AutoFeatureExtractor, DetrForObjectDetection,)
 from visualization import visualize_attention_map, visualize_prediction
 from style import css, description, title
@@ -12,6 +11,7 @@ from style import css, description, title
 from PIL import Image
 
 
+
 def make_prediction(img, feature_extractor, model):
     inputs = feature_extractor(img, return_tensors="pt")
     outputs = model(**inputs)
@@ -29,8 +29,6 @@ def make_prediction(img, feature_extractor, model):
 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"
@@ -80,41 +78,43 @@ with gr.Blocks(css=css) as app:
         with gr.TabItem("Image upload and detections visualization"):
             with gr.Row():
                 with gr.Column():
-                    img_input = gr.Image(type="pil")
+                    with gr.Row():
+                        img_input = gr.Image(type="pil")
                     img_input_mask = gr.Image(type="pil", visible=False)
+                    with gr.Row():
+                        example_images = gr.Dataset(
+                            components=[img_input, img_input_mask],
+                            samples=[
+                                [path.as_posix(), path.as_posix().replace("_HE", "_mask")]
+                                for path in sorted(
+                                    pathlib.Path("cd45rb_test_imgs").rglob("*_HE.png")
+                                )
+                            ],
+                            samples_per_page=2,
+                        )
                 with gr.Column():
-                    example_images = gr.Dataset(
-                        components=[img_input, img_input_mask],
-                        samples=[
-                            [path.as_posix(), path.as_posix().replace("_HE", "_mask")]
-                            for path in sorted(
-                                pathlib.Path("cd45rb_test_imgs").rglob("*_HE.png")
-                            )
-                        ],
-                        samples_per_page=2,
-                    )
-            with gr.Row():
-                with gr.Column():
-                    options = gr.Dropdown(
-                        value=MODELS_NAMES[0],
-                        choices=MODELS_NAMES,
-                        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.Column():
-                    display_mask = gr.Checkbox(
-                        label="Display masks", default=False
-                    )
-                    detect_button = gr.Button("Detect leukocytes")
+                    with gr.Row():
+                        options = gr.Dropdown(
+                            value=MODELS_NAMES[0],
+                            choices=MODELS_NAMES,
+                            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():
+                        display_mask = gr.Checkbox(
+                            label="Display masks", default=False
+                        )
+                    with gr.Row():
+                        detect_button = gr.Button("Detect leukocytes")
             with gr.Row():
                 with gr.Column():
                     gr.Markdown(
                         """The selected image with detected bounding boxes by the model"""
                     )
-                    img_output_from_upload = gr.Image(shape=(850, 850))
+                    img_output_from_upload = gr.Image(shape=(800, 800))
         with gr.TabItem("Attentions visualization"):
             gr.Markdown("""Encoder attentions""")
             with gr.Row():