|
import torch |
|
import cv2 |
|
|
|
from ultralyticsplus import YOLO, render_result |
|
|
|
|
|
model = YOLO('keremberke/yolov8m-hard-hat-detection') |
|
|
|
|
|
model.overrides['conf'] = 0.25 |
|
model.overrides['iou'] = 0.45 |
|
model.overrides['agnostic_nms'] = False |
|
model.overrides['max_det'] = 1000 |
|
|
|
|
|
|
|
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): |
|
js = Javascript(''' |
|
async function takePhoto(quality) { |
|
const div = document.createElement('div'); |
|
const capture = document.createElement('button'); |
|
capture.textContent = 'Capture'; |
|
div.appendChild(capture); |
|
|
|
const video = document.createElement('video'); |
|
video.style.display = 'block'; |
|
const stream = await navigator.mediaDevices.getUserMedia({video: true}); |
|
|
|
document.body.appendChild(div); |
|
div.appendChild(video); |
|
video.srcObject = stream; |
|
await video.play(); |
|
|
|
// Resize the output to fit the video element. |
|
google.colab.output.setIframeHeight(document.documentElement.scrollHeight, true); |
|
|
|
// Wait for Capture to be clicked. |
|
await new Promise((resolve) => capture.onclick = resolve); |
|
|
|
const canvas = document.createElement('canvas'); |
|
canvas.width = video.videoWidth; |
|
canvas.height = video.videoHeight; |
|
canvas.getContext('2d').drawImage(video, 0, 0); |
|
stream.getVideoTracks()[0].stop(); |
|
div.remove(); |
|
return canvas.toDataURL('image/jpeg', quality); |
|
} |
|
''') |
|
display(js) |
|
data = eval_js('takePhoto({})'.format(quality)) |
|
binary = b64decode(data.split(',')[1]) |
|
with open(filename, 'wb') as f: |
|
f.write(binary) |
|
return filename |
|
|
|
|
|
|
|
from ultralyticsplus import YOLO, render_result |
|
|
|
def safety_helmet(): |
|
|
|
|
|
|
|
results = model.predict(filename) |
|
|
|
|
|
|
|
|
|
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) |