import torch
import cv2

from ultralyticsplus import YOLO, render_result

# load model
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



#from IPython.display import display, Javascript
#from google.colab.output import eval_js
#from base64 import b64decode

def take_photo(filename='photo.jpg', quality=0.8):

    # program to capture single image from webcam in python
  
    # importing OpenCV library
    from cv2 import imshow
    from cv2 import imwrite
  
    # initialize the camera
    # If you have multiple camera connected with 
    # current device, assign a value in cam_port 
    # variable according to that
    cam_port = 0
    cam = cv2.VideoCapture(cam_port)
  
    # reading the input using the camera
    result, filename = cam.read()
  
    # If image will detected without any error, 
    # show result
    if result:
  
    # showing result, it take frame name and image 
    # output
        imshow("photo.jpg", filename)
  
    # saving image in local storage
        imwrite("photo.jpg", filename)
  
    # If keyboard interrupt occurs, destroy image 
    # window
        cv2.waitKey(0)
        cv2.destroyAllWindows()
#        cv2.destroyWindow("photo.jpg")

  
# If captured image is corrupted, moving to else part
    else:
        print("No image detected. Please! try again")
       
        
  #with open(filename, 'wb') as f:
  #  f.write(binary)
    return filename

#take_photo()

#from ultralyticsplus import YOLO, render_result

def safety_helmet():

# perform inference
#results = model.predict(image_path)
    results = model.predict(filename)

#    self.assertEqual(array.shape, (1, 224, 224, 3))

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


import gradio as gr

demo = gr.Interface(fn=safety_helmet, inputs=take_photo(), outputs=safety_helmet(), description="Safety Helmet Detection")

demo.launch(share=True)