Update app.py

app.py

@@ -5,21 +5,14 @@ import numpy as np
 import cv2
 from PIL import Image
 
-# Pre-load the base configuration and models (without setting a threshold yet)
-base_config = DetrConfig.from_pretrained("facebook/detr-resnet-101")
-base_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-101", config=base_config)
-base_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-101")
+def load_model(model_name, threshold):
+    config = DetrConfig.from_pretrained(model_name, threshold=threshold)
+    model = DetrForObjectDetection.from_pretrained(model_name, config=config)
+    image_processor = DetrImageProcessor.from_pretrained(model_name)
+    return pipeline(task='object-detection', model=model, image_processor=image_processor)
 
-def load_model(threshold):
-    # Adjust the configuration for the current threshold
-    config = DetrConfig.from_pretrained("facebook/detr-resnet-50", threshold=threshold)
-    # Create a new model instance with the updated configuration
-    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-101", config=config)
-    # Image processor does not need to be re-loaded
-    return pipeline(task='object-detection', model=model, image_processor=base_processor)
-
-# Initialize the pipeline with a default threshold
-od_pipe = load_model(0.25)  # Set a default threshold here
+# Load the initial model with default threshold
+od_pipe = load_model("facebook/detr-resnet-101", 0.25)  # Setting a default threshold
 
 def draw_detections(image, detections):
     np_image = np.array(image)
@@ -31,34 +24,33 @@ def draw_detections(image, detections):
         box = detection['box']
         x_min, y_min = box['xmin'], box['ymin']
         x_max, y_max = box['xmax'], box['ymax']
-        # Draw rectangles and text with a larger font
         cv2.rectangle(np_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
         label_text = f'{label} {score:.2f}'
-        # Increase the font size and text thickness
         cv2.putText(np_image, label_text, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 255, 255), 4)
     
-    # Convert BGR to RGB for displaying
     final_image = cv2.cvtColor(np_image, cv2.COLOR_BGR2RGB)
     final_pil_image = Image.fromarray(final_image)
     return final_pil_image
 
-def get_pipeline_prediction(threshold, pil_image):
+def get_pipeline_prediction(model_name, threshold, pil_image):
     global od_pipe
-    od_pipe = load_model(threshold)  # reload model with the specified threshold
+    od_pipe = load_model(model_name, threshold)  # Reload model with the specified model and threshold
     try:
         if not isinstance(pil_image, Image.Image):
             pil_image = Image.fromarray(np.array(pil_image).astype('uint8'), 'RGB')
         result = od_pipe(pil_image)
         processed_image = draw_detections(pil_image, result)
-        return processed_image, result
+        description = f'Model used: {model_name}, Detection Threshold: {threshold}'
+        return processed_image, result, description
     except Exception as e:
-        return pil_image, {"error": str(e)}
+        return pil_image, {"error": str(e)}, "Failed to process image"
 
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             gr.Markdown("## Object Detection")
             inp_image = gr.Image(label="Upload your image here")
+            model_dropdown = gr.Dropdown(choices=["facebook/detr-resnet-50", "facebook/detr-resnet-50-panoptic", "facebook/detr-resnet-101", "facebook/detr-resnet-101-panoptic"], value="facebook/detr-resnet-101", label="Select Model")
             threshold_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.25, label="Detection Threshold")
             run_button = gr.Button("Detect Objects")
         with gr.Column():
@@ -66,7 +58,9 @@ with gr.Blocks() as demo:
                 output_image = gr.Image()
             with gr.Tab("Detection Results"):
                 output_data = gr.JSON()
+            with gr.Tab("Description"):
+                description_output = gr.Textbox()
 
-    run_button.click(get_pipeline_prediction, inputs=[threshold_slider, inp_image], outputs=[output_image, output_data])
+    run_button.click(get_pipeline_prediction, inputs=[model_dropdown, threshold_slider, inp_image], outputs=[output_image, output_data, description_output])
 
 demo.launch()