fix extractor

This commit is contained in:
Colombo 2020-07-03 16:06:06 +04:00
parent 4ce4997d1a
commit 57ea3e61b7

View file

@ -151,7 +151,9 @@ class ExtractSubprocessor(Subprocessor):
if len(rects) != 0:
data.rects_rotation = rot
break
if max_faces_from_image is not None and len(data.rects) > 1:
if max_faces_from_image is not None and \
max_faces_from_image > 0 and \
len(data.rects) > 0:
data.rects = data.rects[0:max_faces_from_image]
return data
@ -371,7 +373,7 @@ class ExtractSubprocessor(Subprocessor):
self.rect_size = 100
self.rect_locked = False
self.extract_needed = True
self.image = None
self.image_filepath = None
@ -411,7 +413,7 @@ class ExtractSubprocessor(Subprocessor):
data = self.input_data[0]
filepath, data_rects, data_landmarks = data.filepath, data.rects, data.landmarks
is_frame_done = False
if self.image_filepath != filepath:
self.image_filepath = filepath
if self.cache_original_image[0] == filepath:
@ -446,7 +448,7 @@ class ExtractSubprocessor(Subprocessor):
], (1, 1, 1) )*255).astype(np.uint8)
self.cache_text_lines_img = (sh, self.text_lines_img)
if need_remark_face: # need remark image from input data that already has a marked face?
need_remark_face = False
if len(data_rects) != 0: # If there was already a face then lock the rectangle to it until the mouse is clicked
@ -761,15 +763,15 @@ def main(detector=None,
if dflimg is not None and dflimg.has_data():
face_type = FaceType.fromString ( dflimg.get_face_type() )
input_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
input_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
output_images_paths = pathex.get_image_paths(output_path)
output_debug_path = output_path.parent / (output_path.name + '_debug')
continue_extraction = False
if not manual_output_debug_fix and len(output_images_paths) > 0:
if len(output_images_paths) > 128:
if not manual_output_debug_fix and len(output_images_paths) > 0:
if len(output_images_paths) > 128:
continue_extraction = io.input_bool ("Continue extraction?", True, help_message="Extraction can be continued, but you must specify the same options again.")
if len(output_images_paths) > 128 and continue_extraction:
try:
input_image_paths = input_image_paths[ [ Path(x).stem for x in input_image_paths ].index ( Path(output_images_paths[-128]).stem.split('_')[0] ) : ]
@ -805,10 +807,10 @@ def main(detector=None,
io.log_info ("[1] manual")
detector = {0:'s3fd', 1:'manual'}[ io.input_int("", 0, [0,1]) ]
if output_debug is None:
output_debug = io.input_bool (f"Write debug images to {output_debug_path.name}?", False)
if output_debug:
output_debug_path.mkdir(parents=True, exist_ok=True)