-
Notifications
You must be signed in to change notification settings - Fork 0
/
demo.py
68 lines (56 loc) · 2.12 KB
/
demo.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
from facenet_pytorch import InceptionResnetV1, fixed_image_standardization, MTCNN
import torch
from torch.utils.data import DataLoader, SubsetRandomSampler
import numpy as np
import os
import argparse
from copy import deepcopy
import yaml
import cv2
from utils import softmax
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
def run(opt):
with open(opt.names, 'r') as f:
names = yaml.load(f, Loader=yaml.FullLoader)
detector = MTCNN(device=device, post_process=True)
classifier = InceptionResnetV1(classify=True, num_classes=len(names))
classifier.load_state_dict(torch.load(opt.model))
classifier.eval()
try:
source = int(opt.source)
except:
source = opt.source
cap = cv2.VideoCapture(source)
while True:
ret, frame = cap.read()
if not ret:
cap.release()
break
faces, _ = detector.detect(frame[:, :, [2,1,0]])
if not isinstance(faces, type(None)):
for i, face in enumerate(faces):
x1, y1, x2, y2 = list(map(int, face))
img = pre_processing(cv2.resize(frame[y1:y2, x1:x2, [2,1,0]], (opt.size, opt.size), interpolation=cv2.INTER_CUBIC))
result = softmax(classifier(img).detach().cpu().numpy())
if opt.target == np.argmax(result):
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255))
cv2.imshow("frame", frame)
key = cv2.waitKey(1)
if key == ord("q"):
break
def pre_processing(image):
image = np.transpose(image, (2, 0, 1))
image = np.expand_dims(image.astype(np.float32), axis=0)
image = fixed_image_standardization(torch.tensor(image))
return image
def parse():
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, required=True)
parser.add_argument("--names", type=str, required=True)
parser.add_argument("--size", default=256, type=int)
parser.add_argument("--source", type=str, default="0")
parser.add_argument("--target", type=int, default=1)
return parser.parse_args()
if __name__ == "__main__":
opt = parse()
run(opt)