import cv2
import os
from ultralytics import YOLO
import numpy as np
import PIL.Image as Image
import gradio as gr




def predict_image(img):
    # print(img)
    # labels=['awake', 'drowsy']
    model= YOLO(r"C:\Users\sdadi\Desktop\M_L\Drowsiness_detec\yolov8s_drowsy.pt")
    im_array=[]
    (height, width, channels)= img.shape
    img= cv2.resize(img, (640, 640))
    img= cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    preds= model.predict(source= img,  show_labels=True, show_conf=False, imgsz=640)
    
    try:
        for pred in preds:
            im_array= pred.plot()
            # im_array = Image.fromarray(im_array[..., ::-1])
        # pred =preds[0]
        # boxes = pred.boxes.cpu().numpy()
        # argmax= np.argmax(boxes.conf[0])
        # labelmax= boxes.cls[argmax]
        # text= labels[int(labelmax)]
            
    except:
        pass
    im_array= cv2.cvtColor(im_array, cv2.COLOR_BGR2RGB)
    im_array= cv2.resize(im_array, (width, height))
    
    return im_array

image_iface= gr.Interface(fn =predict_image,
                          inputs= gr.Image(label='Upload Image', sources=['upload', 'clipboard']), 
                          outputs= gr.Image(label='Inference Results'),
                          description= "Upload Images for Inference using the trained model")
#   gr.slider(minimum= 0.1, maximum= 0.8, value= 0.5))

demo= gr.TabbedInterface(interface_list=[image_iface], tab_names=['Image Inference'], theme="soft", title= "Drowsiness detection using YOLOv8n")

if __name__ =='__main__':
    demo.launch(debug="True")