20065320rp commited on
Commit
b5af2f3
1 Parent(s): a4d5a9a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +82 -0
app.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import cv2
3
+
4
+ from ultralyticsplus import YOLO, render_result
5
+
6
+ # load model
7
+ model = YOLO('keremberke/yolov8m-hard-hat-detection')
8
+
9
+ # set model parameters
10
+ model.overrides['conf'] = 0.25 # NMS confidence threshold
11
+ model.overrides['iou'] = 0.45 # NMS IoU threshold
12
+ model.overrides['agnostic_nms'] = False # NMS class-agnostic
13
+ model.overrides['max_det'] = 1000 # maximum number of detections per image
14
+
15
+
16
+
17
+ from IPython.display import display, Javascript
18
+ from google.colab.output import eval_js
19
+ from base64 import b64decode
20
+
21
+ def take_photo(filename='photo.jpg', quality=0.8):
22
+ js = Javascript('''
23
+ async function takePhoto(quality) {
24
+ const div = document.createElement('div');
25
+ const capture = document.createElement('button');
26
+ capture.textContent = 'Capture';
27
+ div.appendChild(capture);
28
+
29
+ const video = document.createElement('video');
30
+ video.style.display = 'block';
31
+ const stream = await navigator.mediaDevices.getUserMedia({video: true});
32
+
33
+ document.body.appendChild(div);
34
+ div.appendChild(video);
35
+ video.srcObject = stream;
36
+ await video.play();
37
+
38
+ // Resize the output to fit the video element.
39
+ google.colab.output.setIframeHeight(document.documentElement.scrollHeight, true);
40
+
41
+ // Wait for Capture to be clicked.
42
+ await new Promise((resolve) => capture.onclick = resolve);
43
+
44
+ const canvas = document.createElement('canvas');
45
+ canvas.width = video.videoWidth;
46
+ canvas.height = video.videoHeight;
47
+ canvas.getContext('2d').drawImage(video, 0, 0);
48
+ stream.getVideoTracks()[0].stop();
49
+ div.remove();
50
+ return canvas.toDataURL('image/jpeg', quality);
51
+ }
52
+ ''')
53
+ display(js)
54
+ data = eval_js('takePhoto({})'.format(quality))
55
+ binary = b64decode(data.split(',')[1])
56
+ with open(filename, 'wb') as f:
57
+ f.write(binary)
58
+ return filename
59
+
60
+
61
+
62
+ from ultralyticsplus import YOLO, render_result
63
+
64
+ def safety_helmet():
65
+
66
+ # perform inference
67
+ #results = model.predict(image_path)
68
+ results = model.predict(filename)
69
+
70
+ # self.assertEqual(array.shape, (1, 224, 224, 3))
71
+
72
+ # observe results
73
+ print(results[0].boxes)
74
+ render = render_result(model=model, image=filename, result=results[0])
75
+ render.show()
76
+
77
+
78
+ import gradio as gr
79
+
80
+ demo = gr.Interface(fn=safety_helmet, inputs=take_photo(), outputs=safety_helmet(), description="Safety Helmet Detection")
81
+
82
+ demo.launch(share=True)