import os
import subprocess

# Clone the yolov5 repository and install its requirements
if not os.path.exists('yolov5'):
    subprocess.run(['git', 'clone', 'https://github.com/ultralytics/yolov5'], check=True)
subprocess.run(['pip', 'install', '-r', 'yolov5/requirements.txt'], check=True)

import torch
import torchvision
from torchvision.transforms import functional as F
from PIL import Image
import cv2
import gradio as gr
import numpy as np

from yolov5.models.yolo import Model
from yolov5.utils.general import non_max_suppression

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
model.eval()
print("Model loaded successfully")

def preprocess_image(image):
    try:
        image = Image.fromarray(image)  # Convert numpy array to PIL Image
        image_tensor = F.to_tensor(image).unsqueeze(0).to(device)
        print(f"Preprocessed image tensor: {image_tensor.shape}")
        return image_tensor
    except Exception as e:
        print(f"Error in preprocessing image: {e}")
        return None

def draw_boxes(image, outputs, threshold=0.3):
    try:
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        h, w, _ = image.shape

        for box in outputs:
            if box is not None:
                x1, y1, x2, y2, score, label = box[:6]
                if score > threshold:
                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                    cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
                    text = f"{model.names[int(label)]:s}: {score:.2f}"
                    cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    except Exception as e:
        print(f"Error in drawing boxes: {e}")
        return image

def detect_objects(image):
    image_tensor = preprocess_image(image)
    if image_tensor is None:
        return image
    try:
        outputs = model(image_tensor)[0]  # Get the first element of the output
        print(f"Model raw outputs: {outputs}")
        outputs = non_max_suppression(outputs, conf_thres=0.25, iou_thres=0.45)[0]  # Apply NMS
        if outputs is None or len(outputs) == 0:
            print("No objects detected.")
            return image
        print(f"Filtered outputs: {outputs}")
        result_image = draw_boxes(image, outputs.cpu().numpy())
        return result_image
    except Exception as e:
        print(f"Error in detecting objects: {e}")
        return image

iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="numpy"),
    title="YOLOv5 Object Detection",
    description="Upload an image to detect objects using the YOLOv5 model."
)

if __name__ == "__main__":
    iface.launch()