class DirectoryIterator(Iterator):
def __init__(self, directory, image_data_generator,
resize=None, target_size=None, color_mode='rgb',
dim_ordering='default',
classes=None, class_mode='categorical',
batch_size=32, shuffle=True, seed=None, gt_directory=None,
save_to_dir=None, save_prefix='', save_format='jpeg', model_name=None):
# Check dim order
if dim_ordering == 'default':
dim_ordering = K.image_dim_ordering()
self.dim_ordering = dim_ordering
self.directory = directory
self.gt_directory = gt_directory
self.image_data_generator = image_data_generator
self.resize = resize
self.save_to_dir = save_to_dir
self.save_prefix = save_prefix
self.save_format = save_format
self.model_name = model_name
# Check target size
if target_size is None and batch_size > 1:
raise ValueError('Target_size None works only with batch_size=1')
self.target_size = (None, None) if target_size is None else tuple(target_size)
# Check color mode
if color_mode not in {'rgb', 'grayscale', 'bgr'}:
raise ValueError('Invalid color mode:', color_mode,
'; expected "rgb" or "grayscale".')
self.color_mode = color_mode
if self.color_mode == 'rgb' or self.color_mode == 'bgr':
self.grayscale = False
if self.dim_ordering == 'tf':
self.image_shape = self.target_size + (3,)
self.gt_image_shape = self.target_size + (1,)
else:
self.image_shape = (3,) + self.target_size
self.gt_image_shape = (1,) + self.target_size
else:
self.grayscale = True
if self.dim_ordering == 'tf':
self.image_shape = self.target_size + (1,)
self.gt_image_shape = self.image_shape
else:
self.image_shape = (1,) + self.target_size
self.gt_image_shape = self.image_shape
# Check class mode
if class_mode not in {'categorical', 'binary', 'sparse',
'segmentation', 'detection', None}:
raise ValueError('Invalid class_mode:', class_mode,
'; expected one of "categorical", '
'"binary", "sparse", "segmentation", "detection" or None.')
self.class_mode = class_mode
self.has_gt_image = True if self.class_mode == 'segmentation' else False
# Check class names
if not classes:
if self.class_mode == 'segmentation' or self.class_mode == 'detection':
raise ValueError('You should input the class names')
else:
classes = list_subdirs(directory)
else:
classes = classes.values()
self.nb_class = len(classes)
self.class_indices = dict(zip(classes, range(len(classes))))
self.nb_sample = 0
self.filenames = []
self.classes = []
############################################################################################
## SDD utility implementation by Team 3 ##
##########################################
if "ssd" in self.model_name:
self.ssd_generator = BBoxUtility(self.nb_class)
else:
self.ssd_generator = None
############################################################################################
# Get filenames
if self.class_mode == 'detection':
for fname in os.listdir(directory):
if has_valid_extension(fname):
self.filenames.append(fname)
# Look for the GT filename
gt_fname = os.path.join(directory,fname.replace('jpg','txt'))
if not os.path.isfile(gt_fname):
raise ValueError('GT file not found: ' + gt_fname)
self.filenames = np.sort(self.filenames)
elif not self.class_mode == 'segmentation':
for subdir in classes:
subpath = os.path.join(directory, subdir)
for fname in os.listdir(subpath):
if has_valid_extension(fname):
self.classes.append(self.class_indices[subdir])
self.filenames.append(os.path.join(subdir, fname))
self.classes = np.array(self.classes)
else:
for fname in os.listdir(directory):
if has_valid_extension(fname):
self.filenames.append(fname)
# Look for the GT filename
gt_fname = os.path.join(gt_directory,
os.path.split(fname)[1])
if not os.path.isfile(gt_fname):
raise ValueError('GT file not found: ' + gt_fname)
self.filenames = np.sort(self.filenames)
self.nb_sample = len(self.filenames)
print(' Found %d images belonging to %d classes' % (self.nb_sample,
self.nb_class))
super(DirectoryIterator, self).__init__(self.nb_sample, batch_size,
shuffle, seed)
def next(self):
# Lock the generation of index only. The rest is not under thread
# lock so it can be done in parallel
with self.lock:
index_array, current_index, current_batch_size = next(self.index_generator)
# Create the batch_x and batch_y
if current_batch_size > 1:
batch_x = np.zeros((current_batch_size,) + self.image_shape)
if self.has_gt_image:
batch_y = np.zeros((current_batch_size,) + self.gt_image_shape)
if self.class_mode == 'detection':
batch_y = []
# Build batch of image data
for i, j in enumerate(index_array):
# Load image
fname = self.filenames[j]
# print(fname)
img = load_img(os.path.join(self.directory, fname),
grayscale=self.grayscale,
resize=self.resize, order=1)
x = img_to_array(img, dim_ordering=self.dim_ordering)
# Load GT image if segmentation
if self.has_gt_image:
# Load GT image
gt_img = load_img(os.path.join(self.gt_directory, fname),
grayscale=True,
resize=self.resize, order=0)
y = img_to_array(gt_img, dim_ordering=self.dim_ordering)
else:
y = None
# Load GT image if detection
if self.class_mode == 'detection':
label_path = os.path.join(self.directory, fname).replace('jpg','txt')
gt = np.loadtxt(label_path)
if len(gt.shape) == 1:
gt = gt[np.newaxis,]
y = gt.copy()
y = y[((y[:,1] > 0.) & (y[:,1] < 1.))]
y = y[((y[:,2] > 0.) & (y[:,2] < 1.))]
y = y[((y[:,3] > 0.) & (y[:,3] < 1.))]
y = y[((y[:,4] > 0.) & (y[:,4] < 1.))]
if (y.shape != gt.shape) or (y.shape[0] == 0):
warnings.warn('DirectoryIterator: found an invalid annotation '
'on GT file '+label_path)
# shuffle gt boxes order
np.random.shuffle(y)
# Standarize image
x = self.image_data_generator.standardize(x, y)
# Data augmentation
x, y = self.image_data_generator.random_transform(x, y)
# Add images to batches
if current_batch_size > 1:
batch_x[i] = x
if self.has_gt_image:
batch_y[i] = y
elif self.class_mode == 'detection':
batch_y.append(y)
else:
batch_x = np.expand_dims(x, axis=0)
if self.has_gt_image:
batch_y = np.expand_dims(y, axis=0)
elif self.class_mode == 'detection':
batch_y = [y]
# optionally save augmented images to disk for debugging purposes
if self.save_to_dir:
for i in range(current_batch_size):
fname = '{prefix}_{index}_{hash}.{format}'.format(prefix=self.save_prefix,
index=current_index + i,
hash=np.random.randint(1e4),
format=self.save_format)
if self.class_mode == 'segmentation':
nclasses = self.classes # TODO: Change
color_map = sns.hls_palette(nclasses+1)
void_label = nclasses
save_img2(batch_x[i], batch_y[i],
os.path.join(self.save_to_dir, fname), color_map,
void_label)
else:
img = array_to_img(batch_x[i], self.dim_ordering,
scale=True)
img.save(os.path.join(self.save_to_dir, fname))
# Build batch of labels
if self.class_mode == 'sparse':
batch_y = self.classes[index_array]
elif self.class_mode == 'binary':
batch_y = self.classes[index_array].astype('float32')
elif self.class_mode == 'categorical':
batch_y = np.zeros((len(batch_x), self.nb_class), dtype='float32')
for i, label in enumerate(self.classes[index_array]):
batch_y[i, label] = 1.
elif self.class_mode == 'detection':
# TODO detection: check model, other networks may expect a different batch_y format and shape
# YOLOLoss expects a particular batch_y format and shape
#########################################################################################################
if "yolo" in self.model_name:
batch_y = yolo_build_gt_batch(batch_y, self.image_shape, self.nb_class)
elif "ssd" in self.model_name:
batch_y = self.ssd_generator.ssd_build_gt_batch(batch_y, self.image_shape)
#########################################################################################################
elif self.class_mode == None:
return batch_x
return batch_x, batch_y
def next(self):
# Lock the generation of index only. The rest is not under thread
# lock so it can be done in parallel
with self.lock:
index_array, current_index, current_batch_size = next(self.index_generator)
# Create the batch_x and batch_y
if current_batch_size > 1:
batch_x = np.zeros((current_batch_size,) + self.image_shape)
if self.has_gt_image:
batch_y = np.zeros((current_batch_size,) + self.gt_image_shape)
if self.class_mode == 'detection':
batch_y = []
# Build batch of image data
for i, j in enumerate(index_array):
# Load image
fname = self.filenames[j]
# print(fname)
img = load_img(os.path.join(self.directory, fname),
grayscale=self.grayscale,
resize=self.resize, order=1)
x = img_to_array(img, dim_ordering=self.dim_ordering)
# Load GT image if segmentation
if self.has_gt_image:
# Load GT image
gt_img = load_img(os.path.join(self.gt_directory, fname),
grayscale=True,
resize=self.resize, order=0)
y = img_to_array(gt_img, dim_ordering=self.dim_ordering)
else:
y = None
# Load GT image if detection
if self.class_mode == 'detection':
label_path = os.path.join(self.directory, fname).replace('jpg','txt')
gt = np.loadtxt(label_path)
if len(gt.shape) == 1:
gt = gt[np.newaxis,]
y = gt.copy()
y = y[((y[:,1] > 0.) & (y[:,1] < 1.))]
y = y[((y[:,2] > 0.) & (y[:,2] < 1.))]
y = y[((y[:,3] > 0.) & (y[:,3] < 1.))]
y = y[((y[:,4] > 0.) & (y[:,4] < 1.))]
if (y.shape != gt.shape) or (y.shape[0] == 0):
warnings.warn('DirectoryIterator: found an invalid annotation '
'on GT file '+label_path)
# shuffle gt boxes order
np.random.shuffle(y)
# Standarize image
x = self.image_data_generator.standardize(x, y)
# Data augmentation
x, y = self.image_data_generator.random_transform(x, y)
# Add images to batches
if current_batch_size > 1:
batch_x[i] = x
if self.has_gt_image:
batch_y[i] = y
elif self.class_mode == 'detection':
batch_y.append(y)
else:
batch_x = np.expand_dims(x, axis=0)
if self.has_gt_image:
batch_y = np.expand_dims(y, axis=0)
elif self.class_mode == 'detection':
batch_y = [y]
# optionally save augmented images to disk for debugging purposes
if self.save_to_dir:
for i in range(current_batch_size):
fname = '{prefix}_{index}_{hash}.{format}'.format(prefix=self.save_prefix,
index=current_index + i,
hash=np.random.randint(1e4),
format=self.save_format)
if self.class_mode == 'segmentation':
nclasses = self.classes # TODO: Change
color_map = sns.hls_palette(nclasses+1)
void_label = nclasses
save_img2(batch_x[i], batch_y[i],
os.path.join(self.save_to_dir, fname), color_map,
void_label)
else:
img = array_to_img(batch_x[i], self.dim_ordering,
scale=True)
img.save(os.path.join(self.save_to_dir, fname))
# Build batch of labels
if self.class_mode == 'sparse':
batch_y = self.classes[index_array]
elif self.class_mode == 'binary':
batch_y = self.classes[index_array].astype('float32')
elif self.class_mode == 'categorical':
batch_y = np.zeros((len(batch_x), self.nb_class), dtype='float32')
for i, label in enumerate(self.classes[index_array]):
batch_y[i, label] = 1.
elif self.class_mode == 'detection':
if self.model_name == 'ssd':
bbox_util = BBoxUtility(self.nb_class)
batch_y = bbox_util.ssd_build_gt_batch(batch_y)
elif self.model_name == 'yolo' or self.model_name == 'tiny-yolo':
# TODO detection: check model, other networks may expect a different batch_y format and shape
# YOLOLoss expects a particular batch_y format and shape
batch_y = yolo_build_gt_batch(batch_y, self.image_shape, self.nb_class)
else:
print('Unknown detection model')
elif self.class_mode == None:
return batch_x
return batch_x, batch_y