def yolo5_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) #elif enhance_augment == 'mosaic_v5': # mosaic augment from YOLOv5 #image_data, box_data = random_mosaic_augment_v5(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) #elif self.enhance_augment == 'mosaic_v5': # mosaic augment from YOLOv5 #image_data, box_data = random_mosaic_augment_v5(image_data, box_data, prob=0.2) y_true_data = get_y_true_data(box_data, self.anchors, self.input_shape, self.num_classes, self.multi_anchor_assign) return [image_data, y_true_data], np.zeros(self.batch_size)
def data_generator(): n = len(annotation_lines) i = 0 for num in range(sample_num): image, _ = get_ground_truth_data(annotation_lines[i], model_input_shape, augment=True) i = (i+1) % n image = np.array([image], dtype=np.float32) yield [image]
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 data_generator(): n = len(annotation_lines) i = 0 for num in range(sample_num): image, _ = get_ground_truth_data(annotation_lines[i], model_input_shape, augment=True) i = (i + 1) % n image = np.array([image], dtype=np.float32) if not args.channel: image = image[..., -1] image = np.expand_dims(image, axis=-1) yield [image]
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)