Spaces:
Sleeping
Sleeping
| 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() |