github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-07 05:22:06 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

Extractor: extract from dflimg is now the same as from any image

Browse source
This commit is contained in:
iperov 2021-04-20 16:45:20 +04:00
parent 93fe480eca
commit 7a08c0c1d3
1 changed files with 40 additions and 66 deletions
Download patch file Download diff file Expand all files Collapse all files

28
mainscripts/Extractor.py
View file

@ -97,9 +97,6 @@ class ExtractSubprocessor(Subprocessor):
h, w, c = image.shape h, w, c = image.shape
dflimg = DFLIMG.load (filepath)
extract_from_dflimg = (h == w and (dflimg is not None and dflimg.has_data()) )
if 'rects' in self.type or self.type == 'all': if 'rects' in self.type or self.type == 'all':
data = ExtractSubprocessor.Cli.rects_stage (data=data, data = ExtractSubprocessor.Cli.rects_stage (data=data,
image=image, image=image,
@ -110,7 +107,6 @@ class ExtractSubprocessor(Subprocessor):
if 'landmarks' in self.type or self.type == 'all': if 'landmarks' in self.type or self.type == 'all':
data = ExtractSubprocessor.Cli.landmarks_stage (data=data, data = ExtractSubprocessor.Cli.landmarks_stage (data=data,
image=image, image=image,
extract_from_dflimg=extract_from_dflimg,
landmarks_extractor=self.landmarks_extractor, landmarks_extractor=self.landmarks_extractor,
rects_extractor=self.rects_extractor, rects_extractor=self.rects_extractor,
) )
@ -121,7 +117,6 @@ class ExtractSubprocessor(Subprocessor):
face_type=self.face_type, face_type=self.face_type,
image_size=self.image_size, image_size=self.image_size,
jpeg_quality=self.jpeg_quality, jpeg_quality=self.jpeg_quality,
extract_from_dflimg=extract_from_dflimg,
output_debug_path=self.output_debug_path, output_debug_path=self.output_debug_path,
final_output_path=self.final_output_path, final_output_path=self.final_output_path,
) )
@ -161,7 +156,6 @@ class ExtractSubprocessor(Subprocessor):
@staticmethod @staticmethod
def landmarks_stage(data, def landmarks_stage(data,
image, image,
extract_from_dflimg,
landmarks_extractor, landmarks_extractor,
rects_extractor, rects_extractor,
): ):
@ -176,7 +170,7 @@ class ExtractSubprocessor(Subprocessor):
elif data.rects_rotation == 270: elif data.rects_rotation == 270:
rotated_image = image.swapaxes( 0,1 )[::-1,:,:] rotated_image = image.swapaxes( 0,1 )[::-1,:,:]
data.landmarks = landmarks_extractor.extract (rotated_image, data.rects, rects_extractor if (not extract_from_dflimg and data.landmarks_accurate) else None, is_bgr=True) data.landmarks = landmarks_extractor.extract (rotated_image, data.rects, rects_extractor if (data.landmarks_accurate) else None, is_bgr=True)
if data.rects_rotation != 0: if data.rects_rotation != 0:
for i, (rect, lmrks) in enumerate(zip(data.rects, data.landmarks)): for i, (rect, lmrks) in enumerate(zip(data.rects, data.landmarks)):
new_rect, new_lmrks = rect, lmrks new_rect, new_lmrks = rect, lmrks
@ -207,7 +201,6 @@ class ExtractSubprocessor(Subprocessor):
face_type, face_type,
image_size, image_size,
jpeg_quality, jpeg_quality,
extract_from_dflimg = False,
output_debug_path=None, output_debug_path=None,
final_output_path=None, final_output_path=None,
): ):
@ -219,21 +212,8 @@ class ExtractSubprocessor(Subprocessor):
if output_debug_path is not None: if output_debug_path is not None:
debug_image = image.copy() debug_image = image.copy()
if extract_from_dflimg and len(rects) != 1:
#if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it
print("extract_from_dflimg and len(rects) != 1", filepath )
output_filepath = final_output_path / filepath.name
if filepath != str(output_file):
shutil.copy ( str(filepath), str(output_filepath) )
data.final_output_files.append (output_filepath)
else:
face_idx = 0 face_idx = 0
for rect, image_landmarks in zip( rects, landmarks ): for rect, image_landmarks in zip( rects, landmarks ):
if extract_from_dflimg and face_idx > 1:
#cannot extract more than 1 face from dflimg
break
if image_landmarks is None: if image_landmarks is None:
continue continue
@ -264,12 +244,6 @@ class ExtractSubprocessor(Subprocessor):
if data.force_output_path is not None: if data.force_output_path is not None:
output_path = data.force_output_path output_path = data.force_output_path
if extract_from_dflimg and filepath.suffix == '.jpg':
#if extracting from dflimg and jpg copy it in order not to lose quality
output_filepath = output_path / filepath.name
if filepath != output_filepath:
shutil.copy ( str(filepath), str(output_filepath) )
else:
output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg" output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg"
cv2_imwrite(output_filepath, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality ] ) cv2_imwrite(output_filepath, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality ] )