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,
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.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'Silent start: choosed model "{self.model_name}"')
else:
io.log_info ("Choose one of saved models, or enter a name to create a new model.")
io.log_info ("[r] : rename")
io.log_info ("[d] : delete")
io.log_info ("")
for i, model_name in enumerate(saved_models_names):
s = f"[{i}] : {model_name} "
if i == 0:
s += "- latest"
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"Enter the name of the model you want to rename")
elif is_delete:
name = io.input_str(f"Enter the name of the model you want to delete")
if name in saved_models_names:
if is_rename:
new_model_name = io.input_str(f"Enter new name of the model")
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"No saved models found. Enter a name of a new model", "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"Loading {self.model_name} model...")
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 ("\nModel first run.")
if silent_start:
self.device_config = nn.DeviceConfig.BestGPU()
io.log_info (f"Silent start: choosed device {'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] - next. [space] - switch preview type. [enter] - confirm."
io.log_info (f"Choose image for the preview history. {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_static_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_static_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 ("Press enter in 2 seconds to override model settings.", 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"Autobackup every N hour", default_autobackup_hour, add_info="0..24", help_message="Autobackup model files with preview every N hour. Latest backup located in 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"Write preview history", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
if self.options['write_preview_history']:
if io.is_support_windows():
self.choose_preview_history = io.input_bool("Choose image for the preview history", False)
elif io.is_colab():
self.choose_preview_history = io.input_bool("Randomly choose new image for preview history", False, help_message="Preview image history will stay stuck with old faces if you reuse the same model on different celebs. Choose no unless you are changing src/dst to a new person")
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("Target iteration", 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("Flip faces randomly", default_random_flip, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
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("Flip SRC faces randomly", default_random_src_flip, help_message="Random horizontal flip SRC faceset. Covers more angles, but the face may look less naturally.")
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("Flip DST faces randomly", default_random_dst_flip, help_message="Random horizontal flip DST faceset. Makes generalization of src->dst better, if src random flip is not enabled.")
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("Batch_size", default_batch_size, valid_range=range, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
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):
#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_static_previews(self):
return self.onGetPreview (self.sample_for_preview)
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() )
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 ("Creating backup...", 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_static_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'=={" Model Summary ":=^{width_total}}=='] # Model/status summary
summary_text += [f'=={" "*width_total}==']
summary_text += [f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='] # Name
summary_text += [f'=={" "*width_total}==']
summary_text += [f'=={"Current iteration": >{width_name}}: {str(self.get_iter()): <{width_value}}=='] # Iter
summary_text += [f'=={" "*width_total}==']
summary_text += [f'=={" Model Options ":-^{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'=={" Running On ":-^{width_total}}=='] # Training hardware info
summary_text += [f'=={" "*width_total}==']
if len(self.device_config.devices) == 0:
summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only
else:
for device in self.device_config.devices:
summary_text += [f'=={"Device index": >{width_name}}: {device.index: <{width_value}}=='] # GPU hardware device index
summary_text += [f'=={"Name": >{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'=={"VRAM": >{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 = 'Iter: %d' % (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)