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() |