File size: 3,719 Bytes
2d7bad3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 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) |