mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-07 05:22:06 -07:00
interactive converter now saves the session.
if input frames are changed (amount or filenames) then interactive converter automatically starts a new session. if model is more trained then all frames will be recomputed again with their saved configs.
This commit is contained in:
Colombo
parent
bac9d5a99d
commit
bc0b25cdda
2 changed files with 102 additions and 24 deletions
|
|
@ -137,9 +137,9 @@ class InteractBase(object):
|
|||
else: print("capture_keys: already set for window ", wnd_name)
|
||||
else: print("capture_keys: named_window ", wnd_name, " not found.")
|
||||
|
||||
def progress_bar(self, desc, total, leave=True):
|
||||
def progress_bar(self, desc, total, leave=True, initial=0):
|
||||
if self.pg_bar is None:
|
||||
self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True )
|
||||
self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True, initial=initial )
|
||||
else: print("progress_bar: already set.")
|
||||
|
||||
def progress_bar_inc(self, c):
|
||||
|
|
@ -154,8 +154,8 @@ class InteractBase(object):
|
|||
self.pg_bar = None
|
||||
else: print("progress_bar not set.")
|
||||
|
||||
def progress_bar_generator(self, data, desc, leave=True):
|
||||
self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True )
|
||||
def progress_bar_generator(self, data, desc, leave=True, initial=0):
|
||||
self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True, initial=initial )
|
||||
for x in self.pg_bar:
|
||||
yield x
|
||||
self.pg_bar.close()
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
import multiprocessing
|
||||
import operator
|
||||
import os
|
||||
import pickle
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
|
|
@ -13,7 +14,8 @@ import numpy as np
|
|||
import numpy.linalg as npla
|
||||
|
||||
import imagelib
|
||||
from converters import FrameInfo, ConverterConfig, ConvertMasked, ConvertFaceAvatar
|
||||
from converters import (ConverterConfig, ConvertFaceAvatar, ConvertMasked,
|
||||
FrameInfo)
|
||||
from facelib import FaceType, FANSegmentator, LandmarksProcessor
|
||||
from interact import interact as io
|
||||
from joblib import SubprocessFunctionCaller, Subprocessor
|
||||
|
|
@ -177,7 +179,7 @@ class ConvertSubprocessor(Subprocessor):
|
|||
return pf.frame_info.filename
|
||||
|
||||
#override
|
||||
def __init__(self, is_interactive, converter_config, frames, output_path):
|
||||
def __init__(self, is_interactive, converter_config, frames, output_path, model_iter):
|
||||
if len (frames) == 0:
|
||||
raise ValueError ("len (frames) == 0")
|
||||
|
||||
|
|
@ -203,20 +205,84 @@ class ConvertSubprocessor(Subprocessor):
|
|||
|
||||
self.dcscn_host, self.superres_func = SubprocessFunctionCaller.make_pair(superres_func)
|
||||
|
||||
self.frames = frames
|
||||
self.output_path = output_path
|
||||
self.model_iter = model_iter
|
||||
|
||||
self.prefetch_frame_count = self.process_count = min(6,multiprocessing.cpu_count())
|
||||
|
||||
session_data = None
|
||||
session_dat_path = self.output_path / 'session.dat'
|
||||
if session_dat_path.exists():
|
||||
|
||||
try:
|
||||
with open( str(session_dat_path), "rb") as f:
|
||||
session_data = pickle.loads(f.read())
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
self.frames = frames
|
||||
self.frames_idxs = [ *range(len(self.frames)) ]
|
||||
self.frames_done_idxs = []
|
||||
|
||||
digits = [ str(i) for i in range(10)]
|
||||
if self.is_interactive and session_data is not None:
|
||||
s_frames = session_data.get('frames', None)
|
||||
s_frames_idxs = session_data.get('frames_idxs', None)
|
||||
s_frames_done_idxs = session_data.get('frames_done_idxs', None)
|
||||
s_model_iter = session_data.get('model_iter', None)
|
||||
|
||||
frames_equal = (s_frames is not None) and \
|
||||
(s_frames_idxs is not None) and \
|
||||
(s_frames_done_idxs is not None) and \
|
||||
(s_model_iter is not None) and \
|
||||
(len(frames) == len(s_frames))
|
||||
|
||||
if frames_equal:
|
||||
for i in range(len(frames)):
|
||||
frame = frames[i]
|
||||
s_frame = s_frames[i]
|
||||
if frame.frame_info.filename != s_frame.frame_info.filename:
|
||||
frames_equal = False
|
||||
if not frames_equal:
|
||||
break
|
||||
|
||||
if frames_equal:
|
||||
io.log_info ("Using saved session.")
|
||||
self.frames = s_frames
|
||||
self.frames_idxs = s_frames_idxs
|
||||
self.frames_done_idxs = s_frames_done_idxs
|
||||
|
||||
if self.model_iter != s_model_iter:
|
||||
#model is more trained, recompute all frames
|
||||
for frame in self.frames:
|
||||
frame.is_done = False
|
||||
|
||||
if self.model_iter != s_model_iter or \
|
||||
len(self.frames_idxs) == 0:
|
||||
#rewind to begin if model is more trained or all frames are done
|
||||
|
||||
while len(self.frames_done_idxs) > 0:
|
||||
prev_frame = self.frames[self.frames_done_idxs.pop()]
|
||||
self.frames_idxs.insert(0, prev_frame.idx)
|
||||
|
||||
if len(self.frames_idxs) != 0:
|
||||
cur_frame = self.frames[self.frames_idxs[0]]
|
||||
cur_frame.is_shown = False
|
||||
|
||||
if not frames_equal:
|
||||
session_data = None
|
||||
|
||||
if session_data is None:
|
||||
for filename in Path_utils.get_image_paths(self.output_path): #remove all images in output_path
|
||||
Path(filename).unlink()
|
||||
|
||||
frames[0].cfg = self.converter_config.copy()
|
||||
|
||||
for i in range( len(self.frames) ):
|
||||
frame = self.frames[i]
|
||||
frame.idx = i
|
||||
frame.output_filename = self.output_path / ( Path(frame.frame_info.filename).stem + '.png' )
|
||||
|
||||
frames[0].cfg = self.converter_config.copy()
|
||||
|
||||
|
||||
#override
|
||||
def process_info_generator(self):
|
||||
|
|
@ -232,7 +298,7 @@ class ConvertSubprocessor(Subprocessor):
|
|||
|
||||
#overridable optional
|
||||
def on_clients_initialized(self):
|
||||
io.progress_bar ("Converting", len (self.frames_idxs) )
|
||||
io.progress_bar ("Converting", len (self.frames_idxs), initial=len(self.frames_done_idxs) )
|
||||
|
||||
self.process_remain_frames = not self.is_interactive
|
||||
self.is_interactive_quitting = not self.is_interactive
|
||||
|
|
@ -251,11 +317,25 @@ class ConvertSubprocessor(Subprocessor):
|
|||
|
||||
#overridable optional
|
||||
def on_clients_finalized(self):
|
||||
io.progress_bar_close()
|
||||
|
||||
if self.is_interactive:
|
||||
self.screen_manager.finalize()
|
||||
|
||||
io.progress_bar_close()
|
||||
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,
|
||||
}
|
||||
save_path = self.output_path / 'session.dat'
|
||||
save_path.write_bytes( pickle.dumps(session_data) )
|
||||
|
||||
io.log_info ("Session is saved to " + '/'.join (save_path.parts[-2:]) )
|
||||
|
||||
cfg_change_keys = ['`','1', '2', '3', '4', '5', '6', '7', '8', '9',
|
||||
'q', 'a', 'w', 's', 'e', 'd', 'r', 'f', 't', 'g','y','h','u','j',
|
||||
|
|
@ -501,10 +581,7 @@ def main (args, device_args):
|
|||
io.log_err('Input directory not found. Please ensure it exists.')
|
||||
return
|
||||
|
||||
if output_path.exists():
|
||||
for filename in Path_utils.get_image_paths(output_path):
|
||||
Path(filename).unlink()
|
||||
else:
|
||||
if not output_path.exists():
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if not model_path.exists():
|
||||
|
|
@ -643,6 +720,7 @@ def main (args, device_args):
|
|||
converter_config = cfg,
|
||||
frames = frames,
|
||||
output_path = output_path,
|
||||
model_iter = model.get_iter()
|
||||
).run()
|
||||
|
||||
model.finalize()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue