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)