surabhic's picture
Update app.py
944f4e4 verified
raw
history blame contribute delete
999 Bytes
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()