mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-21 22:13:20 -07:00
support preview or not
This commit is contained in:
plucky
parent
0d4ee90c36
commit
0b562a151e
3 changed files with 147 additions and 146 deletions
|
|
@ -16,7 +16,7 @@ $ docker build -t deepfacelab-cpu -f Dockerfile.cpu .
|
|||
```
|
||||
$ docker run -p 8888:8888 --hostname deepfacelab-cpu --name deepfacelab-cpu -v **your source path**:/srv deepfacelab-cpu
|
||||
# for example
|
||||
$ docker run -p 8888:8888 --hostname deepfacelab-cpu --name deepfacelab-cpu -v /Users/plucky/own/DeepFaceLab:/srv deepfacelab-cpu
|
||||
$ docker run -p 8888:8888 --hostname deepfacelab-cpu --name deepfacelab-cpu -v $PWD:/srv deepfacelab-cpu
|
||||
```
|
||||
|
||||
then you will see the log:
|
||||
|
|
|
|||
123
main.py
123
main.py
|
|
@ -20,19 +20,19 @@ def str2bool(v):
|
|||
return False
|
||||
else:
|
||||
raise argparse.ArgumentTypeError('Boolean value expected.')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os_utils.set_process_lowest_prio()
|
||||
parser = argparse.ArgumentParser()
|
||||
os_utils.set_process_lowest_prio()
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--tf-suppress-std', action="store_true", dest="tf_suppress_std", default=False, help="Suppress tensorflow initialization info. May not works on some python builds such as anaconda python 3.6.4. If you can fix it, you are welcome.")
|
||||
|
||||
|
||||
subparsers = parser.add_subparsers()
|
||||
|
||||
|
||||
def process_extract(arguments):
|
||||
from mainscripts import Extractor
|
||||
from mainscripts import Extractor
|
||||
Extractor.main (
|
||||
input_dir=arguments.input_dir,
|
||||
output_dir=arguments.output_dir,
|
||||
input_dir=arguments.input_dir,
|
||||
output_dir=arguments.output_dir,
|
||||
debug=arguments.debug,
|
||||
face_type=arguments.face_type,
|
||||
detector=arguments.detector,
|
||||
|
|
@ -41,43 +41,43 @@ if __name__ == "__main__":
|
|||
manual_fix=arguments.manual_fix,
|
||||
manual_window_size=arguments.manual_window_size
|
||||
)
|
||||
|
||||
|
||||
extract_parser = subparsers.add_parser( "extract", help="Extract the faces from a pictures.")
|
||||
extract_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
|
||||
extract_parser.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.")
|
||||
extract_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Writes debug images to [output_dir]_debug\ directory.")
|
||||
extract_parser.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")
|
||||
extract_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Writes debug images to [output_dir]_debug\ directory.")
|
||||
extract_parser.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")
|
||||
extract_parser.add_argument('--detector', dest="detector", choices=['dlib','mt','manual'], default='dlib', help="Type of detector. Default 'dlib'. 'mt' (MTCNNv1) - faster, better, almost no jitter, perfect for gathering thousands faces for src-set. It is also good for dst-set, but can generate false faces in frames where main face not recognized! In this case for dst-set use either 'dlib' with '--manual-fix' or '--detector manual'. Manual detector suitable only for dst-set.")
|
||||
extract_parser.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
|
||||
extract_parser.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.")
|
||||
extract_parser.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=0, help="Manual fix window size. Example: 1368. Default: frame size.")
|
||||
extract_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.")
|
||||
|
||||
|
||||
extract_parser.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=0, help="Manual fix window size. Example: 1368. Default: frame size.")
|
||||
extract_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.")
|
||||
|
||||
|
||||
extract_parser.set_defaults (func=process_extract)
|
||||
|
||||
def process_sort(arguments):
|
||||
|
||||
def process_sort(arguments):
|
||||
from mainscripts import Sorter
|
||||
Sorter.main (input_path=arguments.input_dir, sort_by_method=arguments.sort_by_method)
|
||||
|
||||
sort_parser = subparsers.add_parser( "sort", help="Sort faces in a directory.")
|
||||
|
||||
sort_parser = subparsers.add_parser( "sort", help="Sort faces in a directory.")
|
||||
sort_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
|
||||
sort_parser.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "hist", "hist-dissim", "brightness", "hue", "black", "origname"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
|
||||
sort_parser.set_defaults (func=process_sort)
|
||||
|
||||
def process_train(arguments):
|
||||
|
||||
|
||||
def process_train(arguments):
|
||||
|
||||
if 'DFL_TARGET_EPOCH' in os.environ.keys():
|
||||
arguments.target_epoch = int ( os.environ['DFL_TARGET_EPOCH'] )
|
||||
|
||||
|
||||
if 'DFL_BATCH_SIZE' in os.environ.keys():
|
||||
arguments.batch_size = int ( os.environ['DFL_BATCH_SIZE'] )
|
||||
|
||||
from mainscripts import Trainer
|
||||
Trainer.main (
|
||||
training_data_src_dir=arguments.training_data_src_dir,
|
||||
training_data_dst_dir=arguments.training_data_dst_dir,
|
||||
model_path=arguments.model_dir,
|
||||
training_data_src_dir=arguments.training_data_src_dir,
|
||||
training_data_dst_dir=arguments.training_data_dst_dir,
|
||||
model_path=arguments.model_dir,
|
||||
model_name=arguments.model_name,
|
||||
debug = arguments.debug,
|
||||
#**options
|
||||
|
|
@ -91,32 +91,33 @@ if __name__ == "__main__":
|
|||
force_gpu_idxs = arguments.force_gpu_idxs,
|
||||
cpu_only = arguments.cpu_only
|
||||
)
|
||||
|
||||
train_parser = subparsers.add_parser( "train", help="Trainer")
|
||||
|
||||
train_parser = subparsers.add_parser( "train", help="Trainer")
|
||||
train_parser.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of src-set.")
|
||||
train_parser.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of dst-set.")
|
||||
train_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
|
||||
train_parser.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model")
|
||||
train_parser.add_argument('--write-preview-history', action="store_true", dest="write_preview_history", default=False, help="Enable write preview history.")
|
||||
train_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug training.")
|
||||
train_parser.add_argument('--batch-size', type=int, dest="batch_size", default=0, help="Model batch size. Default - auto. Environment variable: ODFS_BATCH_SIZE.")
|
||||
train_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug training.")
|
||||
train_parser.add_argument('--preview', action="store_true", dest="preview", default=True, help="Show preview.")
|
||||
train_parser.add_argument('--batch-size', type=int, dest="batch_size", default=0, help="Model batch size. Default - auto. Environment variable: ODFS_BATCH_SIZE.")
|
||||
train_parser.add_argument('--target-epoch', type=int, dest="target_epoch", default=0, help="Train until target epoch. Default - unlimited. Environment variable: ODFS_TARGET_EPOCH.")
|
||||
train_parser.add_argument('--save-interval-min', type=int, dest="save_interval_min", default=10, help="Save interval in minutes. Default 10.")
|
||||
train_parser.add_argument('--choose-worst-gpu', action="store_true", dest="choose_worst_gpu", default=False, help="Choose worst GPU instead of best.")
|
||||
train_parser.add_argument('--force-best-gpu-idx', type=int, dest="force_best_gpu_idx", default=-1, help="Force to choose this GPU idx as best(worst).")
|
||||
train_parser.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="MultiGPU option. It will select only same best(worst) GPU models.")
|
||||
train_parser.add_argument('--force-gpu-idxs', type=str, dest="force_gpu_idxs", default=None, help="Override final GPU idxs. Example: 0,1,2.")
|
||||
train_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
|
||||
|
||||
train_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
|
||||
|
||||
train_parser.set_defaults (func=process_train)
|
||||
|
||||
|
||||
def process_convert(arguments):
|
||||
if arguments.ask_for_params:
|
||||
try:
|
||||
mode = int ( input ("Choose mode: (1) hist match, (2) hist match bw, (3) seamless (default), (4) seamless hist match : ") )
|
||||
except:
|
||||
mode = 3
|
||||
|
||||
|
||||
if mode == 1:
|
||||
arguments.mode = 'hist-match'
|
||||
elif mode == 2:
|
||||
|
|
@ -125,73 +126,73 @@ if __name__ == "__main__":
|
|||
arguments.mode = 'seamless'
|
||||
elif mode == 4:
|
||||
arguments.mode = 'seamless-hist-match'
|
||||
|
||||
|
||||
if arguments.mode == 'hist-match' or arguments.mode == 'hist-match-bw':
|
||||
try:
|
||||
arguments.masked_hist_match = bool ( {"1":True,"0":False}[input("Masked hist match? [0 or 1] (default 1) : ").lower()] )
|
||||
except:
|
||||
arguments.masked_hist_match = True
|
||||
|
||||
|
||||
if arguments.mode == 'hist-match' or arguments.mode == 'hist-match-bw' or arguments.mode == 'seamless-hist-match':
|
||||
try:
|
||||
hist_match_threshold = int ( input ("Hist match threshold. [0..255] (default - 255) : ") )
|
||||
arguments.hist_match_threshold = hist_match_threshold
|
||||
except:
|
||||
arguments.hist_match_threshold = 255
|
||||
|
||||
arguments.hist_match_threshold = 255
|
||||
|
||||
try:
|
||||
arguments.use_predicted_mask = bool ( {"1":True,"0":False}[input("Use predicted mask? [0 or 1] (default 1) : ").lower()] )
|
||||
except:
|
||||
arguments.use_predicted_mask = False
|
||||
|
||||
arguments.use_predicted_mask = False
|
||||
|
||||
try:
|
||||
arguments.erode_mask_modifier = int ( input ("Choose erode mask modifier [-200..200] (default 0) : ") )
|
||||
except:
|
||||
arguments.erode_mask_modifier = 0
|
||||
|
||||
|
||||
try:
|
||||
arguments.blur_mask_modifier = int ( input ("Choose blur mask modifier [-200..200] (default 0) : ") )
|
||||
except:
|
||||
arguments.blur_mask_modifier = 0
|
||||
|
||||
|
||||
if arguments.mode == 'seamless' or arguments.mode == 'seamless-hist-match':
|
||||
try:
|
||||
arguments.seamless_erode_mask_modifier = int ( input ("Choose seamless erode mask modifier [-100..100] (default 0) : ") )
|
||||
except:
|
||||
arguments.seamless_erode_mask_modifier = 0
|
||||
|
||||
|
||||
try:
|
||||
arguments.output_face_scale_modifier = int ( input ("Choose output face scale modifier [-50..50] (default 0) : ") )
|
||||
except:
|
||||
arguments.output_face_scale_modifier = 0
|
||||
|
||||
|
||||
try:
|
||||
arguments.transfercolor = bool ( {"1":True,"0":False}[input("Transfer color from dst face to converted final face? [0 or 1] (default 0) : ").lower()] )
|
||||
except:
|
||||
arguments.transfercolor = False
|
||||
|
||||
arguments.transfercolor = False
|
||||
|
||||
try:
|
||||
arguments.final_image_color_degrade_power = int ( input ("Degrade color power of final image [0..100] (default 0) : ") )
|
||||
except:
|
||||
arguments.final_image_color_degrade_power = 0
|
||||
|
||||
|
||||
try:
|
||||
arguments.alpha = bool ( {"1":True,"0":False}[input("Export png with alpha channel? [0..1] (default 0) : ").lower()] )
|
||||
except:
|
||||
arguments.alpha = False
|
||||
|
||||
|
||||
arguments.erode_mask_modifier = np.clip ( int(arguments.erode_mask_modifier), -200, 200)
|
||||
arguments.blur_mask_modifier = np.clip ( int(arguments.blur_mask_modifier), -200, 200)
|
||||
arguments.seamless_erode_mask_modifier = np.clip ( int(arguments.seamless_erode_mask_modifier), -100, 100)
|
||||
arguments.output_face_scale_modifier = np.clip ( int(arguments.output_face_scale_modifier), -50, 50)
|
||||
|
||||
|
||||
from mainscripts import Converter
|
||||
Converter.main (
|
||||
input_dir=arguments.input_dir,
|
||||
output_dir=arguments.output_dir,
|
||||
input_dir=arguments.input_dir,
|
||||
output_dir=arguments.output_dir,
|
||||
aligned_dir=arguments.aligned_dir,
|
||||
model_dir=arguments.model_dir,
|
||||
model_name=arguments.model_name,
|
||||
model_dir=arguments.model_dir,
|
||||
model_name=arguments.model_name,
|
||||
debug = arguments.debug,
|
||||
mode = arguments.mode,
|
||||
masked_hist_match = arguments.masked_hist_match,
|
||||
|
|
@ -207,36 +208,36 @@ if __name__ == "__main__":
|
|||
force_best_gpu_idx = arguments.force_best_gpu_idx,
|
||||
cpu_only = arguments.cpu_only
|
||||
)
|
||||
|
||||
convert_parser = subparsers.add_parser( "convert", help="Converter")
|
||||
|
||||
convert_parser = subparsers.add_parser( "convert", help="Converter")
|
||||
convert_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
|
||||
convert_parser.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.")
|
||||
convert_parser.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored. Not used in AVATAR model.")
|
||||
convert_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
|
||||
convert_parser.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model")
|
||||
convert_parser.add_argument('--ask-for-params', action="store_true", dest="ask_for_params", default=False, help="Ask for params.")
|
||||
convert_parser.add_argument('--ask-for-params', action="store_true", dest="ask_for_params", default=False, help="Ask for params.")
|
||||
convert_parser.add_argument('--mode', dest="mode", choices=['seamless','hist-match', 'hist-match-bw','seamless-hist-match'], default='seamless', help="Face overlaying mode. Seriously affects result.")
|
||||
convert_parser.add_argument('--masked-hist-match', type=str2bool, nargs='?', const=True, default=True, help="True or False. Excludes background for hist match. Default - True.")
|
||||
convert_parser.add_argument('--hist-match-threshold', type=int, dest="hist_match_threshold", default=255, help="Hist match threshold. Decrease to hide artifacts of hist match. Valid range [0..255]. Default 255")
|
||||
convert_parser.add_argument('--use-predicted-mask', action="store_true", dest="use_predicted_mask", default=True, help="Use predicted mask by model. Default - True.")
|
||||
convert_parser.add_argument('--erode-mask-modifier', type=int, dest="erode_mask_modifier", default=0, help="Automatic erode mask modifier. Valid range [-200..200].")
|
||||
convert_parser.add_argument('--blur-mask-modifier', type=int, dest="blur_mask_modifier", default=0, help="Automatic blur mask modifier. Valid range [-200..200].")
|
||||
convert_parser.add_argument('--blur-mask-modifier', type=int, dest="blur_mask_modifier", default=0, help="Automatic blur mask modifier. Valid range [-200..200].")
|
||||
convert_parser.add_argument('--seamless-erode-mask-modifier', type=int, dest="seamless_erode_mask_modifier", default=0, help="Automatic seamless erode mask modifier. Valid range [-200..200].")
|
||||
convert_parser.add_argument('--output-face-scale-modifier', type=int, dest="output_face_scale_modifier", default=0, help="Output face scale modifier. Valid range [-50..50].")
|
||||
convert_parser.add_argument('--final-image-color-degrade-power', type=int, dest="final_image_color_degrade_power", default=0, help="Degrades colors of final image to hide face problems. Valid range [0..100].")
|
||||
convert_parser.add_argument('--output-face-scale-modifier', type=int, dest="output_face_scale_modifier", default=0, help="Output face scale modifier. Valid range [-50..50].")
|
||||
convert_parser.add_argument('--final-image-color-degrade-power', type=int, dest="final_image_color_degrade_power", default=0, help="Degrades colors of final image to hide face problems. Valid range [0..100].")
|
||||
convert_parser.add_argument('--transfercolor', action="store_true", dest="transfercolor", default=False, help="Transfer color from dst face to converted final face.")
|
||||
convert_parser.add_argument('--alpha', action="store_true", dest="alpha", default=False, help="Embeds alpha channel of face mask to final PNG. Used in manual composing video by editors such as Sony Vegas or After Effects.")
|
||||
convert_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.")
|
||||
convert_parser.add_argument('--force-best-gpu-idx', type=int, dest="force_best_gpu_idx", default=-1, help="Force to choose this GPU idx as best.")
|
||||
convert_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.")
|
||||
|
||||
|
||||
convert_parser.set_defaults(func=process_convert)
|
||||
|
||||
def bad_args(arguments):
|
||||
parser.print_help()
|
||||
exit(0)
|
||||
parser.set_defaults(func=bad_args)
|
||||
|
||||
|
||||
arguments = parser.parse_args()
|
||||
if arguments.tf_suppress_std:
|
||||
os.environ['TF_SUPPRESS_STD'] = '1'
|
||||
|
|
|
|||
|
|
@ -1,76 +1,76 @@
|
|||
import sys
|
||||
import traceback
|
||||
import queue
|
||||
import colorsys
|
||||
import colorsys
|
||||
import time
|
||||
import numpy as np
|
||||
import itertools
|
||||
|
||||
from pathlib import Path
|
||||
from utils import Path_utils
|
||||
from utils import image_utils
|
||||
from utils import image_utils
|
||||
import cv2
|
||||
|
||||
def trainerThread (input_queue, output_queue, training_data_src_dir, training_data_dst_dir, model_path, model_name, save_interval_min=10, debug=False, target_epoch=0, **in_options):
|
||||
|
||||
while True:
|
||||
try:
|
||||
try:
|
||||
training_data_src_path = Path(training_data_src_dir)
|
||||
training_data_dst_path = Path(training_data_dst_dir)
|
||||
model_path = Path(model_path)
|
||||
|
||||
|
||||
if not training_data_src_path.exists():
|
||||
print( 'Training data src directory is not exists.')
|
||||
return
|
||||
|
||||
|
||||
if not training_data_dst_path.exists():
|
||||
print( 'Training data dst directory is not exists.')
|
||||
return
|
||||
|
||||
|
||||
if not model_path.exists():
|
||||
model_path.mkdir(exist_ok=True)
|
||||
|
||||
|
||||
|
||||
import models
|
||||
|
||||
|
||||
import models
|
||||
model = models.import_model(model_name)(
|
||||
model_path,
|
||||
training_data_src_path=training_data_src_path,
|
||||
training_data_dst_path=training_data_dst_path,
|
||||
model_path,
|
||||
training_data_src_path=training_data_src_path,
|
||||
training_data_dst_path=training_data_dst_path,
|
||||
debug=debug,
|
||||
**in_options)
|
||||
|
||||
|
||||
is_reached_goal = (target_epoch > 0 and model.get_epoch() >= target_epoch)
|
||||
|
||||
|
||||
def model_save():
|
||||
if not debug and not is_reached_goal:
|
||||
model.save()
|
||||
|
||||
|
||||
def send_preview():
|
||||
if not debug:
|
||||
previews = model.get_previews()
|
||||
if not debug:
|
||||
previews = model.get_previews()
|
||||
output_queue.put ( {'op':'show', 'previews': previews, 'epoch':model.get_epoch(), 'loss_history': model.get_loss_history().copy() } )
|
||||
else:
|
||||
previews = [( 'debug, press update for new', model.debug_one_epoch())]
|
||||
output_queue.put ( {'op':'show', 'previews': previews} )
|
||||
|
||||
|
||||
|
||||
|
||||
if model.is_first_run():
|
||||
model_save()
|
||||
|
||||
|
||||
if target_epoch != 0:
|
||||
if is_reached_goal:
|
||||
print ('Model already trained to target epoch. You can use preview.')
|
||||
else:
|
||||
print('Starting. Target epoch: %d. Press "Enter" to stop training and save model.' % (target_epoch) )
|
||||
else:
|
||||
else:
|
||||
print('Starting. Press "Enter" to stop training and save model.')
|
||||
|
||||
|
||||
last_save_time = time.time()
|
||||
for i in itertools.count(0,1):
|
||||
if not debug:
|
||||
if not is_reached_goal:
|
||||
loss_string = model.train_one_epoch()
|
||||
loss_string = model.train_one_epoch()
|
||||
|
||||
print (loss_string, end='\r')
|
||||
if target_epoch != 0 and model.get_epoch() >= target_epoch:
|
||||
|
|
@ -80,39 +80,39 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
|
|||
print ('You can use preview now.')
|
||||
|
||||
if not is_reached_goal and (time.time() - last_save_time) >= save_interval_min*60:
|
||||
last_save_time = time.time()
|
||||
last_save_time = time.time()
|
||||
model_save()
|
||||
send_preview()
|
||||
|
||||
|
||||
if i==0:
|
||||
if is_reached_goal:
|
||||
model.pass_one_epoch()
|
||||
model.pass_one_epoch()
|
||||
send_preview()
|
||||
|
||||
|
||||
if debug:
|
||||
time.sleep(0.005)
|
||||
|
||||
|
||||
while not input_queue.empty():
|
||||
input = input_queue.get()
|
||||
op = input['op']
|
||||
if op == 'save':
|
||||
model_save()
|
||||
elif op == 'preview':
|
||||
elif op == 'preview':
|
||||
if is_reached_goal:
|
||||
model.pass_one_epoch()
|
||||
model.pass_one_epoch()
|
||||
send_preview()
|
||||
elif op == 'close':
|
||||
model_save()
|
||||
i = -1
|
||||
break
|
||||
|
||||
|
||||
if i == -1:
|
||||
break
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
model.finalize()
|
||||
|
||||
|
||||
except Exception as e:
|
||||
print ('Error: %s' % (str(e)))
|
||||
traceback.print_exc()
|
||||
|
|
@ -120,8 +120,8 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
|
|||
output_queue.put ( {'op':'close'} )
|
||||
|
||||
def previewThread (input_queue, output_queue):
|
||||
|
||||
|
||||
|
||||
|
||||
previews = None
|
||||
loss_history = None
|
||||
selected_preview = 0
|
||||
|
|
@ -129,7 +129,7 @@ def previewThread (input_queue, output_queue):
|
|||
is_showing = False
|
||||
is_waiting_preview = False
|
||||
epoch = 0
|
||||
while True:
|
||||
while True:
|
||||
if not input_queue.empty():
|
||||
input = input_queue.get()
|
||||
op = input['op']
|
||||
|
|
@ -145,7 +145,7 @@ def previewThread (input_queue, output_queue):
|
|||
(h, w, c) = preview_rgb.shape
|
||||
max_h = max (max_h, h)
|
||||
max_w = max (max_w, w)
|
||||
|
||||
|
||||
max_size = 800
|
||||
if max_h > max_size:
|
||||
max_w = int( max_w / (max_h / max_size) )
|
||||
|
|
@ -162,101 +162,101 @@ def previewThread (input_queue, output_queue):
|
|||
update_preview = True
|
||||
elif op == 'close':
|
||||
break
|
||||
|
||||
|
||||
if update_preview:
|
||||
update_preview = False
|
||||
(h,w,c) = previews[0][1].shape
|
||||
|
||||
|
||||
selected_preview_name = previews[selected_preview][0]
|
||||
selected_preview_rgb = previews[selected_preview][1]
|
||||
|
||||
|
||||
# HEAD
|
||||
head_text_color = [0.8]*c
|
||||
head_lines = [
|
||||
'[s]:save [enter]:exit',
|
||||
'[p]:update [space]:next preview',
|
||||
'Preview: "%s" [%d/%d]' % (selected_preview_name,selected_preview+1, len(previews) )
|
||||
]
|
||||
]
|
||||
head_line_height = 15
|
||||
head_height = len(head_lines) * head_line_height
|
||||
head = np.ones ( (head_height,w,c) ) * 0.1
|
||||
|
||||
|
||||
for i in range(0, len(head_lines)):
|
||||
t = i*head_line_height
|
||||
b = (i+1)*head_line_height
|
||||
head[t:b, 0:w] += image_utils.get_text_image ( (w,head_line_height,c) , head_lines[i], color=head_text_color )
|
||||
|
||||
|
||||
final = head
|
||||
|
||||
|
||||
if loss_history is not None:
|
||||
# LOSS HISTORY
|
||||
loss_history = np.array (loss_history)
|
||||
|
||||
|
||||
lh_height = 100
|
||||
lh_img = np.ones ( (lh_height,w,c) ) * 0.1
|
||||
loss_count = len(loss_history[0])
|
||||
lh_len = len(loss_history)
|
||||
|
||||
l_per_col = lh_len / w
|
||||
plist_max = [ [ max (0.0, 0.0, *[ loss_history[i_ab][p]
|
||||
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
|
||||
|
||||
l_per_col = lh_len / w
|
||||
plist_max = [ [ max (0.0, 0.0, *[ 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(0,loss_count)
|
||||
]
|
||||
for col in range(0, w)
|
||||
]
|
||||
|
||||
|
||||
plist_min = [ [ min (plist_max[col][p],
|
||||
plist_max[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(0,loss_count)
|
||||
]
|
||||
for col in range(0, w)
|
||||
]
|
||||
|
||||
|
||||
plist_min = [ [ min (plist_max[col][p],
|
||||
plist_max[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(0,loss_count)
|
||||
]
|
||||
for col in range(0, w)
|
||||
)
|
||||
for p in range(0,loss_count)
|
||||
]
|
||||
for col in range(0, w)
|
||||
]
|
||||
plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2
|
||||
|
||||
|
||||
if l_per_col >= 1.0:
|
||||
for col in range(0, w):
|
||||
for p in range(0,loss_count):
|
||||
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)
|
||||
|
||||
|
||||
if epoch != 0:
|
||||
lh_text = 'Loss history. Epoch: %d' % (epoch)
|
||||
else:
|
||||
lh_text = 'Loss history.'
|
||||
|
||||
|
||||
lh_img[last_line_t:last_line_b, 0:w] += image_utils.get_text_image ( (w,last_line_b-last_line_t,c), lh_text, color=head_text_color )
|
||||
|
||||
|
||||
final = np.concatenate ( [final, lh_img], axis=0 )
|
||||
|
||||
|
||||
final = np.concatenate ( [final, selected_preview_rgb], axis=0 )
|
||||
|
||||
|
||||
cv2.imshow ( 'Training preview', final)
|
||||
is_showing = True
|
||||
|
||||
|
||||
if is_showing:
|
||||
key = cv2.waitKey(100)
|
||||
else:
|
||||
|
|
@ -274,16 +274,16 @@ def previewThread (input_queue, output_queue):
|
|||
elif key == ord(' '):
|
||||
selected_preview = (selected_preview + 1) % len(previews)
|
||||
update_preview = True
|
||||
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
def main (training_data_src_dir, training_data_dst_dir, model_path, model_name, **in_options):
|
||||
|
||||
def main (training_data_src_dir, training_data_dst_dir, model_path, model_name,preview, **in_options):
|
||||
print ("Running trainer.\r\n")
|
||||
|
||||
|
||||
output_queue = queue.Queue()
|
||||
input_queue = queue.Queue()
|
||||
import threading
|
||||
thread = threading.Thread(target=trainerThread, args=(output_queue, input_queue, training_data_src_dir, training_data_dst_dir, model_path, model_name), kwargs=in_options )
|
||||
thread.start()
|
||||
|
||||
previewThread (input_queue, output_queue)
|
||||
if preview:
|
||||
previewThread (input_queue, output_queue)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue