import gradio as gr
import torch
from PIL import Image
import subprocess
import os
import PIL
from pathlib import Path
import uuid

# Images
torch.hub.download_url_to_file('https://miro.medium.com/max/1400/1*EYFejGUjvjPcc4PZTwoufw.jpeg', '1*EYFejGUjvjPcc4PZTwoufw.jpeg')
torch.hub.download_url_to_file('https://production-media.paperswithcode.com/tasks/ezgif-frame-001_OZzxdny.jpg', 'ezgif-frame-001_OZzxdny.jpg')
torch.hub.download_url_to_file('https://favtutor.com/resources/images/uploads/Social_Distancing_Covid_19__1.jpg', 'Social_Distancing_Covid_19__1.jpg')
torch.hub.download_url_to_file('https://nkcf.org/wp-content/uploads/2017/11/people.jpg', 'people.jpg')

def yolo(im):
  file_name = str(uuid.uuid4())
  im.save(f'{file_name}.jpg')
  os.system(f"python tools/infer.py --weights yolov6s.pt --source {str(file_name)}.jpg --project ''")
  img = PIL.Image.open(f"exp/{file_name}.jpg")
  os.remove(f"exp/{file_name}.jpg")
  os.remove(f'{file_name}.jpg')
  return img

inputs = gr.inputs.Image(type='pil', label="Original Image")
outputs = gr.outputs.Image(type="pil", label="Output Image")

title = "YOLOv6 - Demo"
description = "YOLOv6 is a single-stage object detection framework dedicated to industrial applications, with hardware-friendly efficient design and high performance. Here is a quick Gradio Demo for testing YOLOv6s model. More details from  <a href='https://github.com/meituan/YOLOv6'>https://github.com/meituan/YOLOv6</a> "
article = "<p>YOLOv6-nano achieves 35.0 mAP on COCO val2017 dataset with 1242 FPS on T4 using TensorRT FP16 for bs32 inference, and YOLOv6-s achieves 43.1 mAP on COCO val2017 dataset with 520 FPS on T4 using TensorRT FP16 for bs32 inference. More information at <a href='https://github.com/meituan/YOLOv6'>https://github.com/meituan/YOLOv6</a></p>"

examples = [['1*EYFejGUjvjPcc4PZTwoufw.jpeg'], ['ezgif-frame-001_OZzxdny.jpg'], ['Social_Distancing_Covid_19__1.jpg'], ['people.jpg']]

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