Shardul010's picture
Add Gradio app, requirements, YOLOv8 best weights, and sample test images
a784cd1
Raw
History Blame Contribute Delete
1.17 kB
import gradio as gr
from ultralytics import YOLO
from PIL import Image
# Load the optimal trained YOLOv8s weights (CutPaste-Poisson n=100)
model = YOLO("best.pt")
def detect_defects(image):
# Run YOLOv8 inference on the uploaded image
results = model(image)
# Generate the annotated image with bounding boxes and confidence scores
annotated_img = results[0].plot()
# Convert BGR (OpenCV format) to RGB (PIL/Gradio format)
return Image.fromarray(annotated_img[..., ::-1])
# Define the Gradio web interface
iface = gr.Interface(
fn=detect_defects,
inputs=gr.Image(type="pil", label="Upload Steel Surface Image"),
outputs=gr.Image(type="pil", label="Detected Defects & Confidence"),
title="Real-Time Steel Surface Defect Detection",
description="Upload a grayscale steel surface image (NEU-DET) to detect 6 types of industrial defects: **crazing, inclusion, patches, pitted_surface, rolled-in_scale, and scratches**.\n\nThis model is powered by **YOLOv8s** trained with an optimal volume of Cut-Paste Poisson blending synthetic data augmentation.",
examples=[["sample_scratch.jpg"], ["sample_patch.jpg"]]
)
iface.launch()