modify visualizer

app.py

@@ -3,18 +3,11 @@ from backend import visualize_image
 
 # gradio inputs
 image_input = gr.inputs.Image(type="pil", label="Input Image")
-color_mode_select = gr.inputs.Radio(["Black/white", "Random", "Segmentation"])
-mode_dropdown = gr.inputs.Dropdown(["Trees", "Buildings", "Both"])
+color_mode_select = gr.inputs.Radio(["Black/white", "Random", "Segmentation"], label="Color Mode")
+mode_dropdown = gr.inputs.Dropdown(["Trees", "Buildings", "Both"], label="Detection Mode")
 
 tree_threshold_slider = gr.inputs.Slider(0, 1, 0.1, 0.7, label='Set confidence threshold "%" for trees')
 building_threshold_slider = gr.inputs.Slider(0, 1, 0.1, 0.7, label='Set confidence threshold "%" for buildings')
-if mode_dropdown == "Trees":
-    tree_threshold_slider.enable()
-if mode_dropdown == "Buildings":
-    building_threshold_slider.enable()
-if mode_dropdown == "Both":
-    tree_threshold_slider.enable()
-    building_threshold_slider.enable()
 
 # gradio outputs
 output_image = gr.outputs.Image(type="pil", label="Output Image")

backend.py

@@ -81,10 +81,21 @@ def visualize_image(im, mode="BOTH", tree_threshold=0.7, building_threshold=0.7,
         building_instances = segment_building(im, building_threshold)
         instances = Instances.cat([tree_instances, building_instances])
 
+    # visualizer = Visualizer(im[:, :, ::-1],
+    #                         scale=0.5,
+    #                         instance_mode=color_mode)
+    
+    metadata = MetadataCatalog.get("your_model_metadata")
+    category_names = metadata.get("thing_classes")
     visualizer = Visualizer(im[:, :, ::-1],
+                            metadata=metadata,
                             scale=0.5,
                             instance_mode=color_mode)
-    
+    # in the visualizer, add category label names to detected instances
+    for instance in instances:
+        label = category_names[instance["category_id"]]
+        visualizer.draw_text(label, instance["bbox"][:2])
+
     output_image = visualizer.draw_instance_predictions(instances)
     
     return Image.fromarray(output_image.get_image()[:, :, ::-1])