import gradio as gr
from ultralytics import YOLO
import cv2

def check_acc(box):
    res_index_list = box.cls.tolist()
    result = ""

    for index in res_index_list:
        if index == 1:
            result = "Bike Bike Accident Detected"
            break
        elif index == 2:
            result = "Bike Object Accident Detected"
            break
        elif index == 3:
            result = "Bike Person Accident Detected"
            break
        elif index == 5:
            result = "Car Bike Accident Detected"
            break
        elif index == 6:
            result = "Car Car Accident Detected"
            break
        elif index == 7:
            result = "Car Object Accident Detected"
            break
        elif index == 8:
            result = "Car Person Accident Detected"
            break
    
    return result

def image_predict(image):
    res = ""
    model_path = "best.pt"
    model = YOLO(model_path)
    results = model.predict(image,conf = 0.8,iou = 0.3,imgsz = 512)
    box = results[0].boxes
    res = check_acc(box)
    annotated_frame = results[0].plot()
    if len(res) >0:
        return (res, annotated_frame)

    return ("No Class Detected", None)

def extract_frames(video):
    vidcap = cv2.VideoCapture(video)
    vidcap = cv2.VideoCapture(video)
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    nof = 4
    frame_no = 0
    while vidcap.isOpened():
        res = ""
        render = None
        success, image = vidcap.read()
        
        if success ==False:
            break
        
        # Check if it's time to process the frame based on the desired interval
        if (frame_no % (int(fps / nof))) == 0:
            model_path = "best.pt"
            model = YOLO(model_path)
            results = model.predict(image,conf = 0.8,iou = 0.3,imgsz = 512)  
            box = results[0].boxes
            res = check_acc(box) 
            annotated_frame = results[0].plot()

            if len(res) >0:
                return (res, annotated_frame)
            
        frame_no += 1  # Increment frame number
        
    return ("No Class Detected", None)


def take_input(image, video):
    if(video != None):
        res = extract_frames(video)
    else:
        res = image_predict(image)
    return res


with gr.Blocks(title="YOLOS Object Detection", css=".gradio-container {background:lightyellow;}") as demo:
    gr.HTML('<h1>Accident Detection Using Yolov8</h1>')
    gr.HTML("<br>")
    with gr.Row():
        input_image = gr.Image(label="Input image")
        input_video = gr.Video(label="Input video")
        output_label = gr.Text(label="output label")
        output_image = gr.Image(label="Output image")
    gr.HTML("<br>")
    send_btn = gr.Button("Detect")
    gr.HTML("<br>")

    send_btn.click(fn=take_input, inputs=[input_image, input_video], outputs=[output_label, output_image])

demo.launch(debug=True)