File size: 3,361 Bytes
2d7bad3
 
 
 
c988546
2d7bad3
 
c988546
 
2d7bad3
 
accfa62
3a603bd
 
2d7bad3
 
3a603bd
 
 
 
 
accfa62
 
 
 
 
 
 
 
 
 
3a603bd
accfa62
3a603bd
 
2d7bad3
3a603bd
2d7bad3
 
 
 
accfa62
 
2d7bad3
 
9ba69a1
2d7bad3
9ba69a1
 
 
 
180137f
9ba69a1
 
 
 
2d7bad3
3a603bd
9ba69a1
2d7bad3
9ba69a1
2d7bad3
 
 
 
9ba69a1
 
c988546
caa4724
9ba69a1
 
82257b6
 
 
 
 
 
99b45ec
 
 
 
 
 
 
9ba69a1
 
 
 
 
 
2d7bad3
 
 
3a603bd
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
import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

# Initialize models
model_path = hf_hub_download(repo_id="brainwavecollective/yolov8n-rubber-duck-detector", filename="yolov8n_rubberducks.pt")
duck_model = YOLO(model_path)
standard_model = YOLO('yolov8n.pt')

def process_image(image, model, is_standard_model=True):
    results = model(image)
    processed_image = image.copy()
    
    for r in results:
        boxes = r.boxes
        for box in boxes:
            cls = int(box.cls[0])
            class_name = model.names[cls]
            
            # For standard model, only show teddy bears
            if is_standard_model and class_name != "teddy bear":
                continue
                
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
            conf = float(box.conf[0])
            
            # Rename class to "rubber duck"
            display_name = "rubber duck" if not is_standard_model else class_name
            
            cv2.rectangle(processed_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f"{display_name} ({conf:.2f})"
            cv2.putText(processed_image, label, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    return processed_image

def compare_models(input_image):
    image = np.array(input_image)
    
    standard_image = process_image(image, standard_model, is_standard_model=True)
    duck_image = process_image(image, duck_model, is_standard_model=False)
    
    height, width = image.shape[:2]
    gap = 20  # Add a 20-pixel gap between images
    
    # Create canvas with extra width for the gap
    canvas = np.zeros((height, width * 2 + gap, 3), dtype=np.uint8)
    
    # Place images with gap in between
    canvas[:, :width] = standard_image
    canvas[:, width + gap:] = duck_image
    
    # Fill gap with white or gray color (optional)
    canvas[:, width:width + gap] = [128, 128, 128]  # Gray gap
    
    # Add model labels
    cv2.putText(canvas, "YOLOv8", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.putText(canvas, "YOLOv8 Fine Tune", (width + gap + 10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    
    return canvas

# Create Gradio interface with stacked layout
with gr.Blocks() as iface:
    gr.Markdown("# Improved YOLO Rubber Duck Detection")
    gr.Markdown("Compare standard YOLOv8 (left) with [Fine-tuned YOLOv8 Rubber Duck Detector](https://huggingface.co/brainwavecollective/yolov8n-rubber-duck-detector) (right)")
    
    with gr.Column():
        # First define the input and output components
        input_image = gr.Image(type="pil", label="Input Image", height=200)
        submit_btn = gr.Button("Compare Models")
        output_image = gr.Image(type="numpy", label="Comparison Result", height=400)
        
        # Then add Examples section that references them
        gr.Examples(
            examples=[["test_image.jpg"]],
            inputs=input_image,
            outputs=output_image,
            fn=compare_models,
            cache_examples=True
        )
    
    submit_btn.click(
        fn=compare_models,
        inputs=input_image,
        outputs=output_image,
    )

# Launch the interface
if __name__ == "__main__":
    iface.launch()