File size: 3,006 Bytes
fcd5579 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import traceback
import cv2
import numpy as np
from core.joblib import SubprocessGenerator, ThisThreadGenerator
from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor,
SampleType)
'''
output_sample_types = [
[SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
...
]
'''
class SampleGeneratorImageTemporal(SampleGeneratorBase):
def __init__ (self, samples_path, debug, batch_size, temporal_image_count, sample_process_options=SampleProcessor.Options(), output_sample_types=[], **kwargs):
super().__init__(debug, batch_size)
self.temporal_image_count = temporal_image_count
self.sample_process_options = sample_process_options
self.output_sample_types = output_sample_types
self.samples = SampleLoader.load (SampleType.IMAGE, samples_path)
self.generator_samples = [ self.samples ]
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )] if self.debug else \
[iter_utils.SubprocessGenerator ( self.batch_func, 0 )]
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, generator_id):
samples = self.generator_samples[generator_id]
samples_len = len(samples)
if samples_len == 0:
raise ValueError('No training data provided.')
mult_max = 4
samples_sub_len = samples_len - ( (self.temporal_image_count)*mult_max - (mult_max-1) )
if samples_sub_len <= 0:
raise ValueError('Not enough samples to fit temporal line.')
shuffle_idxs = []
while True:
batches = None
for n_batch in range(self.batch_size):
if len(shuffle_idxs) == 0:
shuffle_idxs = [ *range(samples_sub_len) ]
np.random.shuffle (shuffle_idxs)
idx = shuffle_idxs.pop()
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]
|