File size: 999 Bytes
5a3de7a b19eb2e 5a3de7a 944f4e4 5a3de7a 944f4e4 5a3de7a 944f4e4 5a3de7a 944f4e4 5a3de7a 944f4e4 |
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 |
import os
import cv2
import gradio as gr
from ultralytics import YOLO
# Load model
MODEL_PATH = "yolov8_model.pt"
if not os.path.exists(MODEL_PATH):
raise FileNotFoundError(f"Model file {MODEL_PATH} not found. Please upload it to your Space.")
model = YOLO(MODEL_PATH)
def detect_defects(image):
"""Run detection and return annotated image + detection data"""
try:
results = model(image)
annotated_img = results[0].plot(line_width=2)
return cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
except Exception as e:
print(f"Error: {e}")
return image
# Simple UI
interface = gr.Interface(
fn=detect_defects,
inputs=gr.Image(label="Upload Steel Surface Image", type="numpy"),
outputs=gr.Image(label="Detected Defects"),
title="🔧 Steel Surface Defect Detector",
description="Upload an image to detect surface defects (crazing, scratches, etc.)",
allow_flagging="never"
)
if __name__ == "__main__":
interface.launch() |