import gradio as gr
import torch
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import matplotlib.pyplot as plt
import os

# Download images
torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg')
torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', 'bus.jpg')

# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')


def yolo(im):
    try:
        # Check if the input is an Image object
        if isinstance(im, Image.Image):
            # Convert the PIL image to a numpy array
            im_array = np.array(im)

            # Perform inference with YOLOv5
            results = model(im_array)  # inference

            # Get the bounding boxes and labels
            boxes = results.xyxy[0].cpu().numpy()

            # Convert the results to a PIL Image
            output_image = Image.fromarray(im_array)

            # Draw the bounding boxes and labels on the output image
            draw = ImageDraw.Draw(output_image)
            font = ImageFont.load_default(45)

            for box in boxes:
                label = results.names[int(box[5])]
                draw.rectangle([(box[0], box[1]), (box[2], box[3])], outline="red", width=3)
                draw.text((box[0], box[1]), label, fill="blue", font=font)

            return output_image
        else:
            raise ValueError("The input should be an Image object.")
    except Exception as e:
        print(f"Error processing image: {e}")
        return None


# Define Gradio interface
inputs = gr.Image(type='pil', label="Original Image")
outputs = gr.Image(type="pil", label="Output Image")
title = "YOLOv5"
description = "YOLOv5 Gradio demo for object detection. Upload an image or click an example image to use."

article = "<p style='text-align: center'>YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset, and includes " \
          "simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, " \
          "and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> |" \
          "<a href='https://apps.apple.com/app/id1452689527'>iOS App</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"

examples = [['zidane.jpg'], ['bus.jpg']]

gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch(debug=True)