from PIL import Image
import torch
from numpy import random

from utils.general import (non_max_suppression, scale_coords)
from utils.plots import plot_one_box

from models.models import *
from utils.general import *

import gradio as gr
import requests

import gdown


url = 'https://drive.google.com/u/0/uc?id=1Tdn3yqpZ79X7R1Ql0zNlNScB1Dv9Fp76&export=download'
output = 'yolor_p6.pt'
gdown.download(url, output, quiet=False)

url1 = 'https://cdn.pixabay.com/photo/2014/09/07/21/52/city-438393_1280.jpg'
r = requests.get(url1, allow_redirects=True)
open("city1.jpg", 'wb').write(r.content)
url2 = 'https://cdn.pixabay.com/photo/2016/02/19/11/36/canal-1209808_1280.jpg'
r = requests.get(url2, allow_redirects=True)
open("city2.jpg", 'wb').write(r.content)


conf_thres = 0.4
iou_thres = 0.5


def load_classes(path):
    # Loads *.names file at 'path'
    with open(path, 'r') as f:
        names = f.read().split('\n')
    return list(filter(None, names))  # filter removes empty strings (such as last line)

def detect(pil_img,names):
    img_np = np.array(pil_img)
    img = torch.from_numpy(img_np)
    img = img.float()
    img /= 255.0  # 0 - 255 to 0.0 - 1.0

    # Inference
    pred = model(img.unsqueeze(0).permute(0,3,1,2), augment=False)[0].detach()

    # Apply NMS
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False)

    # Process detections
    for i, det in enumerate(pred):  # detections per image
        if det is not None and len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_coords(img_np.shape, det[:, :4], img_np.shape).round()

            # Print results
            for c in det[:, -1].unique():
                n = (det[:, -1] == c).sum()  # detections per class

            # Write results
            for *xyxy, conf, cls in det:
                label = '%s %.2f' % (names[int(cls)], conf)
                plot_one_box(xyxy, img_np, label=label, color=colors[int(cls)], line_thickness=3)
    return Image.fromarray(img_np)
    

with torch.no_grad():
    cfg = 'cfg/yolor_p6.cfg'
    imgsz = 1280
    names = 'data/coco.names'
    weights = 'yolor_p6.pt'

    # Load model
    model = Darknet(cfg, imgsz)
    model.load_state_dict(torch.load(weights)['model'])
    model.eval()

    # Get names and colors
    names = load_classes(names)
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

    def inference(image):
        image = image.resize(size=(imgsz, imgsz))
        return detect(image, names)

    title = "YOLOR P6"
    description = "demo for YOLOR. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below.\nModel: YOLOR-P6"
    article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2105.04206'>You Only Learn One Representation: Unified Network for Multiple Tasks</a> | <a href='https://github.com/WongKinYiu/yolor'>Github Repo</a></p>"

    gr.Interface(
        inference, 
        [gr.inputs.Image(type="pil", label="Input")], 
        gr.outputs.Image(type="numpy", label="Output"),
        title=title,
        description=description,
        article=article,
        examples=[
                ["city1.jpg"],
                ["city2.jpg"]
            ]).launch()