| import copy |
| import multiprocessing |
| 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) |
|
|
|
|
|
|
| class Index2DHost(): |
| """ |
| Provides random shuffled 2D indexes for multiprocesses |
| """ |
| def __init__(self, indexes2D): |
| self.sq = multiprocessing.Queue() |
| self.cqs = [] |
| self.clis = [] |
| self.thread = threading.Thread(target=self.host_thread, args=(indexes2D,) ) |
| self.thread.daemon = True |
| self.thread.start() |
|
|
| def host_thread(self, indexes2D): |
| indexes_counts_len = len(indexes2D) |
|
|
| idxs = [*range(indexes_counts_len)] |
| idxs_2D = [None]*indexes_counts_len |
| shuffle_idxs = [] |
| shuffle_idxs_2D = [None]*indexes_counts_len |
| for i in range(indexes_counts_len): |
| idxs_2D[i] = indexes2D[i] |
| shuffle_idxs_2D[i] = [] |
|
|
| sq = self.sq |
|
|
| while True: |
| while not sq.empty(): |
| obj = sq.get() |
| cq_id, cmd = obj[0], obj[1] |
|
|
| if cmd == 0: |
| count = obj[2] |
|
|
| result = [] |
| for i in range(count): |
| if len(shuffle_idxs) == 0: |
| shuffle_idxs = idxs.copy() |
| np.random.shuffle(shuffle_idxs) |
| result.append(shuffle_idxs.pop()) |
| self.cqs[cq_id].put (result) |
| elif cmd == 1: |
| targ_idxs,count = obj[2], obj[3] |
| result = [] |
|
|
| for targ_idx in targ_idxs: |
| sub_idxs = [] |
| for i in range(count): |
| ar = shuffle_idxs_2D[targ_idx] |
| if len(ar) == 0: |
| ar = shuffle_idxs_2D[targ_idx] = idxs_2D[targ_idx].copy() |
| np.random.shuffle(ar) |
| sub_idxs.append(ar.pop()) |
| result.append (sub_idxs) |
| self.cqs[cq_id].put (result) |
|
|
| time.sleep(0.001) |
|
|
| def create_cli(self): |
| cq = multiprocessing.Queue() |
| self.cqs.append ( cq ) |
| cq_id = len(self.cqs)-1 |
| return Index2DHost.Cli(self.sq, cq, cq_id) |
|
|
| |
| def __getstate__(self): |
| return dict() |
| def __setstate__(self, d): |
| self.__dict__.update(d) |
|
|
| class Cli(): |
| def __init__(self, sq, cq, cq_id): |
| self.sq = sq |
| self.cq = cq |
| self.cq_id = cq_id |
|
|
| def get_1D(self, count): |
| self.sq.put ( (self.cq_id,0, count) ) |
|
|
| while True: |
| if not self.cq.empty(): |
| return self.cq.get() |
| time.sleep(0.001) |
|
|
| def get_2D(self, idxs, count): |
| self.sq.put ( (self.cq_id,1,idxs,count) ) |
|
|
| while True: |
| if not self.cq.empty(): |
| return self.cq.get() |
| time.sleep(0.001) |
| |
| ''' |
| arg |
| output_sample_types = [ |
| [SampleProcessor.TypeFlags, size, (optional) {} opts ] , |
| ... |
| ] |
| ''' |
| class SampleGeneratorFacePerson(SampleGeneratorBase): |
| def __init__ (self, samples_path, debug=False, batch_size=1, |
| sample_process_options=SampleProcessor.Options(), |
| output_sample_types=[], |
| person_id_mode=1, |
| **kwargs): |
|
|
| super().__init__(debug, batch_size) |
| self.sample_process_options = sample_process_options |
| self.output_sample_types = output_sample_types |
| self.person_id_mode = person_id_mode |
|
|
| raise NotImplementedError("Currently SampleGeneratorFacePerson is not implemented.") |
|
|
| samples_host = SampleLoader.mp_host (SampleType.FACE, samples_path) |
| samples = samples_host.get_list() |
| self.samples_len = len(samples) |
|
|
| if self.samples_len == 0: |
| raise ValueError('No training data provided.') |
|
|
| unique_person_names = { sample.person_name for sample in samples } |
| persons_name_idxs = { person_name : [] for person_name in unique_person_names } |
| for i,sample in enumerate(samples): |
| persons_name_idxs[sample.person_name].append (i) |
| indexes2D = [ persons_name_idxs[person_name] for person_name in unique_person_names ] |
| index2d_host = Index2DHost(indexes2D) |
|
|
| if self.debug: |
| self.generators_count = 1 |
| self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, (samples_host.create_cli(), index2d_host.create_cli(),) )] |
| else: |
| self.generators_count = np.clip(multiprocessing.cpu_count(), 2, 4) |
| self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, (samples_host.create_cli(), index2d_host.create_cli(),) ) 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 ): |
| samples, index2d_host, = param |
| bs = self.batch_size |
|
|
| while True: |
| person_idxs = index2d_host.get_1D(bs) |
| samples_idxs = index2d_host.get_2D(person_idxs, 1) |
|
|
| batches = None |
| for n_batch in range(bs): |
| person_id = person_idxs[n_batch] |
| sample_idx = samples_idxs[n_batch][0] |
|
|
| sample = samples[ sample_idx ] |
| try: |
| x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug) |
| 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)) ] |
|
|
| batches += [ [] ] |
| i_person_id = len(batches)-1 |
|
|
| for i in range(len(x)): |
| batches[i].append ( x[i] ) |
|
|
| batches[i_person_id].append ( np.array([person_id]) ) |
|
|
| yield [ np.array(batch) for batch in batches] |
|
|
| @staticmethod |
| def get_person_id_max_count(samples_path): |
| return SampleLoader.get_person_id_max_count(samples_path) |
|
|
| """ |
| if self.person_id_mode==1: |
| samples_len = len(samples) |
| samples_idxs = [*range(samples_len)] |
| shuffle_idxs = [] |
| elif self.person_id_mode==2: |
| persons_count = len(samples) |
| |
| person_idxs = [] |
| for j in range(persons_count): |
| for i in range(j+1,persons_count): |
| person_idxs += [ [i,j] ] |
| |
| shuffle_person_idxs = [] |
| |
| samples_idxs = [None]*persons_count |
| shuffle_idxs = [None]*persons_count |
| |
| for i in range(persons_count): |
| samples_idxs[i] = [*range(len(samples[i]))] |
| shuffle_idxs[i] = [] |
| elif self.person_id_mode==3: |
| persons_count = len(samples) |
| |
| person_idxs = [ *range(persons_count) ] |
| shuffle_person_idxs = [] |
| |
| samples_idxs = [None]*persons_count |
| shuffle_idxs = [None]*persons_count |
| |
| for i in range(persons_count): |
| samples_idxs[i] = [*range(len(samples[i]))] |
| shuffle_idxs[i] = [] |
| |
| if self.person_id_mode==2: |
| if len(shuffle_person_idxs) == 0: |
| shuffle_person_idxs = person_idxs.copy() |
| np.random.shuffle(shuffle_person_idxs) |
| person_ids = shuffle_person_idxs.pop() |
| |
| |
| batches = None |
| for n_batch in range(self.batch_size): |
| |
| if self.person_id_mode==1: |
| if len(shuffle_idxs) == 0: |
| shuffle_idxs = samples_idxs.copy() |
| np.random.shuffle(shuffle_idxs) ### |
| |
| idx = shuffle_idxs.pop() |
| sample = samples[ idx ] |
| |
| try: |
| x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug) |
| except: |
| raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) |
| |
| if type(x) != tuple and type(x) != list: |
| raise Exception('SampleProcessor.process returns NOT tuple/list') |
| |
| if batches is None: |
| batches = [ [] for _ in range(len(x)) ] |
| |
| batches += [ [] ] |
| i_person_id = len(batches)-1 |
| |
| for i in range(len(x)): |
| batches[i].append ( x[i] ) |
| |
| batches[i_person_id].append ( np.array([sample.person_id]) ) |
| |
| |
| elif self.person_id_mode==2: |
| person_id1, person_id2 = person_ids |
| |
| if len(shuffle_idxs[person_id1]) == 0: |
| shuffle_idxs[person_id1] = samples_idxs[person_id1].copy() |
| np.random.shuffle(shuffle_idxs[person_id1]) |
| |
| idx = shuffle_idxs[person_id1].pop() |
| sample1 = samples[person_id1][idx] |
| |
| if len(shuffle_idxs[person_id2]) == 0: |
| shuffle_idxs[person_id2] = samples_idxs[person_id2].copy() |
| np.random.shuffle(shuffle_idxs[person_id2]) |
| |
| idx = shuffle_idxs[person_id2].pop() |
| sample2 = samples[person_id2][idx] |
| |
| if sample1 is not None and sample2 is not None: |
| try: |
| x1, = SampleProcessor.process ([sample1], self.sample_process_options, self.output_sample_types, self.debug) |
| except: |
| raise Exception ("Exception occured in sample %s. Error: %s" % (sample1.filename, traceback.format_exc() ) ) |
| |
| try: |
| x2, = SampleProcessor.process ([sample2], self.sample_process_options, self.output_sample_types, self.debug) |
| except: |
| raise Exception ("Exception occured in sample %s. Error: %s" % (sample2.filename, traceback.format_exc() ) ) |
| |
| x1_len = len(x1) |
| if batches is None: |
| batches = [ [] for _ in range(x1_len) ] |
| batches += [ [] ] |
| i_person_id1 = len(batches)-1 |
| |
| batches += [ [] for _ in range(len(x2)) ] |
| batches += [ [] ] |
| i_person_id2 = len(batches)-1 |
| |
| for i in range(x1_len): |
| batches[i].append ( x1[i] ) |
| |
| for i in range(len(x2)): |
| batches[x1_len+1+i].append ( x2[i] ) |
| |
| batches[i_person_id1].append ( np.array([sample1.person_id]) ) |
| |
| batches[i_person_id2].append ( np.array([sample2.person_id]) ) |
| |
| elif self.person_id_mode==3: |
| if len(shuffle_person_idxs) == 0: |
| shuffle_person_idxs = person_idxs.copy() |
| np.random.shuffle(shuffle_person_idxs) |
| person_id = shuffle_person_idxs.pop() |
| |
| if len(shuffle_idxs[person_id]) == 0: |
| shuffle_idxs[person_id] = samples_idxs[person_id].copy() |
| np.random.shuffle(shuffle_idxs[person_id]) |
| |
| idx = shuffle_idxs[person_id].pop() |
| sample1 = samples[person_id][idx] |
| |
| if len(shuffle_idxs[person_id]) == 0: |
| shuffle_idxs[person_id] = samples_idxs[person_id].copy() |
| np.random.shuffle(shuffle_idxs[person_id]) |
| |
| idx = shuffle_idxs[person_id].pop() |
| sample2 = samples[person_id][idx] |
| |
| if sample1 is not None and sample2 is not None: |
| try: |
| x1, = SampleProcessor.process ([sample1], self.sample_process_options, self.output_sample_types, self.debug) |
| except: |
| raise Exception ("Exception occured in sample %s. Error: %s" % (sample1.filename, traceback.format_exc() ) ) |
| |
| try: |
| x2, = SampleProcessor.process ([sample2], self.sample_process_options, self.output_sample_types, self.debug) |
| except: |
| raise Exception ("Exception occured in sample %s. Error: %s" % (sample2.filename, traceback.format_exc() ) ) |
| |
| x1_len = len(x1) |
| if batches is None: |
| batches = [ [] for _ in range(x1_len) ] |
| batches += [ [] ] |
| i_person_id1 = len(batches)-1 |
| |
| batches += [ [] for _ in range(len(x2)) ] |
| batches += [ [] ] |
| i_person_id2 = len(batches)-1 |
| |
| for i in range(x1_len): |
| batches[i].append ( x1[i] ) |
| |
| for i in range(len(x2)): |
| batches[x1_len+1+i].append ( x2[i] ) |
| |
| batches[i_person_id1].append ( np.array([sample1.person_id]) ) |
| |
| batches[i_person_id2].append ( np.array([sample2.person_id]) ) |
| """ |
|
|
