Update app.py

app.py

@@ -2,69 +2,39 @@ import gradio as gr
 import cv2
 import numpy as np
 from ultralytics import YOLO
-import torch
 
 # Initialize models
-duck_model = YOLO('https://huggingface.co/brainwavecollective/yolo8n-rubber-duck-detector/resolve/main/yolov8n_rubberducks4.pt')
+duck_model = YOLO('https://huggingface.co/brainwavecollective/yolo8n-rubber-duck-detector/resolve/main/yolov8n_rubberducks.pt')
 standard_model = YOLO('yolov8n.pt')
 
-def calculate_iou(box1, box2):
-    x1 = max(box1[0], box2[0])
-    y1 = max(box1[1], box2[1])
-    x2 = min(box1[2], box2[2])
-    y2 = min(box1[3], box2[3])
-    intersection = max(0, x2 - x1) * max(0, y2 - y1)
-    
-    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
-    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
-    union = box1_area + box2_area - intersection
-    
-    return intersection / union if union > 0 else 0
-
-def process_image(image, model, is_duck_model=True):
-    results = model(image, conf=0.4)
+def process_image(image, model):
+    results = model(image)
+    processed_image = image.copy()
     
-    valid_boxes = []
     for r in results:
-        for box in r.boxes:
-            class_name = model.names[int(box.cls[0])]
-            # For both models, only show teddy bear class
-            if class_name == "teddy bear":
-                valid_boxes.append({
-                    'coords': box.xyxy[0].cpu().numpy().tolist(),
-                    'confidence': float(box.conf[0])
-                })
-    
-    # Filter overlapping boxes
-    filtered_boxes = []
-    for i, box in enumerate(valid_boxes):
-        should_add = True
-        for existing_box in filtered_boxes:
-            if calculate_iou(box['coords'], existing_box['coords']) > 0.5:
-                if box['confidence'] <= existing_box['confidence']:
-                    should_add = False
-                    break
-        if should_add:
-            filtered_boxes.append(box)
+        boxes = r.boxes
+        for box in boxes:
+            # Get box coordinates and confidence
+            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
+            conf = float(box.conf[0])
+            cls = int(box.cls[0])
+            class_name = model.names[cls]
+            
+            # Draw box and label
+            cv2.rectangle(processed_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            label = f"{class_name} ({conf:.2f})"
+            cv2.putText(processed_image, label, (x1, y1-10), 
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
     
-    # Draw boxes and labels
-    processed_image = image.copy()
-    for box in filtered_boxes:
-        x1, y1, x2, y2 = map(int, box['coords'])
-        cv2.rectangle(processed_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        label = f"{'Duck' if is_duck_model else 'Standard'} ({box['confidence']:.2f})"
-        cv2.putText(processed_image, label, (x1, y1-10), 
-                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-    
-    return processed_image, len(filtered_boxes)
+    return processed_image
 
 def compare_models(input_image):
     # Convert from Gradio's PIL image to OpenCV format
     image = np.array(input_image)
     
     # Process with both models
-    duck_image, duck_detections = process_image(image, duck_model, True)
-    standard_image, std_detections = process_image(image, standard_model, False)
+    duck_image = process_image(image, duck_model)
+    standard_image = process_image(image, standard_model)
     
     # Create side-by-side comparison
     height, width = image.shape[:2]
@@ -74,16 +44,11 @@ def compare_models(input_image):
     canvas[:, :width] = duck_image
     canvas[:, width:] = standard_image
     
-    # Add labels for each model
-    cv2.putText(canvas, "Rubber Duck YOLO", (10, 30), 
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-    cv2.putText(canvas, f"Detections: {duck_detections}", (10, 60), 
+    # Add model labels
+    cv2.putText(canvas, "Duck Detector", (10, 30), 
                 cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-    
     cv2.putText(canvas, "Standard YOLOv8", (width + 10, 30), 
                 cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-    cv2.putText(canvas, f"Detections: {std_detections}", (width + 10, 60), 
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
     
     return canvas
 
@@ -93,11 +58,11 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="numpy"),
     title="YOLO Model Comparison",
-    description="Compare Rubber Duck YOLO detector with standard YOLOv8 model",
-    examples=[["example1.jpg"], ["example2.jpg"]],  # Add your example images here
+    description="Compare Duck Detector with standard YOLOv8 model",
+    examples=[["test_image.jpg"]],
     cache_examples=True
 )
 
 # Launch the interface
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch()