Create app.py

app.py

@@ -0,0 +1,103 @@
+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')
+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)
+    
+    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)
+    
+    # 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)
+
+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)
+    
+    # Create side-by-side comparison
+    height, width = image.shape[:2]
+    canvas = np.zeros((height, width * 2, 3), dtype=np.uint8)
+    
+    # Place images side by side
+    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), 
+                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
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=compare_models,
+    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
+    cache_examples=True
+)
+
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch(share=True)