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()