-
Notifications
You must be signed in to change notification settings - Fork 0
/
video.py
125 lines (107 loc) · 4.13 KB
/
video.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
""" Detect people wearing masks in videos
"""
from pathlib import Path
import click
import cv2
import torch
from skvideo.io import FFmpegWriter, vreader
from torchvision.transforms import Compose, Resize, ToPILImage, ToTensor
from common.facedetector import FaceDetector
from train import MaskDetector
@click.command(help="""
modelPath: path to model.ckpt\n
videoPath: path to video file to annotate
""")
@click.argument('modelpath')
@click.argument('videopath')
@click.option('--output', 'outputPath', type=Path,
help='specify output path to save video with annotations')
@torch.no_grad()
def tagVideo(modelpath, videopath, outputPath=None):
""" detect if persons in video are wearing masks or not
"""
result=-1
model = MaskDetector()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.load_state_dict(torch.load(modelpath,map_location=device)['state_dict'], strict=False)
model = model.to(device)
model.eval()
faceDetector = FaceDetector(
prototype='models/deploy.prototxt.txt',
model='models/res10_300x300_ssd_iter_140000.caffemodel',
)
transformations = Compose([
ToPILImage(),
Resize((100, 100)),
ToTensor(),
])
if outputPath:
writer = FFmpegWriter(str(outputPath))
font = cv2.FONT_HERSHEY_SIMPLEX
#cv2.namedWindow('main', cv2.WINDOW_NORMAL)
labels = ['No mask', 'Mask']
labelColor = [(10, 0, 255), (10, 255, 0)]
for frame in vreader(str(videopath)):
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
faces = faceDetector.detect(frame)
print ("FRAME")
for face in faces:
xStart, yStart, width, height = face
print ("FACE",face)
# clamp coordinates that are outside of the image
xStart, yStart = max(xStart, 0), max(yStart, 0)
# Image is 640x640
print ("DIMS",xStart, yStart ,width,height)
right = min(xStart+width,639)
bottom = min(yStart+height,639)
print ("Right",xStart+width)
print ("Bottom",yStart+height)
area = width*height
inarea = (right-xStart)*(bottom-yStart)
print ("Area",area)
print ("inArea",inarea)
areaperc=inarea/area
#print ("areaperc",areaperc)
# predict mask label on extracted face
faceImg = frame[yStart:yStart+height, xStart:xStart+width]
output = model(transformations(faceImg).unsqueeze(0).to(device))
print ("OUTPUT",output)
_, predicted = torch.max(output.data, 1)
print ("result",_)
# draw face frame
cv2.rectangle(frame,
(xStart, yStart),
(xStart + width, yStart + height),
(126, 65, 64),
thickness=2)
# center text according to the face frame
textSize = cv2.getTextSize(labels[predicted], font, 1, 2)[0]
textX = xStart + width // 2 - textSize[0] // 2
# draw prediction label
if (areaperc > 0.75):
print (labels[predicted])
if predicted:
result=1
elif result == -1:
result=0
else:
print ("Insufficuent")
cv2.putText(frame,
labels[predicted],
(textX, yStart-20),
font, 1, labelColor[predicted], 2)
if outputPath:
writer.writeFrame(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
#cv2.imshow('main', frame)
#if cv2.waitKey(1) & 0xFF == ord('q'):
# break
if outputPath:
writer.close()
#cv2.destroyAllWindows()
if result==0: print ("No Mask Found")
if result==1: print ("Mask Found")
if result==-1: print ("No face in photo")
return result
# pylint: disable=no-value-for-parameter
if __name__ == '__main__':
tagVideo()