forked from viccy5022/ReHiFace-S-
-
Notifications
You must be signed in to change notification settings - Fork 3
/
inference_cam.py
114 lines (92 loc) · 3.42 KB
/
inference_cam.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
import pickle
from multiprocessing.dummy import Process, Manager, Queue
import cv2
import numpy as np
# close onnxruntime warning
import onnxruntime
onnxruntime.set_default_logger_severity(3)
class GenInput(Process):
def __init__(self, feature_src_list_, frame_queue_in_, frame_queue_out_):
super().__init__()
self.frame_queue_in = frame_queue_in_
self.frame_queue_out = frame_queue_out_
self.feature_src_list = feature_src_list_
def run(self):
with open('data/image_feature_dict.pkl', 'rb') as f:
image_feature_src_dict = pickle.load(f)
print(len(image_feature_src_dict))
self.feature_src_list.append([image_feature_src_dict['1']])
cap = cv2.VideoCapture(0) # 640 480 1280 720 1920 1080
cap.set(3, 1920)
cap.set(4, 1080)
print(cv2.CAP_PROP_FOURCC, cv2.CAP_PROP_FPS)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))
cap.set(cv2.CAP_PROP_FPS, 30)
print('CAP_PROP_FPS',cap.get(cv2.CAP_PROP_FPS))
count = index = 0
while True:
_, frame = cap.read()
self.frame_queue_in.put(frame)
count += 1
if count % 500 == 0:
self.feature_src_list.append([image_feature_src_dict['{}'.format(1 + index)],
image_feature_src_dict['{}'.format(10 + index)]])
print('change src face')
index += 1
if count % 5000 == 0:
# 退出条件
self.frame_queue_in.put(None)
break
class GetOutput(Process):
def __init__(self, frame_queue_out_):
super().__init__()
self.frame_queue_out = frame_queue_out_
def run(self):
import time
count = 0
fps_count = 0
start_time = time.time()
while True:
queue_out = self.frame_queue_out.get()
# print(queue_out)
frame_out = queue_out[0]
face_detect_flag = queue_out[1]
# print(face_detect_flag)
fps_count += 1
if fps_count % 300 == 0:
end_time = time.time()
print('fps: {}'.format(fps_count / (end_time - start_time)))
start_time = time.time()
fps_count = 0
count += 1
if count % 2500 == 0:
break
cv2.imshow('output', frame_out)
cv2.waitKey(1)
class FaceSwap(Process):
def __init__(self, feature_src_list_, frame_queue_in_,
frame_queue_out_, model_name=''):
super().__init__()
from HifiFaceAPI_parallel_trt_roi_realtime_api import HifiFaceRealTime
self.hfrt = HifiFaceRealTime(feature_src_list_, frame_queue_in_,
frame_queue_out_, model_name=model_name)
def run(self):
self.hfrt.forward()
if __name__ == '__main__':
frame_queue_in = Queue(2)
frame_queue_out = Queue(2)
manager = Manager()
image_feature_src_list = manager.list()
gi = GenInput(image_feature_src_list, frame_queue_in, frame_queue_out)
go = GetOutput(frame_queue_out)
fs = FaceSwap(image_feature_src_list, frame_queue_in, frame_queue_out, model_name='er8_bs1')
gi.start()
go.start()
fs.start()
gi.join()
print('gi stop')
go.join()
print('go stop')
fs.join()
print('fs stop')
print('final stop')