| | 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]) ) |
| | """ |
| |
|
