OSSForks/faceswap
1
0
Fork 0
mirror of https://github.com/zebrajr/faceswap.git synced 2025-12-06 00:20:09 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

bugfix: Train: Toggle mask when 'learn_mask' is selected

Browse Source
This commit is contained in:
torzdf 2022-09-02 11:50:58 +01:00
parent dc18c74eea
commit 477e3e2013
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
plugins/train/trainer/_base.py
View File

@ -789,8 +789,13 @@ class _Samples(): # pylint:disable=too-few-public-methods
full = self._process_full(side, full, predictions[0].shape[1], (0., 0., 1.0))
images = [faces] + predictions
if self._display_mask:
images = self._compile_masked(images, samples[-1])
elif self._model.config["learn_mask"]:
# Remove masks when learn mask is selected but mask toggle is off
images = [batch[..., :3] for batch in images]
images = [self._overlay_foreground(full.copy(), image) for image in images]
return images