File size: 3,186 Bytes
2d7bad3 3a603bd 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 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 |
import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
# Initialize models
duck_model = YOLO('https://huggingface.co/brainwavecollective/yolo8n-rubber-duck-detector/resolve/main/yolov8n_rubberducks.pt')
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("# YOLO Model Comparison")
gr.Markdown("Compare standard YOLOv8 model (left) with Fine-tuned model (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() |