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

from ultralytics import YOLO

file_urls = [
    'https://www.dropbox.com/scl/fi/34yt1vrl4mc4n9ujdf9gm/all_76.jpg?rlkey=f7b6nq478r2m9yahcalzjzif5&dl=1',
    'https://www.dropbox.com/scl/fi/lns6cewinp7rgf3v2g1n8/all_5.jpg?rlkey=20zvut81b829k9lg5yk8ve99z&dl=1',
    'https://www.dropbox.com/scl/fi/13jr2f1znuzulmsyabl2f/long3.jpg?rlkey=jeyriw5a8c0t42e7y2986y53m&dl=1',
    'https://www.dropbox.com/scl/fi/nglwcza7msjo1vu4kw27r/pot4.jpg?rlkey=1ynm35b4j100ta0p5g3fx7hx4&dl=1',
    'https://www.dropbox.com/s/7sjfwncffg8xej2/video_7.mp4?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 = [['VID-20230809-WA0021.mp4'], ['VID-20230809-WA0022.mp4'], 
                ['VID-20230809-WA0024.mp4'], ['VID-20230809-WA0032.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=1024)
    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="Asphalt Road Pavement Distresses Detector",
    examples=path,
    cache_examples=False,
        description= 'This is a demo app that takes in images or videos of Asphalt pavement surfaces and \
         \n detects the following pavement distresses: \
         \n    \
         \n Alligator cracking \
        \n Longitudinal cracking \
        \n Potholes \
        \n Ravelling \
        \n \
        \n This is specifically for Inference and educational purpose.\
        \n \
        \n The model might ocassionaly give false outputs'
    )

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=1024)
            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()