Spaces:
Sleeping
Sleeping
| 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) |