import colorsys import inspect import json import multiprocessing import operator import os import pickle import shutil import tempfile import time from pathlib import Path import cv2 import numpy as np from core import imagelib, pathex from core.cv2ex import * from core.interact import interact as io from core.leras import nn from samplelib import SampleGeneratorBase class ModelBase(object): def __init__(self, is_training=False, is_exporting=False, saved_models_path=None, training_data_src_path=None, training_data_dst_path=None, pretraining_data_path=None, pretrained_model_path=None, no_preview=False, force_model_name=None, force_gpu_idxs=None, cpu_only=False, debug=False, force_model_class_name=None, silent_start=False, **kwargs): self.is_training = is_training self.is_exporting = is_exporting self.saved_models_path = saved_models_path self.training_data_src_path = training_data_src_path self.training_data_dst_path = training_data_dst_path self.pretraining_data_path = pretraining_data_path self.pretrained_model_path = pretrained_model_path self.no_preview = no_preview self.debug = debug self.model_class_name = model_class_name = Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1] if force_model_class_name is None: if force_model_name is not None: self.model_name = force_model_name else: while True: # gather all model dat files saved_models_names = [] for filepath in pathex.get_file_paths(saved_models_path): filepath_name = filepath.name if filepath_name.endswith(f'{model_class_name}_data.dat'): saved_models_names += [ (filepath_name.split('_')[0], os.path.getmtime(filepath)) ] # sort by modified datetime saved_models_names = sorted(saved_models_names, key=operator.itemgetter(1), reverse=True ) saved_models_names = [ x[0] for x in saved_models_names ] if len(saved_models_names) != 0: if silent_start: self.model_name = saved_models_names[0] io.log_info(f'静默启动: 选择模型 "{self.model_name}"') else: io.log_info ("选择保存的一个模型, 或者输入一个名称新建模型。") io.log_info ("[r] : 重命名") io.log_info ("[d] : 删除") io.log_info ("") for i, model_name in enumerate(saved_models_names): s = f"[{i}] : {model_name} " if i == 0: s += "- 最近使用" io.log_info (s) inp = io.input_str(f"", "0", show_default_value=False ) model_idx = -1 try: model_idx = np.clip ( int(inp), 0, len(saved_models_names)-1 ) except: pass if model_idx == -1: if len(inp) == 1: is_rename = inp[0] == 'r' is_delete = inp[0] == 'd' if is_rename or is_delete: if len(saved_models_names) != 0: if is_rename: name = io.input_str(f"输入你想要重命名的模型名称") elif is_delete: name = io.input_str(f"输入你想要删除的模型名称") if name in saved_models_names: if is_rename: new_model_name = io.input_str(f"输入一个新的模型名称") for filepath in pathex.get_paths(saved_models_path): filepath_name = filepath.name model_filename, remain_filename = filepath_name.split('_', 1) if model_filename == name: if is_rename: new_filepath = filepath.parent / ( new_model_name + '_' + remain_filename ) filepath.rename (new_filepath) elif is_delete: filepath.unlink() continue self.model_name = inp else: self.model_name = saved_models_names[model_idx] else: self.model_name = io.input_str(f"未发现保存的模型. 输入一个名字新建模型", "new") self.model_name = self.model_name.replace('_', ' ') break self.model_name = self.model_name + '_' + self.model_class_name else: self.model_name = force_model_class_name self.iter = 0 self.options = {} self.options_show_override = {} self.loss_history = [] self.sample_for_preview = None self.choosed_gpu_indexes = None model_data = {} self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') ) if self.model_data_path.exists(): io.log_info (f"正在加载 {self.model_name} 模型...") model_data = pickle.loads ( self.model_data_path.read_bytes() ) self.iter = model_data.get('iter',0) if self.iter != 0: self.options = model_data['options'] self.loss_history = model_data.get('loss_history', []) self.sample_for_preview = model_data.get('sample_for_preview', None) self.choosed_gpu_indexes = model_data.get('choosed_gpu_indexes', None) if self.is_first_run(): io.log_info ("\n首次运行模型") if silent_start: self.device_config = nn.DeviceConfig.BestGPU() io.log_info (f"静默启动: 选择设备 {'CPU' if self.device_config.cpu_only else self.device_config.devices[0].name}") else: self.device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_best_multi_gpu=True)) \ if not cpu_only else nn.DeviceConfig.CPU() nn.initialize(self.device_config) #### self.default_options_path = saved_models_path / f'{self.model_class_name}_default_options.dat' self.default_options = {} if self.default_options_path.exists(): try: self.default_options = pickle.loads ( self.default_options_path.read_bytes() ) except: pass self.choose_preview_history = False self.batch_size = self.load_or_def_option('batch_size', 1) ##### io.input_skip_pending() self.on_initialize_options() if self.is_first_run(): # save as default options only for first run model initialize self.default_options_path.write_bytes( pickle.dumps (self.options) ) self.autobackup_hour = self.options.get('autobackup_hour', 0) self.write_preview_history = self.options.get('write_preview_history', False) self.target_iter = self.options.get('target_iter',0) self.random_flip = self.options.get('random_flip',True) self.random_src_flip = self.options.get('random_src_flip', False) self.random_dst_flip = self.options.get('random_dst_flip', True) self.on_initialize() self.options['batch_size'] = self.batch_size self.preview_history_writer = None if self.is_training: self.preview_history_path = self.saved_models_path / ( f'{self.get_model_name()}_history' ) self.autobackups_path = self.saved_models_path / ( f'{self.get_model_name()}_autobackups' ) if self.write_preview_history or io.is_colab(): if not self.preview_history_path.exists(): self.preview_history_path.mkdir(exist_ok=True) else: if self.iter == 0: for filename in pathex.get_image_paths(self.preview_history_path): Path(filename).unlink() if self.generator_list is None: raise ValueError( 'You didnt set_training_data_generators()') else: for i, generator in enumerate(self.generator_list): if not isinstance(generator, SampleGeneratorBase): raise ValueError('training data generator is not subclass of SampleGeneratorBase') self.update_sample_for_preview(choose_preview_history=self.choose_preview_history) if self.autobackup_hour != 0: self.autobackup_start_time = time.time() if not self.autobackups_path.exists(): self.autobackups_path.mkdir(exist_ok=True) io.log_info( self.get_summary_text() ) def update_sample_for_preview(self, choose_preview_history=False, force_new=False): if self.sample_for_preview is None or choose_preview_history or force_new: if choose_preview_history and io.is_support_windows(): wnd_name = "[p] - 下一张. [space] - 切换预览类型. [enter] - 确定." io.log_info (f"为预览历史选择图像. {wnd_name}") io.named_window(wnd_name) io.capture_keys(wnd_name) choosed = False preview_id_counter = 0 while not choosed: self.sample_for_preview = self.generate_next_samples() previews = self.get_history_previews() io.show_image( wnd_name, ( previews[preview_id_counter % len(previews) ][1] *255).astype(np.uint8) ) while True: key_events = io.get_key_events(wnd_name) key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False) if key == ord('\n') or key == ord('\r'): choosed = True break elif key == ord(' '): preview_id_counter += 1 break elif key == ord('p'): break try: io.process_messages(0.1) except KeyboardInterrupt: choosed = True io.destroy_window(wnd_name) else: self.sample_for_preview = self.generate_next_samples() try: self.get_history_previews() except: self.sample_for_preview = self.generate_next_samples() self.last_sample = self.sample_for_preview def load_or_def_option(self, name, def_value): options_val = self.options.get(name, None) if options_val is not None: return options_val def_opt_val = self.default_options.get(name, None) if def_opt_val is not None: return def_opt_val return def_value def ask_override(self): return self.is_training and self.iter != 0 and io.input_in_time ("两秒内按下 enter 修改模型设置...", 5 if io.is_colab() else 2 ) def ask_autobackup_hour(self, default_value=0): default_autobackup_hour = self.options['autobackup_hour'] = self.load_or_def_option('autobackup_hour', default_value) self.options['autobackup_hour'] = io.input_int(f"自动备份时间? ", default_autobackup_hour, add_info="0..24", help_message="每 N 小时自动备份一次模型文件和预览. 最新备份位于 model/<>_autobackups/01") def ask_write_preview_history(self, default_value=False): default_write_preview_history = self.load_or_def_option('write_preview_history', default_value) self.options['write_preview_history'] = io.input_bool(f"保存预览历史记录", default_write_preview_history, help_message="预览历史将写入 history 文件夹.") if self.options['write_preview_history']: if io.is_support_windows(): self.choose_preview_history = io.input_bool("为预览历史选择图像", False) elif io.is_colab(): self.choose_preview_history = io.input_bool("随机选择新图像为预览历史", False, help_message="如果你在不同的人物上重复使用相同的模型, 预览历史图像将停留在旧面部上. 除非你将 src/dst 更改为新人物, 不然请选择 no") def ask_target_iter(self, default_value=0): default_target_iter = self.load_or_def_option('target_iter', default_value) self.options['target_iter'] = max(0, io.input_int("目标迭代次数", default_target_iter)) def ask_random_flip(self): default_random_flip = self.load_or_def_option('random_flip', True) self.options['random_flip'] = io.input_bool("随机翻转面部", default_random_flip, help_message="如果不启用此选项, 预测的面部看起来会更自然, 但是 src 面部数据集应该像 dst 面部数据集一样覆盖所有的面部方向.") def ask_random_src_flip(self): default_random_src_flip = self.load_or_def_option('random_src_flip', False) self.options['random_src_flip'] = io.input_bool("随机翻转 SRC 面部", default_random_src_flip, help_message="随机水平翻转 SRC 面部数据集. 覆盖更多角度, 但面部可能看起来不那么自然.") def ask_random_dst_flip(self): default_random_dst_flip = self.load_or_def_option('random_dst_flip', True) self.options['random_dst_flip'] = io.input_bool("随机翻转 DST 面部", default_random_dst_flip, help_message="随机水平翻转 DST 面部数据集. 如果未启用 src 随机翻转, 则使 src->dst 的泛化更好.") def ask_batch_size(self, suggest_batch_size=None, range=None): default_batch_size = self.load_or_def_option('batch_size', suggest_batch_size or self.batch_size) batch_size = max(0, io.input_int("批量大小", default_batch_size, valid_range=range, help_message="较大的批量大小更适合神经网络的泛化, 但会导致 OOM 错误. 手动为你的显卡调整此值.")) if range is not None: batch_size = np.clip(batch_size, range[0], range[1]) self.options['batch_size'] = self.batch_size = batch_size #overridable def on_initialize_options(self): pass #overridable def on_initialize(self): ''' initialize your models store and retrieve your model options in self.options[''] check example ''' pass #overridable def onSave(self): #save your models here pass #overridable def onTrainOneIter(self, sample, generator_list): #train your models here #return array of losses return ( ('loss_src', 0), ('loss_dst', 0) ) #overridable def onGetPreview(self, sample, for_history=False): #you can return multiple previews #return [ ('preview_name',preview_rgb), ... ] return [] #overridable if you want model name differs from folder name def get_model_name(self): return self.model_name #overridable , return [ [model, filename],... ] list def get_model_filename_list(self): return [] #overridable def get_MergerConfig(self): #return predictor_func, predictor_input_shape, MergerConfig() for the model raise NotImplementedError def get_pretraining_data_path(self): return self.pretraining_data_path def get_target_iter(self): return self.target_iter def is_reached_iter_goal(self): return self.target_iter != 0 and self.iter >= self.target_iter def get_previews(self): return self.onGetPreview ( self.last_sample ) def get_history_previews(self): return self.onGetPreview (self.sample_for_preview, for_history=True) def get_preview_history_writer(self): if self.preview_history_writer is None: self.preview_history_writer = PreviewHistoryWriter() return self.preview_history_writer def save(self): Path( self.get_summary_path() ).write_text( self.get_summary_text(),encoding='UTF-8' ) self.onSave() model_data = { 'iter': self.iter, 'options': self.options, 'loss_history': self.loss_history, 'sample_for_preview' : self.sample_for_preview, 'choosed_gpu_indexes' : self.choosed_gpu_indexes, } pathex.write_bytes_safe (self.model_data_path, pickle.dumps(model_data) ) if self.autobackup_hour != 0: diff_hour = int ( (time.time() - self.autobackup_start_time) // 3600 ) if diff_hour > 0 and diff_hour % self.autobackup_hour == 0: self.autobackup_start_time += self.autobackup_hour*3600 self.create_backup() def create_backup(self): io.log_info ("正在创建备份...", end='\r') if not self.autobackups_path.exists(): self.autobackups_path.mkdir(exist_ok=True) bckp_filename_list = [ self.get_strpath_storage_for_file(filename) for _, filename in self.get_model_filename_list() ] bckp_filename_list += [ str(self.get_summary_path()), str(self.model_data_path) ] for i in range(24,0,-1): idx_str = '%.2d' % i next_idx_str = '%.2d' % (i+1) idx_backup_path = self.autobackups_path / idx_str next_idx_packup_path = self.autobackups_path / next_idx_str if idx_backup_path.exists(): if i == 24: pathex.delete_all_files(idx_backup_path) else: next_idx_packup_path.mkdir(exist_ok=True) pathex.move_all_files (idx_backup_path, next_idx_packup_path) if i == 1: idx_backup_path.mkdir(exist_ok=True) for filename in bckp_filename_list: shutil.copy ( str(filename), str(idx_backup_path / Path(filename).name) ) previews = self.get_previews() plist = [] for i in range(len(previews)): name, bgr = previews[i] plist += [ (bgr, idx_backup_path / ( ('preview_%s.jpg') % (name)) ) ] if len(plist) != 0: self.get_preview_history_writer().post(plist, self.loss_history, self.iter) def debug_one_iter(self): images = [] for generator in self.generator_list: for i,batch in enumerate(next(generator)): if len(batch.shape) == 4: images.append( batch[0] ) return imagelib.equalize_and_stack_square (images) def generate_next_samples(self): sample = [] for generator in self.generator_list: if generator.is_initialized(): sample.append ( generator.generate_next() ) else: sample.append ( [] ) self.last_sample = sample return sample #overridable def should_save_preview_history(self): return (not io.is_colab() and self.iter % 10 == 0) or (io.is_colab() and self.iter % 100 == 0) def train_one_iter(self): iter_time = time.time() losses = self.onTrainOneIter() iter_time = time.time() - iter_time self.loss_history.append ( [float(loss[1]) for loss in losses] ) if self.should_save_preview_history(): plist = [] if io.is_colab(): previews = self.get_previews() for i in range(len(previews)): name, bgr = previews[i] plist += [ (bgr, self.get_strpath_storage_for_file('preview_%s.jpg' % (name) ) ) ] if self.write_preview_history: previews = self.get_history_previews() for i in range(len(previews)): name, bgr = previews[i] path = self.preview_history_path / name plist += [ ( bgr, str ( path / ( f'{self.iter:07d}.jpg') ) ) ] if not io.is_colab(): plist += [ ( bgr, str ( path / ( '_last.jpg' ) )) ] if len(plist) != 0: self.get_preview_history_writer().post(plist, self.loss_history, self.iter) self.iter += 1 return self.iter, iter_time def pass_one_iter(self): self.generate_next_samples() def finalize(self): nn.close_session() def is_first_run(self): return self.iter == 0 def is_debug(self): return self.debug def set_batch_size(self, batch_size): self.batch_size = batch_size def get_batch_size(self): return self.batch_size def get_iter(self): return self.iter def set_iter(self, iter): self.iter = iter self.loss_history = self.loss_history[:iter] def get_loss_history(self): return self.loss_history def set_training_data_generators (self, generator_list): self.generator_list = generator_list def get_training_data_generators (self): return self.generator_list def get_model_root_path(self): return self.saved_models_path def get_strpath_storage_for_file(self, filename): return str( self.saved_models_path / ( self.get_model_name() + '_' + filename) ) def get_summary_path(self): return self.get_strpath_storage_for_file('summary.txt') def get_summary_text(self): visible_options = self.options.copy() visible_options.update(self.options_show_override) ###Generate text summary of model hyperparameters #Find the longest key name and value string. Used as column widths. width_name = max([len(k) for k in visible_options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration" width_value = max([len(str(x)) for x in visible_options.values()] + [len(str(self.get_iter())), len(self.get_model_name())]) + 1 # Single space buffer to right edge if len(self.device_config.devices) != 0: #Check length of GPU names width_value = max([len(device.name)+1 for device in self.device_config.devices] + [width_value]) width_total = width_name + width_value + 2 #Plus 2 for ": " summary_text = [] summary_text += [f'\n=={" 模型参数 ":=^{width_total}}=='] # Model/status summary summary_text += [f'{" "*width_total}'] summary_text += [f'{"模型名称": >{width_name}}: {self.get_model_name(): <{width_value}}'] # Name summary_text += [f'{" "*width_total}'] summary_text += [f'{"迭代数量": >{width_name}}: {str(self.get_iter()): <{width_value}}'] # Iter summary_text += [f'{" "*width_total}'] summary_text += [f'{" 模型选项 ":-^{width_total}}'] # Model options summary_text += [f'{" "*width_total}'] for key in visible_options.keys(): summary_text += [f'{key: >{width_name}}: {str(visible_options[key]): <{width_value}}'] # visible_options key/value pairs summary_text += [f'{" "*width_total}'] summary_text += [f'{" 运行设备 ":-^{width_total}}'] # Training hardware info summary_text += [f'{" "*width_total}'] if len(self.device_config.devices) == 0: summary_text += [f'{"当前设备": >{width_name}}: {"CPU": <{width_value}}'] # cpu_only else: for device in self.device_config.devices: summary_text += [f'{"设备编号": >{width_name}}: {device.index: <{width_value}}'] # GPU hardware device index summary_text += [f'{"设备名称": >{width_name}}: {device.name: <{width_value}}'] # GPU name vram_str = f'{device.total_mem_gb:.2f}GB' # GPU VRAM - Formated as #.## (or ##.##) summary_text += [f'{"可用显存": >{width_name}}: {vram_str: <{width_value}}'] summary_text += [f'{" "*width_total}'] summary_text += [f'{"="*width_total}'] summary_text = "\n".join (summary_text) return summary_text @staticmethod def get_loss_history_preview(loss_history, iter, w, c): loss_history = np.array (loss_history.copy()) lh_height = 100 lh_img = np.ones ( (lh_height,w,c) ) * 0.1 if len(loss_history) != 0: loss_count = len(loss_history[0]) lh_len = len(loss_history) l_per_col = lh_len / w plist_max = [ [ max (0.0, loss_history[int(col*l_per_col)][p], *[ loss_history[i_ab][p] for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) ) ] ) for p in range(loss_count) ] for col in range(w) ] plist_min = [ [ min (plist_max[col][p], loss_history[int(col*l_per_col)][p], *[ loss_history[i_ab][p] for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) ) ] ) for p in range(loss_count) ] for col in range(w) ] plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2 for col in range(0, w): for p in range(0,loss_count): point_color = [1.0]*c point_color[0:3] = colorsys.hsv_to_rgb ( p * (1.0/loss_count), 1.0, 1.0 ) ph_max = int ( (plist_max[col][p] / plist_abs_max) * (lh_height-1) ) ph_max = np.clip( ph_max, 0, lh_height-1 ) ph_min = int ( (plist_min[col][p] / plist_abs_max) * (lh_height-1) ) ph_min = np.clip( ph_min, 0, lh_height-1 ) for ph in range(ph_min, ph_max+1): lh_img[ (lh_height-ph-1), col ] = point_color lh_lines = 5 lh_line_height = (lh_height-1)/lh_lines for i in range(0,lh_lines+1): lh_img[ int(i*lh_line_height), : ] = (0.8,)*c last_line_t = int((lh_lines-1)*lh_line_height) last_line_b = int(lh_lines*lh_line_height) lh_text = '迭代: %d 更多好玩的AI应用，访问https://deepface.cc' % (iter) if iter != 0 else '' lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image ( (last_line_b-last_line_t,w,c), lh_text, color=[0.8]*c ) return lh_img class PreviewHistoryWriter(): def __init__(self): self.sq = multiprocessing.Queue() self.p = multiprocessing.Process(target=self.process, args=( self.sq, )) self.p.daemon = True self.p.start() def process(self, sq): while True: while not sq.empty(): plist, loss_history, iter = sq.get() preview_lh_cache = {} for preview, filepath in plist: filepath = Path(filepath) i = (preview.shape[1], preview.shape[2]) preview_lh = preview_lh_cache.get(i, None) if preview_lh is None: preview_lh = ModelBase.get_loss_history_preview(loss_history, iter, preview.shape[1], preview.shape[2]) preview_lh_cache[i] = preview_lh img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8) filepath.parent.mkdir(parents=True, exist_ok=True) cv2_imwrite (filepath, img ) time.sleep(0.01) def post(self, plist, loss_history, iter): self.sq.put ( (plist, loss_history, iter) ) # disable pickling def __getstate__(self): return dict() def __setstate__(self, d): self.__dict__.update(d)