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