Split out first_face vs multi-face calls

This commit is contained in:
Andras Fekete
2023-06-19 13:21:11 -04:00
parent 8e9035f712
commit 19d039d145

View File

@@ -149,17 +149,14 @@ 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):
if self.first_face: faces = self.__get_faces(frame,max_num=0)
for face in faces:
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:
for i in range(len(faces) - 1, -1, -1): for i in range(len(faces) - 1, -1, -1):
sim = self.rec_app.compute_sim(rep_face[0], faces[i].embedding) sim = self.rec_app.compute_sim(rep_face[0], faces[i].embedding)
@@ -178,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)