|
import os |
|
import cv2 |
|
import gradio as gr |
|
from ultralytics import YOLO |
|
|
|
|
|
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 |
|
|
|
|
|
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() |