Spaces:
Sleeping
Sleeping
File size: 2,241 Bytes
3a487cd c511cb2 0a0f7e5 c511cb2 1f17d7a c511cb2 1f17d7a c511cb2 1f17d7a c511cb2 1f17d7a c511cb2 1b9a030 c511cb2 0a0f7e5 | 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 | import os
os.system('pip install --upgrade gradio')
import cv2
import gradio as gr
import requests
import pickle
import numpy as np
# Function to predict and show bounding boxes
def predict_and_show_bounding_boxes(image_path):
try:
# Load the image using cv2
img = cv2.imread(image_path)
if img is None:
print(f"Error: Could not load image at {image_path}")
return None, "Error: Could not load image"
# Perform inference using the loaded YOLO model
results = model.predict(source=image_path, save=False, conf=0.5)
result = results[0]
boxes = result.boxes
if len(boxes) == 0:
# No defects found, show the zero components image
zero_components_img = cv2.imread('zero_components.png')
if zero_components_img is not None:
img = zero_componentss_img
return img
else:
return None, "Error: Could not load zero components image"
for box in boxes:
xyxy = box.xyxy[0].tolist()
x_min, y_min, x_max, y_max = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])
conf = box.conf[0].item()
cls = int(box.cls[0])
cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
label = f"{result.names[cls]}: {conf:.2f}"
cv2.putText(img, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
# Return the processed image
return img
except Exception as e:
print(f"An error occurred during prediction: {e}")
return None, str(e)
try:
with open('pcb_component_detection.pkl', 'rb') as file:
model = pickle.load(file)
print("YOLO model loaded successfully.")
except FileNotFoundError:
print("Error: 'pcb_component_detection.pkl' not found.")
except Exception as e:
print(f"An error occurred while loading the model: {e}")
# Create Gradio interface
iface = gr.Interface(
fn=predict_and_show_bounding_boxes,
inputs=gr.Image(type="filepath"),
outputs=[gr.Image()],
title="Components Detection",
description="Upload an image to detect defects"
)
iface.launch(share=True) |