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

#image=[gr.Image(label="Input Image", source="webcam")


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

    return render


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



#def snap(image):
#    return np.flipud(image)


#iface = gr.Interface(PPE, gr.inputs.Image(source="webcam", tool=None), "image")
#iface.launch()


#demo = gr.Interface(
#    fn=PPE,
#    inputs=gr.inputs.Image(source="webcam", tool=None),
#    outputs="image",
#)
#demo.launch()

demo = gr.Interface(
    fn=PPE,
    inputs="image",
    outputs="image",
)
demo.launch()