github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-07 05:22:06 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix DFLJPG,

Browse source 
SAE: added "rare sample booster"
SAE: pixel loss replaced to smooth transition from DSSIM to PixelLoss in 15k epochs by default
This commit is contained in:
iperov 2019-02-09 18:53:37 +04:00
parent f93b4713a9
commit 4d37fd62cd
11 changed files with 174 additions and 101 deletions
Download patch file Download diff file Expand all files Collapse all files

105
samples/SampleGeneratorFace.py
View file

@ -2,7 +2,7 @@ import traceback
import numpy as np
import random
import cv2
import multiprocessing
from utils import iter_utils
from samples import SampleType
@ -18,10 +18,11 @@ output_sample_types = [
]
'''
class SampleGeneratorFace(SampleGeneratorBase):
def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, with_close_to_self=False, sample_process_options=SampleProcessor.Options(), output_sample_types=[], generators_count=2, **kwargs):
def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, with_close_to_self=False, sample_process_options=SampleProcessor.Options(), output_sample_types=[], add_sample_idx=False, generators_count=2, **kwargs):
super().__init__(samples_path, debug, batch_size)
self.sample_process_options = sample_process_options
self.output_sample_types = output_sample_types
self.add_sample_idx = add_sample_idx
if sort_by_yaw_target_samples_path is not None:
self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET
@ -34,13 +35,15 @@ class SampleGeneratorFace(SampleGeneratorBase):
self.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)
self.generators_count = min ( generators_count, len(self.samples) )
if self.debug:
self.generators_count = 1
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )]
else:
self.generators_count = min ( generators_count, len(self.samples) )
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, i ) for i in range(self.generators_count) ]
self.generators_sq = [ multiprocessing.Queue() for _ in range(self.generators_count) ]
self.generator_counter = -1
def __iter__(self):
@ -50,47 +53,73 @@ class SampleGeneratorFace(SampleGeneratorBase):
self.generator_counter += 1
generator = self.generators[self.generator_counter % len(self.generators) ]
return next(generator)
def batch_func(self, generator_id):
samples = self.samples[generator_id::self.generators_count]
data_len = len(samples)
if data_len == 0:
def repeat_sample_idxs(self, idxs): # [ idx, ... ]
#send idxs list to all sub generators.
for gen_sq in self.generators_sq:
gen_sq.put (idxs)
def batch_func(self, generator_id):
gen_sq = self.generators_sq[generator_id]
samples = self.samples
samples_len = len(samples)
samples_idxs = [ *range(samples_len) ] [generator_id::self.generators_count]
repeat_samples_idxs = []
if len(samples_idxs) == 0:
raise ValueError('No training data provided.')
if self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
if all ( [ x == None for x in samples] ):
if all ( [ samples[idx] == None for idx in samples_idxs] ):
raise ValueError('Not enough training data. Gather more faces!')
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
shuffle_idxs = []
shuffle_idxs = []
elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
shuffle_idxs = []
shuffle_idxs_2D = [[]]*data_len
while True:
shuffle_idxs = []
shuffle_idxs_2D = [[]]*samples_len
while True:
while not gen_sq.empty():
idxs = gen_sq.get()
for idx in idxs:
if idx in samples_idxs:
repeat_samples_idxs.append(idx)
batches = None
for n_batch in range(self.batch_size):
while True:
sample = None
if len(repeat_samples_idxs) > 0:
idx = repeat_samples_idxs.pop()
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
sample = samples[idx]
elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
sample = samples[(idx >> 16) & 0xFFFF][idx & 0xFFFF]
else:
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
if len(shuffle_idxs) == 0:
shuffle_idxs = samples_idxs.copy()
np.random.shuffle(shuffle_idxs)
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
if len(shuffle_idxs) == 0:
shuffle_idxs = random.sample( range(data_len), data_len )
idx = shuffle_idxs.pop()
sample = samples[ idx ]
elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
if len(shuffle_idxs) == 0:
shuffle_idxs = random.sample( range(data_len), data_len )
idx = shuffle_idxs.pop()
if samples[idx] != None:
if len(shuffle_idxs_2D[idx]) == 0:
shuffle_idxs_2D[idx] = random.sample( range(len(samples[idx])), len(samples[idx]) )
idx = shuffle_idxs.pop()
sample = samples[ idx ]
elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
if len(shuffle_idxs) == 0:
shuffle_idxs = samples_idxs.copy()
np.random.shuffle(shuffle_idxs)
idx = shuffle_idxs.pop()
if samples[idx] != None:
if len(shuffle_idxs_2D[idx]) == 0:
shuffle_idxs_2D[idx] = random.sample( range(len(samples[idx])), len(samples[idx]) )
idx2 = shuffle_idxs_2D[idx].pop()
sample = samples[idx][idx2]
idx2 = shuffle_idxs_2D[idx].pop()
sample = samples[idx][idx2]
idx = (idx << 16) | (idx2 & 0xFFFF)
if sample is not None:
try:
@ -103,10 +132,14 @@ class SampleGeneratorFace(SampleGeneratorBase):
if batches is None:
batches = [ [] for _ in range(len(x)) ]
if self.add_sample_idx:
batches += [ [] ]
for i in range(len(x)):
batches[i].append ( x[i] )
if self.add_sample_idx:
batches[-1].append (idx)
break
yield [ np.array(batch) for batch in batches]