def sort_by_blur(input_path): io.log_info("Sorting by blur...") img_list = [(filename, []) for filename in Path_utils.get_image_paths(input_path)] img_list, trash_img_list = BlurEstimatorSubprocessor(img_list).run() io.log_info("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) return img_list, trash_img_list
def refix(workspace): dst = get_workspace_dst(workspace) dst_aligned = os.path.join(dst, "aligned") recover_filename_if_nessesary(dst_aligned) extract_imgs = [ f if f.endswith(".jpg") or f.endswith(".png") else "" for f in os.listdir(dst) ] max_img_no = int(max(extract_imgs).split(".")[0]) ext = extract_imgs[0].split(".")[1] aligned_imgs = list( sorted( filter(lambda x: x is not None, [ f if f.endswith(".jpg") or f.endswith(".png") else None for f in os.listdir(dst_aligned) ]))) need_fix_no = [] i = 0 # 当前文件下标 j = 1 # 期望文件名 while i <= max_img_no and j <= max_img_no and i < len(aligned_imgs): if aligned_imgs[i].startswith("%05d" % j): i += 1 j += 1 else: # print(aligned_imgs[i], j) need_fix_no.append(j) j += 1 for k in range(j, max_img_no + 1): need_fix_no.append(k) # print(k) if len(need_fix_no) == 0: return fix_workspace = os.path.join(dst, "fix") import shutil if os.path.exists(fix_workspace): shutil.rmtree(fix_workspace) os.mkdir(fix_workspace) for no in need_fix_no: f = os.path.join(dst, "%05d.%s" % (no, ext)) io.log_info(f) shutil.copy(f, fix_workspace) fix_workspace_aligned = os.path.join(fix_workspace, "aligned") from mainscripts import Extractor Extractor.main(fix_workspace, fix_workspace_aligned, detector="manual", manual_fix=False) Extractor.extract_fanseg(fix_workspace_aligned) for f in os.listdir(fix_workspace_aligned): f = os.path.join(fix_workspace_aligned, f) io.log_info(f) shutil.move(f, dst_aligned) shutil.rmtree(fix_workspace)
def sort_by_hue(input_path): io.log_info("Sorting by hue...") img_list = [[ x, np.mean( cv2.cvtColor(cv2_imread(x), cv2.COLOR_BGR2HSV)[..., 0].flatten()) ] for x in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading")] io.log_info("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) return img_list
def recover_original_aligned_filename(input_path): io.log_info ("Recovering original aligned filename...") files = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"): filepath = Path(filepath) dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue files += [ [filepath, None, dflimg.get_source_filename(), False] ] files_len = len(files) for i in io.progress_bar_generator( range(files_len), "Sorting" ): fp, _, sf, converted = files[i] if converted: continue sf_stem = Path(sf).stem files[i][1] = fp.parent / ( sf_stem + '_0' + fp.suffix ) files[i][3] = True c = 1 for j in range(i+1, files_len): fp_j, _, sf_j, converted_j = files[j] if converted_j: continue if sf_j == sf: files[j][1] = fp_j.parent / ( sf_stem + ('_%d' % (c)) + fp_j.suffix ) files[j][3] = True c += 1 for file in io.progress_bar_generator( files, "Renaming", leave=False ): fs, _, _, _ = file dst = fs.parent / ( fs.stem + '_tmp' + fs.suffix ) try: fs.rename (dst) except: io.log_err ('fail to rename %s' % (fs.name) ) for file in io.progress_bar_generator( files, "Renaming" ): fs, fd, _, _ = file fs = fs.parent / ( fs.stem + '_tmp' + fs.suffix ) try: fs.rename (fd) except: io.log_err ('fail to rename %s' % (fs.name) )
def auto(workspace): import subprocess for f in os.listdir(workspace): if os.path.isdir(os.path.join(get_root_path(), "workspace", f)) and f.startswith("data_dst_"): train_bat = os.path.join(get_root_path(), "auto_train.bat") convert_bat = os.path.join(get_root_path(), "auto_convert.bat") step_bat = os.path.join(get_root_path(), "auto_step.bat") subprocess.call([train_bat]) subprocess.call([convert_bat]) subprocess.call([step_bat]) io.log_info("Finish " + f)
def sort_by_black(input_path): io.log_info ("Sorting by amount of black pixels...") img_list = [] for x in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): img = cv2_imread(x) img_list.append ([x, img[(img == 0)].size ]) io.log_info ("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1), reverse=False) return img_list
def sync_trash(trash_path, pool_path): import shutil count = 0 for f in io.progress_bar_generator(os.listdir(trash_path), "Trash Files"): if f.endswith(".jpg") or f.endswith(".png"): img_name = f.split("_")[-1] img_path = os.path.join(pool_path, img_name) dst_path = os.path.join(trash_path, "_origin") if os.path.exists(img_path): shutil.move(img_path, dst_path) count += 1 io.log_info("Trash %d" % count)
def get_pitch_yaw_roll(input_path, r=0.05): import os import numpy as np import cv2 from shutil import copyfile from pathlib import Path from utils import Path_utils from utils.DFLPNG import DFLPNG from utils.DFLJPG import DFLJPG from facelib import LandmarksProcessor from joblib import Subprocessor import multiprocessing from interact import interact as io from imagelib import estimate_sharpness io.log_info("Sorting by face yaw...") img_list = [] trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) if filepath.suffix == '.png': dflimg = DFLPNG.load(str(filepath)) elif filepath.suffix == '.jpg': dflimg = DFLJPG.load(str(filepath)) else: dflimg = None if dflimg is None: io.log_err("%s is not a dfl image file" % (filepath.name)) trash_img_list.append([str(filepath)]) continue pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll( dflimg.get_landmarks()) img_list.append([str(filepath), pitch, yaw, roll]) img_list.sort(key=lambda item: item[1]) with open(os.path.join(input_path, "_pitch_yaw_roll.csv"), "w") as f: for i in img_list: f.write("%s,%f,%f,%f\n" % (os.path.basename(i[0]), i[1], i[2], i[3])) import cv width = 800 img = cv.cv_new((width, width)) xs = [i[1] for i in img_list] ys = [i[2] for i in img_list] cs = [(128, 128, 128)] * len(xs) rs = [int(r * width / 2)] * len(xs) cv.cv_scatter(img, xs, ys, [-1, 1], [-1, 1], cs, rs) cs = [(0xcc, 0x66, 0x33)] * len(xs) rs = [2] * len(xs) cv.cv_scatter(img, xs, ys, [-1, 1], [-1, 1], cs, rs) cv.cv_save(img, os.path.join(input_path, "_pitch_yaw_roll.bmp")) return img_list
def unpack(samples_path): samples_dat_path = samples_path / packed_faceset_filename if not samples_dat_path.exists(): io.log_info(f"{samples_dat_path} : file not found.") return samples = PackedFaceset.load(samples_path) for sample in io.progress_bar_generator(samples, "Unpacking"): with open(samples_path / sample.filename, "wb") as f: f.write( sample.read_raw_file() ) samples_dat_path.unlink()
def save(self): io.log_info ("Saving...") Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text) self.onSave() model_data = { 'epoch': self.epoch, 'options': self.options, 'loss_history': self.loss_history, 'sample_for_preview' : self.sample_for_preview } self.model_data_path.write_bytes( pickle.dumps(model_data) )
def pack(samples_path): samples_dat_path = samples_path / packed_faceset_filename if samples_dat_path.exists(): io.log_info(f"{samples_dat_path} : file already exists !") io.input_bool("Press enter to continue and overwrite.", False) of = open(samples_dat_path, "wb") image_paths = Path_utils.get_image_paths(samples_path) samples = samplelib.SampleHost.load_face_samples(image_paths) samples_len = len(samples) samples_configs = [] for sample in samples: sample.filename = str(Path(sample.filename).relative_to(samples_path)) samples_configs.append ( sample.get_config() ) samples_bytes = pickle.dumps(samples_configs, 4) of.write ( struct.pack ("Q", PackedFaceset.VERSION ) ) of.write ( struct.pack ("Q", len(samples_bytes) ) ) of.write ( samples_bytes ) sample_data_table_offset = of.tell() of.write ( bytes( 8*(samples_len+1) ) ) #sample data offset table data_start_offset = of.tell() offsets = [] for sample in io.progress_bar_generator(samples, "Packing"): try: with open( samples_path / sample.filename, "rb") as f: b = f.read() offsets.append ( of.tell() - data_start_offset ) of.write(b) except: raise Exception(f"error while processing sample {sample.filename}") offsets.append ( of.tell() ) of.seek(sample_data_table_offset, 0) for offset in offsets: of.write ( struct.pack("Q", offset) ) of.seek(0,2) of.close() for filename in io.progress_bar_generator(image_paths,"Deleting"): Path(filename).unlink()
def ask_settings(self): s = """Choose mode: \n""" for key in self.mode_dict.keys(): s += f"""({key}) {self.mode_dict[key]}\n""" s += f"""Default: {self.default_mode} : """ mode = io.input_int (s, self.default_mode) self.mode = self.mode_dict.get (mode, self.mode_dict[self.default_mode] ) if 'raw' not in self.mode: if self.mode == 'hist-match' or self.mode == 'hist-match-bw': self.masked_hist_match = io.input_bool("Masked hist match? (y/n skip:y) : ", True) if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match': self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255) if self.face_type == FaceType.FULL: s = """Choose mask mode: \n""" for key in self.full_face_mask_mode_dict.keys(): s += f"""({key}) {self.full_face_mask_mode_dict[key]}\n""" s += f"""?:help Default: 1 : """ self.mask_mode = io.input_int (s, 1, valid_list=self.full_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks.") else: s = """Choose mask mode: \n""" for key in self.half_face_mask_mode_dict.keys(): s += f"""({key}) {self.half_face_mask_mode_dict[key]}\n""" s += f"""?:help , Default: 1 : """ self.mask_mode = io.input_int (s, 1, valid_list=self.half_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images.") if 'raw' not in self.mode: self.erode_mask_modifier = self.base_erode_mask_modifier + np.clip ( io.input_int ("Choose erode mask modifier [-200..200] (skip:%d) : " % (self.default_erode_mask_modifier), self.default_erode_mask_modifier), -200, 200) self.blur_mask_modifier = self.base_blur_mask_modifier + np.clip ( io.input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (self.default_blur_mask_modifier), self.default_blur_mask_modifier), -200, 200) self.motion_blur_power = np.clip ( io.input_int ("Choose motion blur power [0..100] (skip:%d) : " % (0), 0), 0, 100) self.output_face_scale = np.clip (io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0), -50, 50) if 'raw' not in self.mode: self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct']) self.color_transfer_mode = self.ctm_str_dict[self.color_transfer_mode] super().ask_settings() if 'raw' not in self.mode: self.color_degrade_power = np.clip ( io.input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100) self.export_mask_alpha = io.input_bool("Export png with alpha channel of the mask? (y/n skip:n) : ", False) io.log_info ("")
def process_folder(dirpath, multi_gpu=False, cpu_only=False): output_dirpath = dirpath.parent / (dirpath.name + '_enhanced') output_dirpath.mkdir(exist_ok=True, parents=True) dirpath_parts = '/'.join(dirpath.parts[-2:]) output_dirpath_parts = '/'.join(output_dirpath.parts[-2:]) io.log_info(f"Enhancing faceset in {dirpath_parts}") io.log_info(f"Processing to {output_dirpath_parts}") output_images_paths = Path_utils.get_image_paths(output_dirpath) if len(output_images_paths) > 0: for filename in output_images_paths: Path(filename).unlink() image_paths = [Path(x) for x in Path_utils.get_image_paths(dirpath)] result = FacesetEnhancerSubprocessor(image_paths, output_dirpath, multi_gpu=multi_gpu, cpu_only=cpu_only).run() is_merge = io.input_bool( f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ? (y/n skip:y) : ", True) if is_merge: io.log_info(f"Copying processed files to {dirpath_parts}") for (filepath, output_filepath) in result: try: shutil.copy(output_filepath, filepath) except: pass io.log_info(f"Removing {output_dirpath_parts}") shutil.rmtree(output_dirpath)
def auto_extract_to_img(): workspace = os.path.join(get_root_path(), "workspace") data_dst = None for f in os.listdir(workspace): if f.startswith("data_dst_"): data_dst = f break io.log_info(data_dst) video_name = None if data_dst is not None: name = "_".join(data_dst.split("_")[8:]) print(name) for f in os.listdir(os.path.join(workspace, "data_trash")): if f.startswith(name): video_name = f break io.log_info(video_name) if video_name is not None: video_path = os.path.join(workspace, "data_trash", video_name) data_dst_path = os.path.join(workspace, data_dst) io.log_info(video_path) io.log_info(data_dst_path) for f in io.progress_bar_generator(os.listdir(data_dst_path), "Remove"): if f.endswith(".jpg") or f.endswith(".png"): os.remove(os.path.join(data_dst_path, f)) dfl.dfl_extract_video(video_path, data_dst_path)
def final_process(input_path, img_list, trash_img_list): if len(trash_img_list) != 0: parent_input_path = input_path.parent trash_path = parent_input_path / (input_path.stem + '_trash') trash_path.mkdir(exist_ok=True) io.log_info("Trashing %d items to %s" % (len(trash_img_list), str(trash_path))) for filename in Path_utils.get_image_paths(trash_path): Path(filename).unlink() for i in io.progress_bar_generator(range(len(trash_img_list)), "Moving trash", leave=False): src = Path(trash_img_list[i][0]) dst = trash_path / src.name try: src.rename(dst) except: io.log_info('fail to trashing %s' % (src.name)) io.log_info("") if len(img_list) != 0: for i in io.progress_bar_generator([*range(len(img_list))], "Renaming", leave=False): src = Path(img_list[i][0]) dst = "%s/%.5d_%s" % (input_path, i, src.name) # dst = input_path / ('%.5d_%s' % (i, src.name )) try: src.rename(dst) except: io.log_info('fail to rename %s' % (src.name)) for i in io.progress_bar_generator([*range(len(img_list))], "Renaming"): src = Path(img_list[i][0]) src = "%s/%.5d_%s" % (input_path, i, src.name) src = Path(src) dst = "%s/%.5d%s" % (input_path, i, src.suffix) # src = input_path / ('%.5d_%s' % (i, src.name)) # dst = input_path / ('%.5d%s' % (i, src.suffix)) try: src.rename(dst) except: io.log_info('fail to rename %s' % (src.name))
def sort_by_oneface_in_image(input_path): io.log_info ("Sort by one face in images...") image_paths = Path_utils.get_image_paths(input_path) a = np.array ([ ( int(x[0]), int(x[1]) ) \ for x in [ Path(filepath).stem.split('_') for filepath in image_paths ] if len(x) == 2 ]) if len(a) > 0: idxs = np.ndarray.flatten ( np.argwhere ( a[:,1] != 0 ) ) idxs = np.unique ( a[idxs][:,0] ) idxs = np.ndarray.flatten ( np.argwhere ( np.array([ x[0] in idxs for x in a ]) == True ) ) if len(idxs) > 0: io.log_info ("Found %d images." % (len(idxs)) ) img_list = [ (path,) for i,path in enumerate(image_paths) if i not in idxs ] trash_img_list = [ (image_paths[x],) for x in idxs ] return img_list, trash_img_list return [], []
def train(workspace, target_loss=0.01): import os model_dir = os.path.join(workspace, "model") for f in os.listdir(workspace): if not os.path.isdir(os.path.join(workspace, f)) or not f.startswith("data_dst_"): continue io.log_info(f) data_dst = os.path.join(workspace, f) data_src_aligned = os.path.join(data_dst, "src") if not os.path.exists(data_src_aligned): data_src_aligned = os.path.join(workspace, "data_src", "aligned") data_dst_aligned = os.path.join(data_dst, "aligned") # 训练 dfl.dfl_train(data_src_aligned, data_dst_aligned, model_dir) return
def mp4(workspace, skip=False): import os for f in os.listdir(workspace): if not os.path.isdir(os.path.join(workspace, f)) or not f.startswith("data_dst_"): continue io.log_info(f) data_dst = os.path.join(workspace, f) data_dst_merged = os.path.join(data_dst, "merged") data_dst_aligned = os.path.join(data_dst, "aligned") data_dst_video = os.path.join(data_dst, "video") refer_path = None for v in os.listdir(data_dst_video): if v.split(".")[-1] in ["mp4", "avi", "wmv", "mkv"]: refer_path = os.path.join(data_dst_video, v) break if not refer_path: io.log_err("No Refer File In " + data_dst_video) return io.log_info("Refer File " + refer_path) # 恢复排序 need_recover = True for img in os.listdir(data_dst_aligned): if img.endswith("_0.jpg") or img.endswith("_0.png"): need_recover = False if need_recover: recover_filename(data_dst_aligned) # 如果data_dst里没有脸则extract has_img = False for img in os.listdir(data_dst): if img.endswith(".jpg") or img.endswith(".png"): has_img = True break if not has_img: dfl.dfl_extract_video(refer_path, data_dst) # 去掉没有脸的 # if skip: # skip_no_face(data_dst) # 转mp4 refer_name = ".".join(os.path.basename(refer_path).split(".")[:-1]) result_path = os.path.join( workspace, "result_%s_%s.mp4" % (get_time_str(), refer_name)) dfl.dfl_video_from_sequence(data_dst_merged, result_path, refer_path) # 移动到trash trash_dir = os.path.join(workspace, "../trash_workspace") import shutil shutil.move(data_dst, trash_dir)
def on_initialize(self, client_dict): io.log_info('Running on %s.' % (client_dict['device_name'])) self.device_idx = client_dict['device_idx'] self.device_name = client_dict['device_name'] self.converter = client_dict['converter'] self.output_path = Path( client_dict['output_dir']) if 'output_dir' in client_dict.keys( ) else None self.alignments = client_dict['alignments'] self.debug = client_dict['debug'] from nnlib import nnlib #model process ate all GPU mem, #so we cannot use GPU for any TF operations in converter processes #therefore forcing active_DeviceConfig to CPU only nnlib.active_DeviceConfig = nnlib.DeviceConfig(cpu_only=True) return None
def extract_fanseg(input_dir, device_args={} ): multi_gpu = device_args.get('multi_gpu', False) cpu_only = device_args.get('cpu_only', False) input_path = Path(input_dir) if not input_path.exists(): raise ValueError('Input directory not found. Please ensure it exists.') paths_to_extract = [] for filename in Path_utils.get_image_paths(input_path) : filepath = Path(filename) dflimg = DFLIMG.load ( filepath ) if dflimg is not None: paths_to_extract.append (filepath) paths_to_extract_len = len(paths_to_extract) if paths_to_extract_len > 0: io.log_info ("Performing extract fanseg for %d files..." % (paths_to_extract_len) ) data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in paths_to_extract ], 'fanseg', multi_gpu=multi_gpu, cpu_only=cpu_only).run()
def main(input_path, sort_by_method): input_path = Path(input_path) sort_by_method = sort_by_method.lower() io.log_info("Running sort tool.\r\n") img_list = [] trash_img_list = [] if sort_by_method == 'blur': img_list, trash_img_list = sort_by_blur(input_path) elif sort_by_method == 'face': img_list, trash_img_list = sort_by_face(input_path) elif sort_by_method == 'face-dissim': img_list, trash_img_list = sort_by_face_dissim(input_path) elif sort_by_method == 'face-yaw': img_list, trash_img_list = sort_by_face_yaw(input_path) elif sort_by_method == 'face-pitch': img_list, trash_img_list = sort_by_face_pitch(input_path) elif sort_by_method == 'hist': img_list = sort_by_hist(input_path) elif sort_by_method == 'hist-dissim': img_list, trash_img_list = sort_by_hist_dissim(input_path) elif sort_by_method == 'brightness': img_list = sort_by_brightness(input_path) elif sort_by_method == 'hue': img_list = sort_by_hue(input_path) elif sort_by_method == 'black': img_list = sort_by_black(input_path) elif sort_by_method == 'origname': img_list, trash_img_list = sort_by_origname(input_path) elif sort_by_method == 'oneface': img_list, trash_img_list = sort_by_oneface_in_image(input_path) elif sort_by_method == 'vggface': img_list, trash_img_list = sort_by_vggface(input_path) elif sort_by_method == 'absdiff': img_list, trash_img_list = sort_by_absdiff(input_path) elif sort_by_method == 'final': img_list, trash_img_list = sort_final(input_path) elif sort_by_method == 'final-no-blur': img_list, trash_img_list = sort_final(input_path, include_by_blur=False) final_process(input_path, img_list, trash_img_list)
def load(sample_type, samples_path): samples_cache = SampleHost.samples_cache if str(samples_path) not in samples_cache.keys(): samples_cache[str(samples_path)] = [None] * SampleType.QTY samples = samples_cache[str(samples_path)] if sample_type == SampleType.IMAGE: if samples[sample_type] is None: samples[sample_type] = [ Sample(filename=filename) for filename in io.progress_bar_generator( Path_utils.get_image_paths(samples_path), "Loading") ] elif sample_type == SampleType.FACE: if samples[sample_type] is None: result = None try: result = samplelib.PackedFaceset.load(samples_path) except: io.log_err( f"Error occured while loading samplelib.PackedFaceset.load {str(samples_dat_path)}, {traceback.format_exc()}" ) if result is not None: io.log_info( f"Loaded {len(result)} packed faces from {samples_path}" ) if result is None: result = SampleHost.load_face_samples( Path_utils.get_image_paths(samples_path)) samples[sample_type] = result elif sample_type == SampleType.FACE_TEMPORAL_SORTED: if samples[sample_type] is None: samples[ sample_type] = SampleHost.upgradeToFaceTemporalSortedSamples( SampleHost.load(SampleType.FACE, samples_path)) return samples[sample_type]
def sort_by_face(input_path): io.log_info("根据相似度[similarity]排序...") img_list = [] trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) if filepath.suffix == '.png': dflimg = DFLPNG.load(str(filepath)) elif filepath.suffix == '.jpg': dflimg = DFLJPG.load(str(filepath)) else: dflimg = None if dflimg is None: io.log_err("%s 不是DeepFaceLab的图片格式,请使用DeepFaceLab提取脸图" % (filepath.name)) trash_img_list.append([str(filepath)]) continue img_list.append([str(filepath), dflimg.get_landmarks()]) img_list_len = len(img_list) for i in io.progress_bar_generator(range(0, img_list_len - 1), "Sorting"): min_score = float("inf") j_min_score = i + 1 for j in range(i + 1, len(img_list)): fl1 = img_list[i][1] fl2 = img_list[j][1] score = np.sum(np.absolute((fl2 - fl1).flatten())) if score < min_score: min_score = score j_min_score = j img_list[i + 1], img_list[j_min_score] = img_list[j_min_score], img_list[i + 1] return img_list, trash_img_list
def on_clients_finalized(self): io.progress_bar_close() if self.is_interactive: self.screen_manager.finalize() for frame in self.frames: frame.output_filename = None frame.image = None session_data = { 'frames': self.frames, 'frames_idxs': self.frames_idxs, 'frames_done_idxs': self.frames_done_idxs, 'model_iter' : self.model_iter, } self.converter_session_filepath.write_bytes( pickle.dumps(session_data) ) io.log_info ("Session is saved to " + '/'.join (self.converter_session_filepath.parts[-2:]) )
def step(workspace): import shutil for f in os.listdir(workspace): if os.path.isdir(os.path.join(workspace, f)) and f.startswith("data_dst_"): model = os.path.join(workspace, "model") model_dst = os.path.join(workspace, f, "model") if not os.path.exists(model_dst): io.log_info("Move Model Files To %s" % f) os.mkdir(model_dst) for m in os.listdir(model): mf = os.path.join(model, m) if os.path.isfile(mf): shutil.copy(os.path.join(model, m), model_dst) src = os.path.join(workspace, f) dst = os.path.join(workspace, "data_trash") io.log_info("Move %s To %s" % (src, dst)) shutil.move(src, dst) return
def add_landmarks_debug_images(input_path): io.log_info ("Adding landmarks debug images...") for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"): filepath = Path(filepath) img = cv2_imread(str(filepath)) dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue if img is not None: face_landmarks = dflimg.get_landmarks() LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True, ie_polys=IEPolys.load(dflimg.get_ie_polys()) ) output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem), '_debug.jpg') cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )
def change_workspace(): wss = [] for f in os.listdir(get_root_path()): fpath = os.path.join(get_root_path(), f) if os.path.isfile(fpath) and f.startswith("@workspace"): os.remove(fpath) elif os.path.isdir(fpath) and f.startswith("workspace"): wss.append(f) inputs = "1234567890"[0:len(wss)] for i in range(0, len(wss)): io.log_info("[ %s ] %s" % (inputs[i], wss[i])) no = io.input_str("Select Workspace:", 0)[0] idx = inputs.find(no) if idx < 0: raise Exception("Invalid Idx " + no) ws = wss[idx] io.log_info("Select " + ws) f = open(os.path.join(get_root_path(), "@" + ws), 'w') f.write(ws) f.close()
def sort_by_face_dissim(input_path): io.log_info("根据差异度[dissimilarity]排序...") img_list = [] trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) if filepath.suffix == '.png': dflimg = DFLPNG.load(str(filepath)) elif filepath.suffix == '.jpg': dflimg = DFLJPG.load(str(filepath)) else: dflimg = None if dflimg is None: io.log_err("%s 不是DeepFaceLab的图片格式,请使用DeepFaceLab提取脸图" % (filepath.name)) trash_img_list.append([str(filepath)]) continue img_list.append([str(filepath), dflimg.get_landmarks(), 0]) img_list_len = len(img_list) for i in io.progress_bar_generator(range(img_list_len - 1), "Sorting"): score_total = 0 for j in range(i + 1, len(img_list)): if i == j: continue fl1 = img_list[i][1] fl2 = img_list[j][1] score_total += np.sum(np.absolute((fl2 - fl1).flatten())) img_list[i][2] = score_total io.log_info("排序...") img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True) return img_list, trash_img_list
def sort_by_origname(input_path): io.log_info("Sort by original filename...") img_list = [] trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) dflimg = DFLIMG.load(filepath) if dflimg is None: io.log_err("%s is not a dfl image file" % (filepath.name)) trash_img_list.append([str(filepath)]) continue img_list.append([str(filepath), dflimg.get_source_filename()]) io.log_info("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1)) return img_list, trash_img_list
def skip_by_pitch(src_path, dst_path): import os import shutil import cv size = 800 r = 20 src_img_list = get_pitch_yaw_roll(src_path) dst_img_list = get_pitch_yaw_roll(dst_path) trash_path = dst_path + "_trash" if not os.path.exists(trash_path): os.makedirs(trash_path) img = cv.cv_new((size + 1, size + 1)) trans: Callable[[Any], int] = lambda v: int((v + 1) * size / 2) count = 0 for [_, pitch, yaw, _] in src_img_list: x = trans(pitch) y = trans(yaw) cv.cv_point(img, (x, y), (128, 128, 128), r) # cv.cv_show(img) xys = [] for [path, pitch, yaw, _] in dst_img_list: x = trans(pitch) y = trans(yaw) c = img[y, x] c_ = img[-y, x] if sum(c) == 255 * 3 and sum(c_) == 255 * 3: xys.append((x, y, (0, 0, 0xff))) if not os.path.exists(path) or not os.path.exists(trash_path): continue count += 1 shutil.move(path, trash_path) else: xys.append((x, y, (0xcc, 0x66, 0x33))) for (x, y, color) in xys: cv.cv_point(img, (x, y), color, 2) # cv.cv_show(img) io.log_info("Out Of Pitch, %d / %d" % (count, len(dst_img_list))) save_path = os.path.join(dst_path, "_skip_by_pitch.bmp") cv.cv_save(img, save_path)