|
import torch |
|
import cv2 |
|
import numpy as np |
|
import gradio as gr |
|
from ultralytics import YOLO |
|
from PIL import Image |
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
model_path = hf_hub_download("keremberke/yolov8n-pcb-defect-segmentation", "best.pt") |
|
|
|
|
|
model = YOLO(model_path) |
|
|
|
def detect_defects(image): |
|
results = model.predict(image, save=False) |
|
annotated_image = results[0].plot() |
|
return annotated_image |
|
|
|
iface = gr.Interface( |
|
fn=detect_defects, |
|
inputs=gr.Image(type="pil"), |
|
outputs=gr.Image(type="pil"), |
|
title="PCB Defect Segmentation" |
|
) |
|
|
|
iface.launch() |