def yolo3_data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes, enhance_augment, rescale_interval, multi_anchor_assign):
'''data generator for fit_generator'''
n = len(annotation_lines)
i = 0
# prepare multiscale config
rescale_step = 0
input_shape_list = get_multiscale_list()
while True:
if rescale_interval > 0:
# Do multi-scale training on different input shape
rescale_step = (rescale_step + 1) % rescale_interval
if rescale_step == 0:
input_shape = input_shape_list[random.randint(0, len(input_shape_list)-1)]
image_data = []
box_data = []
for b in range(batch_size):
if i==0:
np.random.shuffle(annotation_lines)
image, box = get_ground_truth_data(annotation_lines[i], input_shape, augment=True)
image_data.append(image)
box_data.append(box)
i = (i+1) % n
image_data = np.array(image_data)
box_data = np.array(box_data)
if enhance_augment == 'mosaic':
# add random mosaic augment on batch ground truth data
image_data, box_data = random_mosaic_augment(image_data, box_data, prob=0.2)
y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes, multi_anchor_assign)
yield [image_data, *y_true], np.zeros(batch_size)
def __getitem__(self, index):
# generate annotation indexes for every batch
batch_indexs = self.indexes[index*self.batch_size : (index+1)*self.batch_size]
# fetch annotation lines based on index
batch_annotation_lines = [self.annotation_lines[i] for i in batch_indexs]
if self.rescale_interval > 0:
# Do multi-scale training on different input shape
self.rescale_step = (self.rescale_step + 1) % self.rescale_interval
if self.rescale_step == 0:
self.input_shape = self.input_shape_list[random.randint(0, len(self.input_shape_list)-1)]
image_data = []
box_data = []
for b in range(self.batch_size):
image, box = get_ground_truth_data(batch_annotation_lines[b], self.input_shape, augment=True)
image_data.append(image)
box_data.append(box)
image_data = np.array(image_data)
box_data = np.array(box_data)
if self.enhance_augment == 'mosaic':
# add random mosaic augment on batch ground truth data
image_data, box_data = random_mosaic_augment(image_data, box_data, prob=0.2)
y_true = preprocess_true_boxes(box_data, self.input_shape, self.anchors, self.num_classes, self.multi_anchor_assign)
return [image_data, *y_true], np.zeros(self.batch_size)
def main():
parser = argparse.ArgumentParser(
argument_default=argparse.SUPPRESS,
description='Test tool for mosaic data augment function')
parser.add_argument('--annotation_file',
type=str,
required=True,
help='data annotation txt file')
parser.add_argument('--classes_path',
type=str,
required=True,
help='path to class definitions')
parser.add_argument(
'--output_path',
type=str,
required=False,
help='output path for augmented images, default is ./test',
default='./test')
parser.add_argument('--batch_size',
type=int,
required=False,
help="batch size for test data, default=16",
default=16)
parser.add_argument(
'--model_image_size',
type=str,
required=False,
help='model image input size as <num>x<num>, default 416x416',
default='416x416')
args = parser.parse_args()
class_names = get_classes(args.classes_path)
height, width = args.model_image_size.split('x')
model_image_size = (int(height), int(width))
annotation_lines = get_dataset(args.annotation_file)
os.makedirs(args.output_path, exist_ok=True)
image_data = []
boxes_data = []
for i in range(args.batch_size):
annotation_line = annotation_lines[i]
image, boxes = get_ground_truth_data(annotation_line,
input_shape=model_image_size,
augment=True)
#un-normalize image
image = image * 255.0
image = image.astype(np.uint8)
image_data.append(image)
boxes_data.append(boxes)
image_data = np.array(image_data)
boxes_data = np.array(boxes_data)
image_data, boxes_data = random_mosaic_augment(image_data,
boxes_data,
jitter=1)
draw_boxes(image_data, boxes_data, class_names, args.output_path)
def main():
parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS, description='Test tool for enhance mosaic data augment function')
parser.add_argument('--annotation_file', type=str, required=True, help='data annotation txt file')
parser.add_argument('--classes_path', type=str, required=True, help='path to class definitions')
parser.add_argument('--output_path', type=str, required=False, help='output path for augmented images, default is ./test', default='./test')
parser.add_argument('--batch_size', type=int, required=False, help = "batch size for test data, default=16", default=16)
parser.add_argument('--model_image_size', type=str, required=False, help='model image input size as <height>x<width>, default 416x416', default='416x416')
parser.add_argument('--augment_type', type=str, required=False, help = "enhance data augmentation type (mosaic/cutmix), default=mosaic", default='mosaic')
args = parser.parse_args()
class_names = get_classes(args.classes_path)
height, width = args.model_image_size.split('x')
model_image_size = (int(height), int(width))
assert (model_image_size[0]%32 == 0 and model_image_size[1]%32 == 0), 'model_image_size should be multiples of 32'
annotation_lines = get_dataset(args.annotation_file)
os.makedirs(args.output_path, exist_ok=True)
image_data = []
boxes_data = []
for i in range(args.batch_size):
annotation_line = annotation_lines[i]
image, boxes = get_ground_truth_data(annotation_line, input_shape=model_image_size, augment=True)
#un-normalize image
image = image*255.0
image = image.astype(np.uint8)
image_data.append(image)
boxes_data.append(boxes)
image_data = np.array(image_data)
boxes_data = np.array(boxes_data)
if args.augment_type == 'mosaic':
image_data, boxes_data = random_mosaic_augment(image_data, boxes_data, prob=1)
elif args.augment_type == 'cutmix':
image_data, boxes_data = random_cutmix_augment(image_data, boxes_data, prob=1)
else:
raise ValueError('Unsupported augment type')
draw_boxes(image_data, boxes_data, class_names, args.output_path)
def main():
parser = argparse.ArgumentParser(
argument_default=argparse.SUPPRESS,
description='Test tool for enhance mosaic data augment function')
parser.add_argument('--annotation_file',
type=str,
required=True,
help='data annotation txt file')
parser.add_argument('--classes_path',
type=str,
required=True,
help='path to class definitions')
parser.add_argument(
'--output_path',
type=str,
required=False,
help='output path for augmented images, default=%(default)s',
default='./test')
parser.add_argument('--batch_size',
type=int,
required=False,
help="batch size for test data, default=%(default)s",
default=16)
parser.add_argument(
'--model_input_shape',
type=str,
required=False,
help='model image input shape as <height>x<width>, default=%(default)s',
default='416x416')
parser.add_argument(
'--enhance_augment',
type=str,
required=False,
help="enhance data augmentation type, default=%(default)s",
default=None,
choices=['mosaic', 'mosaic_v5', 'cutmix', None])
args = parser.parse_args()
class_names = get_classes(args.classes_path)
height, width = args.model_input_shape.split('x')
model_input_shape = (int(height), int(width))
assert (model_input_shape[0] % 32 == 0 and model_input_shape[1] % 32
== 0), 'model_input_shape should be multiples of 32'
annotation_lines = get_dataset(args.annotation_file)
os.makedirs(args.output_path, exist_ok=True)
image_data = []
boxes_data = []
pbar = tqdm(total=args.batch_size, desc='Generate augment image')
for i in range(args.batch_size):
pbar.update(1)
annotation_line = annotation_lines[i]
image, boxes = get_ground_truth_data(annotation_line,
input_shape=model_input_shape,
augment=True)
#denormalize image
image = denormalize_image(image)
image_data.append(image)
boxes_data.append(boxes)
pbar.close()
image_data = np.array(image_data)
boxes_data = np.array(boxes_data)
if args.enhance_augment == 'mosaic':
image_data, boxes_data = random_mosaic_augment(image_data,
boxes_data,
prob=1)
elif args.enhance_augment == 'mosaic_v5':
image_data, boxes_data = random_mosaic_augment_v5(image_data,
boxes_data,
prob=1)
elif args.enhance_augment == 'cutmix':
image_data, boxes_data = random_cutmix_augment(image_data,
boxes_data,
prob=1)
elif args.enhance_augment == None:
print('No enhance augment type. Will only apply base augment')
else:
raise ValueError('Unsupported augment type')
draw_boxes(image_data, boxes_data, class_names, args.output_path)
print('Done. augment images have been saved in', args.output_path)
