| import multiprocessing |
| import pickle |
| import time |
| import traceback |
|
|
| import cv2 |
| import numpy as np |
|
|
| from core import mplib |
| from core.joblib import SubprocessGenerator, ThisThreadGenerator |
| from facelib import LandmarksProcessor |
| from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor, |
| SampleType) |
|
|
|
|
| ''' |
| arg |
| output_sample_types = [ |
| [SampleProcessor.TypeFlags, size, (optional) {} opts ] , |
| ... |
| ] |
| ''' |
| class SampleGeneratorFaceDebug(SampleGeneratorBase): |
| def __init__ (self, samples_path, debug=False, batch_size=1, |
| random_ct_samples_path=None, |
| sample_process_options=SampleProcessor.Options(), |
| output_sample_types=[], |
| add_sample_idx=False, |
| generators_count=4, |
| rnd_seed=None, |
| **kwargs): |
|
|
| super().__init__(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 rnd_seed is None: |
| rnd_seed = np.random.randint(0x80000000) |
|
|
| if self.debug: |
| self.generators_count = 1 |
| else: |
| self.generators_count = max(1, generators_count) |
|
|
| samples = SampleLoader.load (SampleType.FACE, samples_path) |
| self.samples_len = len(samples) |
|
|
| if self.samples_len == 0: |
| raise ValueError('No training data provided.') |
|
|
| if random_ct_samples_path is not None: |
| ct_samples = SampleLoader.load (SampleType.FACE, random_ct_samples_path) |
| else: |
| ct_samples = None |
|
|
| pickled_samples = pickle.dumps(samples, 4) |
| ct_pickled_samples = pickle.dumps(ct_samples, 4) if ct_samples is not None else None |
|
|
| if self.debug: |
| self.generators = [ThisThreadGenerator ( self.batch_func, (pickled_samples, ct_pickled_samples, rnd_seed) )] |
| else: |
| self.generators = [SubprocessGenerator ( self.batch_func, (pickled_samples, ct_pickled_samples, rnd_seed+i), start_now=False ) \ |
| for i in range(self.generators_count) ] |
| |
| SubprocessGenerator.start_in_parallel( self.generators ) |
|
|
| self.generator_counter = -1 |
|
|
| def __iter__(self): |
| return self |
|
|
| def __next__(self): |
| self.generator_counter += 1 |
| generator = self.generators[self.generator_counter % len(self.generators) ] |
| return next(generator) |
|
|
| def batch_func(self, param ): |
| pickled_samples, ct_pickled_samples, rnd_seed = param |
| |
| rnd_state = np.random.RandomState(rnd_seed) |
|
|
| samples = pickle.loads(pickled_samples) |
| idxs = [*range(len(samples))] |
| shuffle_idxs = [] |
| |
| if ct_pickled_samples is not None: |
| ct_samples = pickle.loads(ct_pickled_samples) |
| ct_idxs = [*range(len(ct_samples))] |
| ct_shuffle_idxs = [] |
| else: |
| ct_samples = None |
| |
|
|
| bs = self.batch_size |
| while True: |
| batches = None |
|
|
| for n_batch in range(bs): |
| |
| if len(shuffle_idxs) == 0: |
| shuffle_idxs = idxs.copy() |
| rnd_state.shuffle(shuffle_idxs) |
| |
| sample_idx = shuffle_idxs.pop() |
| sample = samples[sample_idx] |
|
|
| ct_sample = None |
| if ct_samples is not None: |
| if len(ct_shuffle_idxs) == 0: |
| ct_shuffle_idxs = ct_idxs.copy() |
| rnd_state.shuffle(ct_shuffle_idxs) |
| ct_sample_idx = ct_shuffle_idxs.pop() |
| ct_sample = ct_samples[ct_sample_idx] |
|
|
| try: |
| x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug, ct_sample=ct_sample, rnd_state=rnd_state) |
| except: |
| raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) |
|
|
| if batches is None: |
| batches = [ [] for _ in range(len(x)) ] |
| if self.add_sample_idx: |
| batches += [ [] ] |
| i_sample_idx = len(batches)-1 |
|
|
| for i in range(len(x)): |
| batches[i].append ( x[i] ) |
|
|
| if self.add_sample_idx: |
| batches[i_sample_idx].append (sample_idx) |
|
|
| yield [ np.array(batch) for batch in batches] |
|
|
