import gradio as gr
import cv2
import os
from ultralyticsplus import YOLO, render_result

model_path= 'best.pt'

def preds_image(image, conf_thres, iou_thres):

    model = YOLO(model_path)

    result = model.predict(image, 
                           conf= conf_thres,
                           iou= iou_thres)

    box = result[0].boxes
    print("Object type: ", box.cls)
    print("Coordinates: ", box.xyxy)
    print("Probability: ", box.conf)

    render = render_result(model=model, image=image, result=result[0])
    return render
 
inputs_image = [
    gr.Image(label="Input Image"),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.40, step=0.05, label="Confidence threshold"),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU threshold")
]
outputs_image = [
    gr.Image(label="Output Image"),
]

interface_image = gr.Interface(
    fn=preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Chili Leaf Disease Detector (Image)"
)

def preds_video(video):

    video_path = video
    video_path_out = '{}_out.mp4'.format(video_path)
    
    cap = cv2.VideoCapture(video_path)
    ret, frame = cap.read()
    H, W, _ = frame.shape
    out = cv2.VideoWriter(video_path_out, cv2.VideoWriter_fourcc(*'MP4V'), int(cap.get(cv2.CAP_PROP_FPS)), (W, H))

    model= YOLO(model_path)

    threshold = 0.4

    while ret:

        results = model(frame)[0]

        for result in results.boxes.data.tolist():
            x1, y1, x2, y2, score, class_id = result

            if score > threshold:
                cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 4)
            
                #object details
                org = (int(x1), int(y1 - 10))
                font = cv2.FONT_HERSHEY_SIMPLEX
                fontscale = 0.5
                color = (0, 0, 0)
            
                cv2.putText(frame, results.names[int(class_id)].lower(), org,
                            font, fontscale, color, 1, cv2.LINE_AA)

        out.write(frame)
        ret, frame = cap.read()

    cap.release()
    out.release()
    cv2.destroyAllWindows()

    return video_path_out
    
inputs_video = gr.Video(label= 'Original chili leaf video')
outputs_video = gr.Video(label= 'Predicted leaf')

interface_video = gr.Interface(
    fn=preds_video,
    inputs=inputs_video,
    outputs=outputs_video,
    title="Chili Leaf Disease Detector (Video)"
)

gr.TabbedInterface(
    [interface_image, interface_video],
    tab_names=['Image inference', 'Video inference']
).queue().launch()