Update app.py

app.py

@@ -5,82 +5,63 @@ import numpy as np
 import cv2
 from PIL import Image
 
-# Initialize the model
-config = DetrConfig.from_pretrained("facebook/detr-resnet-50")
-model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config=config)
-image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+# Pre-load the base configuration and models (without setting a threshold yet)
+base_config = DetrConfig.from_pretrained("facebook/detr-resnet-50")
+base_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", config=base_config)
+base_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 
 def load_model(threshold):
-    # Reinitialize the model with the desired detection 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-50", config=config)
-    image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-    return pipeline(task='object-detection', model=model, image_processor=image_processor)
+    # Image processor does not need to be re-loaded
+    return pipeline(task='object-detection', model=model, image_processor=base_processor)
 
-od_pipe = load_model(0.5)  # Default threshold
+# Initialize the pipeline with a default threshold
+od_pipe = load_model(0.25)  # Set a default threshold here
 
 def draw_detections(image, detections):
-    # Convert PIL image to a numpy array
     np_image = np.array(image)
-    
-    # Convert RGB to BGR for OpenCV
     np_image = cv2.cvtColor(np_image, cv2.COLOR_RGB2BGR)
     
-    # Draw detections
     for detection in detections:
         score = detection['score']
         label = detection['label']
         box = detection['box']
-        x_min = box['xmin']
-        y_min = box['ymin']
-        x_max = box['xmax']
-        y_max = box['ymax']
-
-        # Increase font size for better visibility
+        x_min, y_min = box['xmin'], box['ymin']
+        x_max, y_max = box['xmax'], box['ymax']
         cv2.rectangle(np_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
-        label_text = f'{label} {score:.2f}'
-        cv2.putText(np_image, label_text, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 255, 255), 4)
+        cv2.putText(np_image, f"{label} {score:.2f}", (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
     
-    # 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
+    return Image.fromarray(final_image)
 
 def get_pipeline_prediction(threshold, pil_image):
     global od_pipe
+    od_pipe = load_model(threshold)  # reload model with the specified threshold
     try:
-        # Check if the model threshold needs adjusting
-        if od_pipe.config.threshold != threshold:
-            od_pipe = load_model(threshold)
-            print("Model reloaded with new threshold:", threshold)
-        
-        # Ensure input is a PIL image
         if not isinstance(pil_image, Image.Image):
             pil_image = Image.fromarray(np.array(pil_image).astype('uint8'), 'RGB')
-        
-        # Run detection and return annotated image and results
-        pipeline_output = od_pipe(pil_image)
-        processed_image = draw_detections(pil_image, pipeline_output)
-        return processed_image, pipeline_output
+        result = od_pipe(pil_image)
+        processed_image = draw_detections(pil_image, result)
+        return processed_image, result
     except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
-        print(error_message)
-        return pil_image, {"error": error_message}
+        return pil_image, {"error": str(e)}
 
-# Gradio interface
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            inp_image = gr.Image(label="Input image")
-            slider = gr.Slider(minimum=0, maximum=1, step=0.05, label="Detection Sensitivity", value=0.5)
-            gr.Markdown("Adjust the slider to change detection sensitivity.")
-            btn_run = gr.Button('Run Detection')
+            gr.Markdown("## Object Detection")
+            inp_image = gr.Image(label="Upload your image here")
+            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():
             with gr.Tab("Annotated Image"):
-                out_image = gr.Image()
+                output_image = gr.Image()
             with gr.Tab("Detection Results"):
-                out_json = gr.JSON()
+                output_data = gr.JSON()
 
-    btn_run.click(get_pipeline_prediction, inputs=[slider, inp_image], outputs=[out_image, out_json])
+    run_button.click(get_pipeline_prediction, inputs=[threshold_slider, inp_image], outputs=[output_image, output_data])
 
 demo.launch()