diff --git a/samplelib/SampleGeneratorFaceTemporal.py b/samplelib/SampleGeneratorFaceTemporal.py new file mode 100644 index 0000000..b0280ad --- /dev/null +++ b/samplelib/SampleGeneratorFaceTemporal.py @@ -0,0 +1,88 @@ +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, SampleHost, SampleProcessor, + SampleType) + + +class SampleGeneratorFaceTemporal(SampleGeneratorBase): + def __init__ (self, samples_path, debug, batch_size, + temporal_image_count=3, + sample_process_options=SampleProcessor.Options(), + output_sample_types=[], + generators_count=2, + **kwargs): + super().__init__(samples_path, debug, batch_size) + + self.temporal_image_count = temporal_image_count + self.sample_process_options = sample_process_options + self.output_sample_types = output_sample_types + + if self.debug: + self.generators_count = 1 + else: + self.generators_count = generators_count + + samples = SampleHost.load (SampleType.FACE_TEMPORAL_SORTED, self.samples_path) + samples_len = len(samples) + if samples_len == 0: + raise ValueError('No training data provided.') + + mult_max = 1 + l = samples_len - ( (self.temporal_image_count)*mult_max - (mult_max-1) ) + index_host = mplib.IndexHost(l+1) + + pickled_samples = pickle.dumps(samples, 4) + if self.debug: + self.generators = [ThisThreadGenerator ( self.batch_func, (pickled_samples, index_host.create_cli(),) )] + else: + self.generators = [SubprocessGenerator ( self.batch_func, (pickled_samples, index_host.create_cli(),), start_now=True ) for i in range(self.generators_count) ] + + 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): + mult_max = 1 + bs = self.batch_size + pickled_samples, index_host = param + samples = pickle.loads(pickled_samples) + + while True: + batches = None + + indexes = index_host.multi_get(bs) + + for n_batch in range(self.batch_size): + idx = indexes[n_batch] + + temporal_samples = [] + mult = np.random.randint(mult_max)+1 + for i in range( self.temporal_image_count ): + sample = samples[ idx+i*mult ] + try: + temporal_samples += SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug)[0] + except: + raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) + + if batches is None: + batches = [ [] for _ in range(len(temporal_samples)) ] + + for i in range(len(temporal_samples)): + batches[i].append ( temporal_samples[i] ) + + yield [ np.array(batch) for batch in batches] diff --git a/samplelib/__init__.py b/samplelib/__init__.py index ecfbfec..67630c5 100644 --- a/samplelib/__init__.py +++ b/samplelib/__init__.py @@ -5,5 +5,6 @@ from .SampleProcessor import SampleProcessor from .SampleGeneratorBase import SampleGeneratorBase from .SampleGeneratorFace import SampleGeneratorFace from .SampleGeneratorFacePerson import SampleGeneratorFacePerson +from .SampleGeneratorFaceTemporal import SampleGeneratorFaceTemporal from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal from .PackedFaceset import PackedFaceset \ No newline at end of file