def threaded_io(self, task, io_args=None):
""" Perform I/O task in a background thread """
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
if task == "load":
func = self.load_images
elif task == "save":
func = self.save_faces
elif task == "reload":
func = self.reload_images
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
return io_thread
def _start_thread(self):
""" Put the :func:`_training` into a background thread so we can keep control.
Returns
-------
:class:`lib.multithreading.MultiThread`
The background thread for running training
"""
logger.debug("Launching Trainer thread")
thread = MultiThread(target=self._training)
thread.start()
logger.debug("Launched Trainer thread")
return thread
def _launch_predictor(self):
""" Launch the prediction process in a background thread.
Starts the prediction thread and returns the thread.
Returns
-------
:class:`~lib.multithreading.MultiThread`
The started Faceswap model prediction thread.
"""
thread = MultiThread(self._predict_faces, thread_count=1)
thread.start()
return thread
def save_models(self):
""" Backup and save the models """
logger.debug("Backing up and saving models")
should_backup = self.get_save_averages()
save_threads = list()
for network in self.networks.values():
name = "save_{}".format(network.name)
save_threads.append(
MultiThread(network.save,
name=name,
should_backup=should_backup))
save_threads.append(
MultiThread(self.state.save,
name="save_state",
should_backup=should_backup))
for thread in save_threads:
thread.start()
for thread in save_threads:
if thread.has_error:
logger.error(thread.errors[0])
thread.join()
# Put in a line break to avoid jumbled console
print("\n")
logger.info("saved models")
def _start_thread(self, task):
""" Create the thread for the given task, add it it :attr:`self._threads` and start it.
Parameters
----------
task: {"load", "save"}
The task that the thread is to be created for
"""
logger.debug("Starting thread: '%s'", task)
args = self._completion_event if task == "save" else None
func = getattr(self, "_{}".format(task))
io_thread = MultiThread(func, args, thread_count=1)
io_thread.start()
self._threads[task] = io_thread
logger.debug("Started thread: '%s'", task)
def save_models(self, snapshot_iteration):
""" Backup and save the models """
logger.debug("Backing up and saving models")
should_backup = self.get_save_averages()
save_threads = list()
for network in self.networks.values():
name = "save_{}".format(network.name)
save_threads.append(
MultiThread(network.save,
name=name,
should_backup=should_backup))
save_threads.append(
MultiThread(self.state.save,
name="save_state",
should_backup=should_backup))
for thread in save_threads:
thread.start()
for thread in save_threads:
if thread.has_error:
logger.error(thread.errors[0])
thread.join()
logger.info("saved models")
if snapshot_iteration:
self.snapshot_models()
def _threaded_redirector(self, task, io_args=None):
""" Redirect image input/output tasks to relevant queues in background thread
Parameters
----------
task: str
The name of the task to be put into a background thread
io_args: tuple, optional
Any arguments that need to be provided to the background function
"""
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
func = getattr(self, "_{}".format(task))
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
self._threads.append(io_thread)
def __init__(self, in_queue, queue_size, arguments):
logger.debug("Initializing %s: (args: %s, queue_size: %s, in_queue: %s)",
self.__class__.__name__, arguments, queue_size, in_queue)
self.batchsize = self.get_batchsize(queue_size)
self.args = arguments
self.in_queue = in_queue
self.out_queue = queue_manager.get_queue("patch")
self.serializer = Serializer.get_serializer("json")
self.faces_count = 0
self.verify_output = False
self.model = self.load_model()
self.predictor = self.model.converter(self.args.swap_model)
self.queues = dict()
self.thread = MultiThread(self.predict_faces, thread_count=1)
self.thread.start()
logger.debug("Initialized %s: (out_queue: %s)", self.__class__.__name__, self.out_queue)
def _set_thread(self, io_args=None):
""" Set the load/save thread
Parameters
----------
io_args: tuple, optional
The arguments to be passed to the load or save thread. Default: `None`.
Returns
-------
:class:`lib.multithreading.MultiThread`: Thread containing the load/save function.
"""
io_args = (self._queue) if io_args is None else (self._queue, *io_args)
retval = MultiThread(getattr(self, "_{}".format(self._task)),
*io_args,
thread_count=1)
logger.trace(retval)
return retval
def _feed_extractor(self):
""" Feed the input queue to the Extractor from a faces folder or from source frames in a
background thread
Returns
-------
:class:`lib.multithreading.Multithread`:
The thread that is feeding the extractor.
"""
masker_input = getattr(self,
"_input_{}".format("faces" if self._input_is_faces else "frames"))
logger.debug("masker_input: %s", masker_input)
args = tuple() if self._update_type == "output" else (self._extractor.input_queue, )
input_thread = MultiThread(masker_input, *args, thread_count=1)
input_thread.start()
logger.debug(input_thread)
return input_thread
def _launch_folder(self):
""" Launch :class:`lib.multithreading.MultiThread` to retrieve faces from a
folder of images.
Goes through the file list one at a time, passing each file to a separate background
thread for some speed up.
"""
reader = SingleFrameLoader(self._location)
num_threads = min(reader.count, self._num_threads)
frame_split = reader.count // self._num_threads
logger.debug("total images: %s, num_threads: %s, frames_per_thread: %s",
reader.count, num_threads, frame_split)
for idx in range(num_threads):
is_final = idx == num_threads - 1
start_idx = idx * frame_split
end_idx = reader.count if is_final else start_idx + frame_split
thread = MultiThread(self._load_from_folder, reader, start_idx, end_idx)
thread.start()
self._threads.append(thread)
def extract(self):
""" Extract the current faces to a folder.
To stop the GUI becoming completely unresponsive (particularly in Windows) the extract is
done in a background thread, with the process count passed back in a queue to the main
thread to update the progress bar.
"""
dirname = FileHandler("dir", None,
initial_folder=os.path.dirname(self._input_location),
title="Select output folder...").return_file
if not dirname:
return
logger.debug(dirname)
queue = Queue()
pbar = PopupProgress("Extracting Faces...", self._alignments.frames_count + 1)
thread = MultiThread(self._background_extract, dirname, queue)
thread.start()
self._monitor_extract(thread, queue, pbar)
def convert_images(self):
""" Convert the images """
logger.debug("Converting images")
save_queue = queue_manager.get_queue("convert_out")
patch_queue = queue_manager.get_queue("patch")
self.patch_threads = MultiThread(self.converter.process,
patch_queue,
save_queue,
thread_count=self.pool_processes,
name="patch")
self.patch_threads.start()
while True:
self.check_thread_error()
if self.disk_io.completion_event.is_set():
logger.debug("DiskIO completion event set. Joining Pool")
break
sleep(1)
self.patch_threads.join()
logger.debug("Putting EOF")
save_queue.put("EOF")
logger.debug("Converted images")
def __init__(self, arguments, samples, display, lock, trigger,
config_tools, tk_vars):
logger.debug(
"Initializing %s: (arguments: '%s', samples: %s: display: %s, lock: %s,"
" trigger: %s, config_tools: %s, tk_vars %s)",
self.__class__.__name__, arguments, samples, display, lock,
trigger, config_tools, tk_vars)
self.samples = samples
self.queue_patch_in = queue_manager.get_queue("preview_patch_in")
self.display = display
self.lock = lock
self.trigger = trigger
self.current_config = config_tools.config
self.converter_arguments = None # Updated converter arguments dict
configfile = arguments.configfile if hasattr(arguments,
"configfile") else None
self.converter = Converter(
output_dir=None,
output_size=self.samples.predictor.output_size,
output_has_mask=self.samples.predictor.has_predicted_mask,
draw_transparent=False,
pre_encode=None,
configfile=configfile,
arguments=self.generate_converter_arguments(arguments))
self.shutdown = Event()
self.thread = MultiThread(self.process,
self.trigger,
self.shutdown,
self.queue_patch_in,
self.samples,
tk_vars,
thread_count=1,
name="patch_thread")
self.thread.start()
def _convert_images(self):
""" Start the multi-threaded patching process, monitor all threads for errors and join on
completion. """
logger.debug("Converting images")
save_queue = queue_manager.get_queue("convert_out")
patch_queue = queue_manager.get_queue("patch")
self._patch_threads = MultiThread(self._converter.process, patch_queue, save_queue,
thread_count=self._pool_processes, name="patch")
self._patch_threads.start()
while True:
self._check_thread_error()
if self._disk_io.completion_event.is_set():
logger.debug("DiskIO completion event set. Joining Pool")
break
if self._patch_threads.completed():
logger.debug("All patch threads completed")
break
sleep(1)
self._patch_threads.join()
logger.debug("Putting EOF")
save_queue.put("EOF")
logger.debug("Converted images")
def in_thread(self, action):
""" Perform selected action inside a thread """
logger.debug("Performing help action: %s", action)
thread = MultiThread(getattr(self, action), thread_count=1)
thread.start()
logger.debug("Performed help action: %s", action)