import gradio as gr
import numpy as np

from ultralyticsplus import YOLO, render_result
import cv2
from PIL import Image

from cv2 import imshow
from cv2 import imwrite


#def flip(im):
#    return np.flipud(im)

def PPE(image):
    # load model
    #model = YOLO('keremberke/yolov8m-protective-equipment-detection')
    model = YOLO('keremberke/yolov8m-hard-hat-detection')
    # set model parameters
    model.overrides['conf'] = 0.25  # NMS confidence threshold
    model.overrides['iou'] = 0.45  # NMS IoU threshold
    model.overrides['agnostic_nms'] = False  # NMS class-agnostic
    model.overrides['max_det'] = 1000  # maximum number of detections per image



    # perform inference
    results = model.predict(image)

    # observe results
    print(results[0].boxes)
    render = render_result(model=model, image=image, result=results[0])
    render.show()






demo = gr.Interface(
    PPE, 
    gr.Image(source="webcam", streaming=True), 
    "image",
    live=True
)
demo.launch()