import gradio as gr
import cv2 
import requests
import os 
import random

from ultralytics import YOLO

file_urls = [
    'https://www.dropbox.com/scl/fi/5pavu4vvkprrtkwktvei7/DSC02373.JPG?rlkey=fpj636qtkf3vrqfxy45n2d9ii&dl=1',
    'https://www.dropbox.com/scl/fi/56pbn4r3ohk85rchcvwdj/DSC02813.JPG?rlkey=jnbsidqtthk6p4ysld6o6kc4t&dl=1',
    'https://www.dropbox.com/scl/fi/av9g5zbmrrzg9064zivat/image_2.jpg?rlkey=ldocvzz5lq98zffqf1lmhbhv1&dl=1',
    'https://www.dropbox.com/scl/fi/izo2eqqnqzcsaxis1qrbx/IMG_7612.JPG?rlkey=6wfjaux44khtlx454ex0ng0hp&dl=1',
    'https://www.dropbox.com/scl/fi/e6vgy1et6vjr61uypk5yu/VID-20230809-WA0021.mp4?rlkey=khv8rw074vezzlg8ob38bpmbx&dl=1'
]

def download_file(url, save_name):
    url = url
    if not os.path.exists(save_name):
        file = requests.get(url)
        open(save_name, 'wb').write(file.content)
 
for i, url in enumerate(file_urls):
    if 'mp4' in file_urls[i]:
        download_file(
            file_urls[i],
            f"video.mp4"
        )
    else:
        download_file(
            file_urls[i],
            f"image_{i}.jpg"
        )


model = YOLO('best.pt')
path  = [['image_0.jpg'], ['image_1.jpg'], ['image_2.jpg'], ['image_3.jpg']]

# path = [['IMG_7612.JPG'], ['IMG_7678.JPG'], ['all_33.jpg'], ['all_80.jpg'], 
#     ['DSC02813.JPG'], ['DSC02373.JPG']]


# path  = [['sc_1_0 (1) (1).JPG'], ['sc_1_0 (16) (1).JPG'], 
#         ['sc_1_0 (18) (1).JPG'], ['sc_1_0 (18).JPG']]

video_path = [['video.mp4']]

classes = ['alligator_cracking', 'longitudinal_cracking', 'potholes', 'ravelling']

def show_preds_image(image_path):
    image = cv2.imread(image_path)
    outputs = model.predict(source=image_path, agnostic_nms=True, conf=0.25, iou=0.4, imgsz=640)
    results = outputs[0].cpu().numpy()

    re_boxes = results.boxes.data.tolist()

    class_colors = {1 : (95, 255, 54), 2: (242, 210, 100), 3: (96, 7, 70), 4:(221, 59, 41)}
    random.seed(42)
    # class_colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for _ in range(4)]

    for i, det in enumerate(results.boxes.xyxy):
        x1, y1, x2, y2 = int(det[0]),  int(det[1]),  int(det[2]),  int(det[3])

        class_label = int(re_boxes[i][-1])
        rectangle_color = class_colors.get(class_label)
        # rectangle_color = class_colors[class_label]
        text_color = rectangle_color
        cv2.rectangle(
            image,
            (int(det[0]), int(det[1])),
            (int(det[2]), int(det[3])),
            color=rectangle_color,
            thickness=3,
            lineType=cv2.LINE_AA
        )

        text_position = (x1, y1+100)
        conf = re_boxes[i][-2]
        class_name = classes[class_label]
        # class_label = class_name.split('_')[0] + '\n' + class_name.split('_')[1] if '_' in class_name else class_name
        cv2.putText(image, classes[class_label] + f' = {round(conf, 2)}',
            text_position, cv2.FONT_HERSHEY_SIMPLEX, 1.5, text_color, 3)
        
    
    # print(class_ids)
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
inputs_image = [
    gr.components.Image(type="filepath", label="Input Image"),
]
outputs_image = [
    gr.components.Image(type="numpy", label="Output Image"),
]
interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Pavement Distress Detector for developing countries",
    examples=path,
    cache_examples=False,
        description= ''
    )

def show_preds_video(video_path):
    cap = cv2.VideoCapture(video_path)
    while(cap.isOpened()):
        ret, frame = cap.read()
        if ret:
            frame_copy = frame.copy()
            outputs = model.predict(source=frame, agnostic_nms=True, conf=0.25, iou=0.4, imgsz=640)
            results = outputs[0].cpu().numpy()
            re_boxes = results.boxes.data.tolist()

            class_colors = {1 : (95, 255, 54), 2: (242, 210, 100), 3: (96, 7, 70), 4:(221, 59, 41)}
            random.seed(42)
            # class_colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for _ in range(4)]
            
            for i, det in enumerate(results.boxes.xyxy):
                x1, y1, x2, y2 = int(det[0]),  int(det[1]),  int(det[2]),  int(det[3])

                class_label = int(re_boxes[i][-1])
                rectangle_color = class_colors.get(class_label)
                # rectangle_color = class_colors[class_label]
                text_color = rectangle_color

                cv2.rectangle(
                    frame_copy,
                    (int(det[0]), int(det[1])),
                    (int(det[2]), int(det[3])),
                    color=rectangle_color,
                    thickness=2,
                    lineType=cv2.LINE_AA
                )

                
                text_position = (x1, y1+100)
                conf = re_boxes[i][-2]
                class_name = classes[class_label]
                # class_label = class_name.split('_')[0] + '\n' + class_name.split('_')[1] if '_' in class_name else class_name
                cv2.putText(frame_copy, classes[class_label] + f' = {round(conf, 2)}',
                text_position, cv2.FONT_HERSHEY_SIMPLEX, 1.5, text_color, 3)

            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
 
inputs_video = [
    gr.components.Video(type="filepath", label="Input Video"),
 
]
outputs_video = [
    gr.components.Image(type="numpy", label="Output Video"),
]
interface_video = gr.Interface(
    fn=show_preds_video,
    inputs=inputs_video,
    outputs=outputs_video,
    title="Asphalt Road Pavement Distresses Detector",
    examples=video_path,
    cache_examples=False,
    # live=True
)
gr.TabbedInterface(
    [interface_image, interface_video],
    tab_names=['Image inference', 'Video inference'],
).queue().launch()