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 'yolo' in self.model_name:
batch_y = yolo_build_gt_batch(batch_y, self.image_shape, self.nb_class)
elif self.model_name == 'ssd':
batch_y = self.bbox_util.ssd_build_gt_batch_v2(batch_y)
# targets = []
# for boxes in batch_y:
# boxes_corrected = np.zeros((boxes.shape[0], 4+self.nb_class))
# for b,box in enumerate(boxes):
# boxes_corrected[b,0] = box[1] - box[3]/2
# boxes_corrected[b,1] = box[2] - box[4]/2
# boxes_corrected[b,2] = box[1] + box[3]/2
# boxes_corrected[b,3] = box[2] + box[4]/2
# c = 4+int(box[0])
# boxes_corrected[b,c] = 1.
# boxes_corrected = self.bbox_util.assign_boxes(boxes_corrected)
# targets.append(boxes_corrected)
# batch_y = np.array(targets)
elif self.class_mode == None:
return batch_x
return batch_x, batch_y
def next(self):
with self.lock:
index_array, current_index, current_batch_size = next(
self.index_generator)
# The transformation of images is not under thread lock so it can
# be done in parallel
if self.target_size[0] is None:
batch_x = 0
else:
batch_x = np.zeros((current_batch_size, ) + self.image_shape)
if self.class_mode == 'seg_map':
if self.target_size[0] is None:
batch_y = 0
else:
batch_y = np.zeros((current_batch_size, 1, self.image_shape[1],
self.image_shape[2]))
grayscale = self.color_mode == 'grayscale'
# build batch of image data
for i, j in enumerate(index_array):
fname = self.filenames[j]
gtname = self.gt_filenames[j]
img = load_img(os.path.join(self.directory, fname),
grayscale=grayscale,
target_size=self.target_size)
x = img_to_array(img, dim_ordering=self.dim_ordering)
if not self.class_mode == 'seg_map':
x = self.image_data_generator.standardize(x)
x = self.image_data_generator.random_transform(x)
if current_batch_size > 1:
batch_x[i] = x
else:
batch_x = np.expand_dims(x, axis=0)
else:
GT = load_img(os.path.join(self.gt_directory, gtname),
grayscale=False,
target_size=self.target_size)
y = img_to_array(GT, dim_ordering=self.dim_ordering)
x = self.image_data_generator.standardize(x, y)
x, y = self.image_data_generator.random_transform(x, y)
if current_batch_size > 1:
batch_x[i] = x
batch_y[i] = y
else:
batch_x = np.expand_dims(x, axis=0)
batch_y = np.expand_dims(y, axis=0)
# 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 == 'seg_map':
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 not self.class_mode == 'seg_map':
return batch_x
# print(self.directory, np.shape(batch_x), np.shape(batch_y))
return batch_x, batch_y