improve app design

app.py

@@ -16,11 +16,16 @@ with demo:
        
     with gr.Box():
         
+        
         with gr.Row():
-            segmentation_task = gr.Dropdown(["semantic", "instance", "panoptic"], value="panoptic", label="Segmentation Task", show_label=True)
-        with gr.Box():
-            with gr.Row():
+            with gr.Column():
+                gr.Markdown("**Inputs**")
+                segmentation_task = gr.Dropdown(["semantic", "instance", "panoptic"], value="panoptic", label="Segmentation Task", show_label=True)
                 input_image = gr.Image(type='filepath',label="Input Image", show_label=True)
+            
+            with gr.Column():
+                gr.Markdown("**Outputs**")
+                output_heading = gr.Textbox(label="Output Type", show_label=True)
                 output_mask = gr.Image(label="Predicted Masks", show_label=True)
     
     gr.Markdown("**Predict**")
@@ -32,10 +37,10 @@ with demo:
     gr.Markdown("**Examples:**")
 
     with gr.Column():
-        gr.Examples(example_list, [input_image, segmentation_task], output_mask, predict_masks)
+        gr.Examples(example_list, [input_image, segmentation_task], [output_mask,output_heading], predict_masks)
         
     
-    submit_button.click(predict_masks, inputs=[input_image, segmentation_task], outputs=output_mask)
+    submit_button.click(predict_masks, inputs=[input_image, segmentation_task], outputs=[output_mask,output_heading])
     
     gr.Markdown('\n Demo created by: <a href=\"https://www.linkedin.com/in/shivalika-singh/\">Shivalika Singh</a>')
 

predict.py

@@ -4,7 +4,8 @@ import numpy as np
 from PIL import Image
 from collections import defaultdict
 import os
-# install detectron this way to avoid torch not available error if mentioned directly as dependency in requirements.txt
+# Mentioning detectron2 as a dependency directly in requirements.txt tries to install detectron2 before torch and results in an error even if torch is listed as a dependency before detectron2.
+# Hence, installing detectron2 this way when using Gradio HF spaces.
 os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
 
 from detectron2.data import MetadataCatalog
@@ -21,12 +22,12 @@ def load_model_and_processor(model_ckpt: str):
 
 def load_default_ckpt(segmentation_task: str):
     if segmentation_task == "semantic":
-        default_pretrained_ckpt = "facebook/mask2former-swin-tiny-ade-semantic"
+        default_ckpt = "facebook/mask2former-swin-tiny-ade-semantic"
     elif segmentation_task == "instance":
-        default_pretrained_ckpt = "facebook/mask2former-swin-small-coco-instance"
+        default_ckpt = "facebook/mask2former-swin-small-coco-instance"
     else:
-        default_pretrained_ckpt = "facebook/mask2former-swin-tiny-coco-panoptic"
-    return default_pretrained_ckpt
+        default_ckpt = "facebook/mask2former-swin-tiny-coco-panoptic"
+    return default_ckpt
 
 def draw_panoptic_segmentation(predicted_segmentation_map, seg_info, image):
     metadata = MetadataCatalog.get("coco_2017_val_panoptic")
@@ -73,8 +74,8 @@ def visualize_instance_seg_mask(mask, input_image):
 def predict_masks(input_img_path: str, segmentation_task: str):
     
     #load model and image processor
-    default_pretrained_ckpt = load_default_ckpt(segmentation_task)
-    model, image_processor = load_model_and_processor(default_pretrained_ckpt)
+    default_ckpt = load_default_ckpt(segmentation_task)
+    model, image_processor = load_model_and_processor(default_ckpt)
     
     ## pass input image through image processor
     image = Image.open(input_img_path)
@@ -90,16 +91,20 @@ def predict_masks(input_img_path: str, segmentation_task: str):
         predicted_segmentation_map = result.cpu().numpy()
         palette = ade_palette()
         output_result = draw_semantic_segmentation(predicted_segmentation_map, image, palette)
+        output_heading = "Semantic Segmentation Output"
 
     elif segmentation_task == "instance":
         result = image_processor.post_process_instance_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
         predicted_instance_map = result["segmentation"].cpu().detach().numpy()
         output_result = visualize_instance_seg_mask(predicted_instance_map, image)
+        output_heading = "Instance Segmentation Output"
 
     else:
         result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
         predicted_segmentation_map = result["segmentation"]
         seg_info = result['segments_info']
         output_result = draw_panoptic_segmentation(predicted_segmentation_map, seg_info, image)
+        output_heading = "Panoptic Segmentation Output"
 
-    return output_result
+
+    return output_result, output_heading