def data_generator(dataset: ClassificationDataset,
config,
batch_size,
shuffle=True):
preprocess_input = Classifiers.get_preprocessing(config.BACKBONE)
b = 0 # batch item index
image_index = -1
image_ids = np.arange(dataset.num_images)
error_count = 0
lock = threading.Lock()
# Keras requires a generator to run indefinitely.
while True:
try:
with lock:
# Increment index to pick next image. Shuffle if at the start of an epoch.
image_index = (image_index + 1) % len(image_ids)
if shuffle and image_index == 0:
np.random.shuffle(image_ids)
# Get GT bounding boxes and masks for image.
image_id = image_ids[image_index]
image = dataset.load_image(image_id)
label = dataset.load_label(image_id)
# Init batch arrays
if b == 0:
batch_images = np.zeros((batch_size, ) + image.shape,
dtype=np.float32)
batch_labels = np.zeros((batch_size, dataset.num_classes),
dtype=np.float32)
# Add to batch
batch_images[b] = preprocess_input(image.astype(np.float32))
batch_labels[b] = label.astype(np.float32)
b += 1
# Batch full?
if b >= batch_size:
inputs = [batch_images, batch_labels]
outputs = []
yield inputs, outputs
# start a new batch
b = 0
except (GeneratorExit, KeyboardInterrupt):
raise
except:
# Log it and skip the image
logging.exception("Error processing image {}".format(image_id))
error_count += 1
if error_count > 5:
raise
def get_preprocessing(name):
return Classifiers.get_preprocessing(name)
def mold_image(images, backbone_name):
if 'vgg' in backbone_name:
return images.astype(np.float32) / 127.0 - 1.0
else:
preprocess_input = Classifiers.get_preprocessing(backbone_name)
return preprocess_input(images)
def data_generator(dataset, config, shuffle=True, batch_size=1):
"""A generator that returns images and corresponding target class ids,
bounding box deltas, and masks.
dataset: The Dataset object to pick data from
config: The model config object
shuffle: If True, shuffles the samples before every epoch
batch_size: How many images to return in each call
Returns a Python generator. Upon calling next() on it, the
generator returns two lists, inputs and outputs. The contents
of the lists differs depending on the received arguments:
inputs list:
- images: [batch, H, W, C]
- gt_mask: [batch, height, width, NUM_CLASSES]. The height and width
are those of the image and NUM_CLASSES doesn't include background
outputs list: empty
"""
b = 0 # batch item index
image_index = -1
image_ids = np.copy(dataset.image_ids)
error_count = 0
preprocess_input = Classifiers.get_preprocessing(config.BACKBONE)
lock = threading.Lock()
# Keras requires a generator to run indefinitely.
while True:
try:
with lock:
# Increment index to pick next image. Shuffle if at the start of an epoch.
image_index = (image_index + 1) % len(image_ids)
if shuffle and image_index == 0:
np.random.shuffle(image_ids)
# Get GT bounding boxes and masks for image.
image_id = image_ids[image_index]
image, gt_mask = load_image_gt(dataset, image_id,
config.IMAGE_SHAPE,
config.USE_MINI_MASK,
config.MINI_MASK_SHAPE)
# Init batch arrays
if b == 0:
batch_images = np.zeros((batch_size, ) + image.shape,
dtype=np.float32)
batch_gt_mask = np.zeros(
(batch_size, gt_mask.shape[0], gt_mask.shape[1],
config.NUM_CLASSES - 1),
dtype=gt_mask.dtype)
# Add to batch
batch_images[b] = preprocess_input(image.astype(np.float32))
batch_gt_mask[b, :, :, :] = gt_mask
b += 1
# Batch full?
if b >= batch_size:
inputs = [batch_images, batch_gt_mask]
outputs = []
yield inputs, outputs
# start a new batch
b = 0
except (GeneratorExit, KeyboardInterrupt):
raise
except:
# Log it and skip the image
logging.exception("Error processing image {}".format(image_id))
error_count += 1
if error_count > 5:
raise
