Spaces:
Sleeping
Sleeping
File size: 1,778 Bytes
2a05d78 | 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 | import gradio as gr
import numpy as np
from ultralytics import YOLO
from PIL import Image
import cv2
# ----------------------------------
# Load Model (must be in Space root)
# ----------------------------------
model = YOLO("best.pt")
# Automatically use model's trained class names
CLASS_NAMES = model.names
# ----------------------------------
# Inference Function
# ----------------------------------
def detect_keys(image, conf_threshold):
# Convert PIL → numpy
img_np = np.array(image)
# Run YOLO inference
results = model(img_np, conf=float(conf_threshold))
# Get annotated image with bounding boxes
annotated = results[0].plot()
# Convert BGR → RGB for Gradio display
annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
# Extract detected class names
detected_classes = set()
if results[0].boxes is not None:
for box in results[0].boxes:
cls_id = int(box.cls[0])
detected_classes.add(CLASS_NAMES[cls_id])
if detected_classes:
detected_text = "Detected: " + ", ".join(sorted(detected_classes))
else:
detected_text = "No keys detected"
return annotated, detected_text
# ----------------------------------
# Gradio Interface
# ----------------------------------
demo = gr.Interface(
fn=detect_keys,
inputs=[
gr.Image(type="pil", label="Upload Image"),
gr.Slider(0.1, 1.0, value=0.25, step=0.05, label="Confidence Threshold"),
],
outputs=[
gr.Image(type="numpy", label="Detection Result"),
gr.Textbox(label="Detected Labels"),
],
title="🔑 Keys Detection (YOLO)",
description="Upload an image to detect keys defects (cracking / missing).",
)
if __name__ == "__main__":
demo.launch()
|