Spaces:
Sleeping
Sleeping
File size: 2,184 Bytes
7d592db ffe3d85 7d592db ffe3d85 eeac6fb 7d592db ffe3d85 7d592db | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import gradio as gr
from ultralytics import YOLO
import os
# Load YOLO model
model = YOLO("best.pt")
def detect_defects(image):
results = model.predict(image, conf=0.25)
return results[0].plot()
# ---------------------------
# Defect Classes (10)
# ---------------------------
classes = [
"punching_hole",
"welding_line",
"crescent_gap",
"water_spot",
"oil_spot",
"silk_spot",
"inclusion",
"rolled_pit",
"crease",
"waist_folding"
]
# ---------------------------
# Load category-wise examples
# ---------------------------
example_dir = "examples"
category_examples = {}
for cls in classes:
class_folder = os.path.join(example_dir, cls)
if os.path.exists(class_folder):
imgs = [
os.path.join(class_folder, f)
for f in sorted(os.listdir(class_folder))
if f.lower().endswith((".jpg", ".png", ".jpeg"))
]
category_examples[cls] = imgs
else:
category_examples[cls] = []
# ---------------------------
# Build Gradio Interface
# ---------------------------
with gr.Blocks(title="Metal Surface Defect Detection (YOLOv8)") as demo:
gr.Markdown("""
# 🔍 Metal Surface Defect Detection (YOLOv8)
Upload an image or choose an example from the defect categories below.
""")
with gr.Row():
input_img = gr.Image(type="numpy", label="Input Image")
output_img = gr.Image(type="numpy", label="Detection Result")
detect_btn = gr.Button("Run Detection")
detect_btn.click(detect_defects, inputs=input_img, outputs=output_img)
gr.Markdown("## 📂 Choose Example Images by Category")
with gr.Tabs():
for cls in classes:
with gr.Tab(cls):
if len(category_examples[cls]) == 0:
gr.Markdown("_No example images found for this category._")
else:
gr.Examples(
examples=category_examples[cls],
inputs=input_img,
outputs=output_img,
fn=detect_defects,
cache_examples=False
)
demo.launch()
|