reorder gr blocks

app.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Tuple
 
 import gradio as gr
 import spaces
@@ -26,16 +26,15 @@ It can be used as an extention to Segment-Anything Model (SAM) or used as a stan
 """
 EXAMPLES = [
     [
-        "tiny",
-        0.5,
         "https://raw.githubusercontent.com/AnasIbrahim/image_agnostic_segmentation/master/demo/rgb_images/000000.png",
         "https://raw.githubusercontent.com/AnasIbrahim/image_agnostic_segmentation/CASE_release/demo/objects_gallery/obj_000001/1.jpg",
         "https://raw.githubusercontent.com/AnasIbrahim/image_agnostic_segmentation/CASE_release/demo/objects_gallery/obj_000001/2.jpg",
         "https://raw.githubusercontent.com/AnasIbrahim/image_agnostic_segmentation/CASE_release/demo/objects_gallery/obj_000001/3.jpg",
         "https://raw.githubusercontent.com/AnasIbrahim/image_agnostic_segmentation/CASE_release/demo/objects_gallery/obj_000001/4.jpg",
         "https://raw.githubusercontent.com/AnasIbrahim/image_agnostic_segmentation/CASE_release/demo/objects_gallery/obj_000001/5.jpg",
-        "https://raw.githubusercontent.com/AnasIbrahim/image_agnostic_segmentation/CASE_release/demo/objects_gallery/obj_000001/6.jpg"
-     ]
+        "https://raw.githubusercontent.com/AnasIbrahim/image_agnostic_segmentation/CASE_release/demo/objects_gallery/obj_000001/6.jpg",
+        "tiny",
+    ]
 ]
 
 DEVICE = torch.device('cuda')
@@ -54,11 +53,10 @@ DOUNSEEN_MODEL = load_dounseen_model(device=DEVICE)
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process(
-    checkpoint_dropdown,
-    threshold,
     image_input,
-    object_image1=None, object_image2=None, object_image3=None, object_image4=None, object_image5=None, object_image6=None
-) -> Optional[Image.Image]:
+    object_image1=None, object_image2=None, object_image3=None, object_image4=None, object_image5=None, object_image6=None,
+    checkpoint_dropdown=CHECKPOINT_NAMES[0],
+) -> Tuple[Image.Image, Image.Image]:
     model = MASK_GENERATORS[checkpoint_dropdown]
     image = np.array(image_input.convert("RGB"))
     sam2_result = model.generate(image)
@@ -85,19 +83,25 @@ def process(
 
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
-    with gr.Row():
-        checkpoint_dropdown_component = gr.Dropdown(
-            choices=CHECKPOINT_NAMES,
-            value=CHECKPOINT_NAMES[0],
-            label="Checkpoint", info="Select a SAM2 checkpoint to use.",
-            interactive=True
-        )
-        threshold = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="Threshold", value=0.5)
     with gr.Row():
         with gr.Column():
             image_input_component = gr.Image(type='pil', label='Upload image')
+            with gr.Row():
+                object_image_1 = gr.Image(type='pil', label=f'Object Image 1')
+                object_image_2 = gr.Image(type='pil', label=f'Object Image 2')
+                object_image_3 = gr.Image(type='pil', label=f'Object Image 3')
+            with gr.Row():
+                object_image_4 = gr.Image(type='pil', label=f'Object Image 4')
+                object_image_5 = gr.Image(type='pil', label=f'Object Image 5')
+                object_image_6 = gr.Image(type='pil', label=f'Object Image 6')
+            object_images = [object_image_1, object_image_2, object_image_3, object_image_4, object_image_5, object_image_6]
+            checkpoint_dropdown_component = gr.Dropdown(
+                choices=CHECKPOINT_NAMES,
+                value=CHECKPOINT_NAMES[0],
+                label="Checkpoint", info="Select a SAM2 checkpoint to use.",
+                interactive=True
+            )
             submit_button_component = gr.Button(value='Submit', variant='primary')
-            object_images = [gr.Image(type="pil", label=f"Object Image {i + 1}", width=256, height=256) for i in range(6)]  # Set a smaller display size for the object images
         with gr.Column():
             image_output_sam = gr.Image(type='pil', label='SAM2 Output')
             image_output_dounseen = gr.Image(type='pil', label='DoUnseen Output')
@@ -105,11 +109,7 @@ with gr.Blocks() as demo:
         gr.Examples(
             fn=process,
             examples=EXAMPLES,
-            inputs=[
-                checkpoint_dropdown_component,
-                threshold,
-                image_input_component,
-            ] + object_images,
+            inputs= [image_input_component] + object_images + [checkpoint_dropdown_component],
             outputs=[image_output_sam, image_output_dounseen],
             cache_examples=False,
             run_on_click=True
@@ -117,11 +117,7 @@ with gr.Blocks() as demo:
 
     submit_button_component.click(
         fn=process,
-        inputs=[
-            checkpoint_dropdown_component,
-            threshold,
-            image_input_component,
-        ] + object_images,
+        inputs=[image_input_component] + object_images + [checkpoint_dropdown_component],
         outputs=[image_output_sam, image_output_dounseen]
     )