diff --git a/DFLIMG/DFLJPG.py b/DFLIMG/DFLJPG.py index 4e2663f..1c6a7c5 100644 --- a/DFLIMG/DFLJPG.py +++ b/DFLIMG/DFLJPG.py @@ -245,6 +245,9 @@ class DFLJPG(object): return None def set_image_to_face_mat(self, image_to_face_mat): self.dfl_dict['image_to_face_mat'] = image_to_face_mat + def has_seg_ie_polys(self): + return self.dfl_dict.get('seg_ie_polys',None) is not None + def get_seg_ie_polys(self): d = self.dfl_dict.get('seg_ie_polys',None) if d is not None: @@ -266,6 +269,9 @@ class DFLJPG(object): self.dfl_dict['seg_ie_polys'] = seg_ie_polys + def has_xseg_mask(self): + return self.dfl_dict.get('xseg_mask',None) is not None + def get_xseg_mask(self): mask_buf = self.dfl_dict.get('xseg_mask',None) if mask_buf is None: diff --git a/main.py b/main.py index ce1540c..922726d 100644 --- a/main.py +++ b/main.py @@ -278,8 +278,7 @@ if __name__ == "__main__": p.set_defaults (func=process_xsegapply) - p = xseg_parser.add_parser( "remove", help="Remove XSeg from the extracted faces.") - + p = xseg_parser.add_parser( "remove", help="Remove applied XSeg masks from the extracted faces.") def process_xsegremove(arguments): osex.set_process_lowest_prio() from mainscripts import XSegUtil @@ -288,6 +287,15 @@ if __name__ == "__main__": p.set_defaults (func=process_xsegremove) + p = xseg_parser.add_parser( "remove_labels", help="Remove XSeg labels from the extracted faces.") + def process_xsegremovelabels(arguments): + osex.set_process_lowest_prio() + from mainscripts import XSegUtil + XSegUtil.remove_xseg_labels (Path(arguments.input_dir) ) + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") + p.set_defaults (func=process_xsegremovelabels) + + p = xseg_parser.add_parser( "fetch", help="Copies faces containing XSeg polygons in _xseg dir.") def process_xsegfetch(arguments): diff --git a/mainscripts/XSegUtil.py b/mainscripts/XSegUtil.py index e557e0b..9f9f893 100644 --- a/mainscripts/XSegUtil.py +++ b/mainscripts/XSegUtil.py @@ -54,20 +54,7 @@ def apply_xseg(input_path, model_path): dflimg.set_xseg_mask(mask) dflimg.save() -def remove_xseg(input_path): - if not input_path.exists(): - raise ValueError(f'{input_path} not found. Please ensure it exists.') - - images_paths = pathex.get_image_paths(input_path, return_Path_class=True) - - for filepath in io.progress_bar_generator(images_paths, "Processing"): - dflimg = DFLIMG.load(filepath) - if dflimg is None or not dflimg.has_data(): - io.log_info(f'{filepath} is not a DFLIMG') - continue - - dflimg.set_xseg_mask(None) - dflimg.save() + def fetch_xseg(input_path): if not input_path.exists(): @@ -93,4 +80,55 @@ def fetch_xseg(input_path): files_copied += 1 shutil.copy ( str(filepath), str(output_path / filepath.name) ) - io.log_info(f'Files copied: {files_copied}') \ No newline at end of file + io.log_info(f'Files copied: {files_copied}') + +def remove_xseg(input_path): + if not input_path.exists(): + raise ValueError(f'{input_path} not found. Please ensure it exists.') + + io.log_info(f'Processing folder {input_path}') + io.log_info('!!! WARNING : APPLIED XSEG MASKS WILL BE REMOVED FROM THE FRAMES !!!') + io.log_info('!!! WARNING : APPLIED XSEG MASKS WILL BE REMOVED FROM THE FRAMES !!!') + io.log_info('!!! WARNING : APPLIED XSEG MASKS WILL BE REMOVED FROM THE FRAMES !!!') + io.input_str('Press enter to continue.') + + images_paths = pathex.get_image_paths(input_path, return_Path_class=True) + + files_processed = 0 + for filepath in io.progress_bar_generator(images_paths, "Processing"): + dflimg = DFLIMG.load(filepath) + if dflimg is None or not dflimg.has_data(): + io.log_info(f'{filepath} is not a DFLIMG') + continue + + if dflimg.has_xseg_mask(): + dflimg.set_xseg_mask(None) + dflimg.save() + files_processed += 1 + io.log_info(f'Files processed: {files_processed}') + +def remove_xseg_labels(input_path): + if not input_path.exists(): + raise ValueError(f'{input_path} not found. Please ensure it exists.') + + io.log_info(f'Processing folder {input_path}') + io.log_info('!!! WARNING : LABELED XSEG POLYGONS WILL BE REMOVED FROM THE FRAMES !!!') + io.log_info('!!! WARNING : LABELED XSEG POLYGONS WILL BE REMOVED FROM THE FRAMES !!!') + io.log_info('!!! WARNING : LABELED XSEG POLYGONS WILL BE REMOVED FROM THE FRAMES !!!') + io.input_str('Press enter to continue.') + + images_paths = pathex.get_image_paths(input_path, return_Path_class=True) + + files_processed = 0 + for filepath in io.progress_bar_generator(images_paths, "Processing"): + dflimg = DFLIMG.load(filepath) + if dflimg is None or not dflimg.has_data(): + io.log_info(f'{filepath} is not a DFLIMG') + continue + + if dflimg.has_seg_ie_polys(): + dflimg.set_seg_ie_polys(None) + dflimg.save() + files_processed += 1 + + io.log_info(f'Files processed: {files_processed}') \ No newline at end of file