Merge pull request #46 from bandi13/returnWhenFaceFound
Return when face found
This commit is contained in:
29
refacer.py
29
refacer.py
@@ -61,7 +61,7 @@ class Refacer:
|
|||||||
print(f"CoreML mode with providers {self.providers}")
|
print(f"CoreML mode with providers {self.providers}")
|
||||||
elif 'CUDAExecutionProvider' in self.providers:
|
elif 'CUDAExecutionProvider' in self.providers:
|
||||||
self.mode = RefacerMode.CUDA
|
self.mode = RefacerMode.CUDA
|
||||||
self.use_num_cpus = 1
|
self.use_num_cpus = 2
|
||||||
self.sess_options.intra_op_num_threads = 1
|
self.sess_options.intra_op_num_threads = 1
|
||||||
if 'TensorrtExecutionProvider' in self.providers:
|
if 'TensorrtExecutionProvider' in self.providers:
|
||||||
self.providers.remove('TensorrtExecutionProvider')
|
self.providers.remove('TensorrtExecutionProvider')
|
||||||
@@ -149,21 +149,21 @@ class Refacer:
|
|||||||
ret.append(face)
|
ret.append(face)
|
||||||
return ret
|
return ret
|
||||||
|
|
||||||
def process_faces(self,frame):
|
def process_first_face(self,frame):
|
||||||
max_num=0
|
faces = self.__get_faces(frame,max_num=1)
|
||||||
if self.first_face:
|
if len(faces) != 0:
|
||||||
max_num=1
|
frame = self.face_swapper.get(frame, faces[0], self.replacement_faces[0][1], paste_back=True)
|
||||||
|
return frame
|
||||||
|
|
||||||
faces = self.__get_faces(frame,max_num=max_num)
|
def process_faces(self,frame):
|
||||||
for face in faces:
|
faces = self.__get_faces(frame,max_num=0)
|
||||||
if self.first_face:
|
|
||||||
frame = self.face_swapper.get(frame, face, self.replacement_faces[0][1], paste_back=True)
|
|
||||||
break
|
|
||||||
else:
|
|
||||||
for rep_face in self.replacement_faces:
|
for rep_face in self.replacement_faces:
|
||||||
sim = self.rec_app.compute_sim(rep_face[0], face.embedding)
|
for i in range(len(faces) - 1, -1, -1):
|
||||||
|
sim = self.rec_app.compute_sim(rep_face[0], faces[i].embedding)
|
||||||
if sim>=rep_face[2]:
|
if sim>=rep_face[2]:
|
||||||
frame = self.face_swapper.get(frame, face, rep_face[1], paste_back=True)
|
frame = self.face_swapper.get(frame, faces[i], rep_face[1], paste_back=True)
|
||||||
|
del faces[i]
|
||||||
|
break
|
||||||
return frame
|
return frame
|
||||||
|
|
||||||
def __check_video_has_audio(self,video_path):
|
def __check_video_has_audio(self,video_path):
|
||||||
@@ -175,6 +175,9 @@ class Refacer:
|
|||||||
|
|
||||||
def reface_group(self, faces, frames, output):
|
def reface_group(self, faces, frames, output):
|
||||||
with ThreadPoolExecutor(max_workers = self.use_num_cpus) as executor:
|
with ThreadPoolExecutor(max_workers = self.use_num_cpus) as executor:
|
||||||
|
if self.first_face:
|
||||||
|
results = list(tqdm(executor.map(self.process_first_face, frames), total=len(frames),desc="Processing frames"))
|
||||||
|
else:
|
||||||
results = list(tqdm(executor.map(self.process_faces, frames), total=len(frames),desc="Processing frames"))
|
results = list(tqdm(executor.map(self.process_faces, frames), total=len(frames),desc="Processing frames"))
|
||||||
for result in results:
|
for result in results:
|
||||||
output.write(result)
|
output.write(result)
|
||||||
|
|||||||
Reference in New Issue
Block a user