mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-07 10:43:27 -04:00
7f53911453
* Convert prints to logger. Further logging improvements. Tidy up * Fix system verbosity. Allow SystemExit * Fix reload extract bug * Child Traceback handling * Safer shutdown procedure * Add shutdown event to queue manager * landmarks_as_xy > property. GUI notes + linting. Aligner bugfix * fix FaceFilter. Enable nFilter when no Filter is supplied * Fix blurry face filter * Continue on IO error. Better error handling * Explicitly print stack trace tocrash log * Windows Multiprocessing bugfix * Add git info and conda version to crash log * Windows/Anaconda mp bugfix * Logging fixes for training
53 lines
2 KiB
Python
53 lines
2 KiB
Python
# Improved-AutoEncoder base classes
|
|
|
|
import logging
|
|
|
|
from lib.utils import backup_file
|
|
|
|
hdf = {'encoderH5': 'IAE_encoder.h5',
|
|
'decoderH5': 'IAE_decoder.h5',
|
|
'inter_AH5': 'IAE_inter_A.h5',
|
|
'inter_BH5': 'IAE_inter_B.h5',
|
|
'inter_bothH5': 'IAE_inter_both.h5'}
|
|
|
|
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
class AutoEncoder:
|
|
def __init__(self, model_dir, gpus):
|
|
self.model_dir = model_dir
|
|
self.gpus = gpus
|
|
|
|
self.encoder = self.Encoder()
|
|
self.decoder = self.Decoder()
|
|
self.inter_A = self.Intermidiate()
|
|
self.inter_B = self.Intermidiate()
|
|
self.inter_both = self.Intermidiate()
|
|
|
|
self.initModel()
|
|
|
|
def load(self, swapped):
|
|
(face_A,face_B) = (hdf['inter_AH5'], hdf['inter_BH5']) if not swapped else (hdf['inter_BH5'], hdf['inter_AH5'])
|
|
|
|
try:
|
|
self.encoder.load_weights(str(self.model_dir / hdf['encoderH5']))
|
|
self.decoder.load_weights(str(self.model_dir / hdf['decoderH5']))
|
|
self.inter_both.load_weights(str(self.model_dir / hdf['inter_bothH5']))
|
|
self.inter_A.load_weights(str(self.model_dir / face_A))
|
|
self.inter_B.load_weights(str(self.model_dir / face_B))
|
|
logger.info('loaded model weights')
|
|
return True
|
|
except Exception:
|
|
logger.warning('Failed loading existing training data. Starting a fresh model: %s', self.model_dir)
|
|
return False
|
|
|
|
def save_weights(self):
|
|
model_dir = str(self.model_dir)
|
|
for model in hdf.values():
|
|
backup_file(model_dir, model)
|
|
self.encoder.save_weights(str(self.model_dir / hdf['encoderH5']))
|
|
self.decoder.save_weights(str(self.model_dir / hdf['decoderH5']))
|
|
self.inter_both.save_weights(str(self.model_dir / hdf['inter_bothH5']))
|
|
self.inter_A.save_weights(str(self.model_dir / hdf['inter_AH5']))
|
|
self.inter_B.save_weights(str(self.model_dir / hdf['inter_BH5']))
|
|
logger.info('saved model weights')
|