def dataset(self, config, debug=False): # #### load dataset train_dataset = coco.CocoDataSet('./COCO2017', 'train', flip_ratio=0., pad_mode='fixed', config=config, debug=False) print(train_dataset.get_categories()) assert config.NUM_CLASSES == len( train_dataset.get_categories() ), f"NUM_CLASSES must be compare with dataset, set:{config.NUM_CLASSES} != {len(train_dataset.get_categories())}" train_generator = data_generator.DataGenerator(train_dataset) train_tf_dataset = tf.data.Dataset.from_generator( train_generator, (tf.float32, tf.float32, tf.float32, tf.int32, tf.float32, tf.float32)) self.train_tf_dataset = train_tf_dataset.padded_batch( config.IMAGES_PER_GPU, padded_shapes=( [None, None, None], # img [None], #img_meta [None, None], #bboxes [None], #labels [None, None, None], #masks [None, None, 1])) # global_mask eval_dataset = coco.CocoDataSet('./COCO2017', 'val', flip_ratio=0., pad_mode='fixed', config=config, debug=False) eval_generator = data_generator.DataGenerator(eval_dataset) eval_tf_dataset = tf.data.Dataset.from_generator( eval_generator, (tf.float32, tf.float32, tf.float32, tf.int32, tf.float32, tf.float32)) self.eval_tf_dataset = eval_tf_dataset.padded_batch( config.IMAGES_PER_GPU, padded_shapes=( [None, None, None], # img [None], #img_meta [None, None], #bboxes [None], #labels [None, None, None], #masks [None, None, 1])) # global_mask if debug: idx = np.random.choice(range(len(train_dataset))) img, img_meta, bboxes, labels, masks, global_mask = train_dataset[ idx] rgb_img = np.round(img + config.MEAN_PIXEL) ori_img = utils.get_original_image(img, img_meta, config.MEAN_PIXEL) visualize.display_instances(rgb_img, bboxes, labels, train_dataset.get_categories()) self.train_dataset = train_dataset
# load dataset from detection.datasets import coco, data_generator img_mean = (123.675, 116.28, 103.53) # img_std = (58.395, 57.12, 57.375) img_std = (1., 1., 1.) train_dataset = coco.CocoDataSet('./COCO2017/', 'val', flip_ratio=0.5, pad_mode='fixed', mean=img_mean, std=img_std, scale=(640, 832), debug=True) train_generator = data_generator.DataGenerator(train_dataset) # load model from detection.models.detectors import faster_rcnn strategy = tf.distribute.MirroredStrategy() with strategy.scope(): model = faster_rcnn.FasterRCNN( num_classes=len(train_dataset.get_categories())) img, img_meta, _, _ = train_dataset[0] batch_imgs = tf.Variable(np.expand_dims(img, 0)) batch_metas = tf.Variable(np.expand_dims(img_meta, 0)) _ = model((batch_imgs, batch_metas), training=False)
train_dataset = coco.CocoDataSet('dataset', 'train', flip_ratio=0, pad_mode='fixed', mean=img_mean, std=img_std, scale=(800, 1216)) val_dataset = coco.CocoDataSet('dataset', 'train', flip_ratio=0, pad_mode='non-fixed', mean=img_mean, std=img_std, scale=(800, 1216)) train_generator = data_generator.DataGenerator(train_dataset, shuffle=False) train_tf_dataset = tf.data.Dataset.from_generator( train_generator, (tf.float32, tf.float32, tf.float32, tf.int32)) train_tf_dataset = train_tf_dataset.batch(BATCH_SIZE).prefetch(100) # 模型设置 num_classes = len(train_dataset.get_categories()) model = faster_rcnn.FasterRCNN(num_classes=num_classes) optimizer = keras.optimizers.SGD(LEARNING_RATE, momentum=0.9, nesterov=True) if CHECKPOINT != '': batch_imgs, batch_metas, batch_bboxes, batch_labels = train_dataset[0] with tf.GradientTape() as tape: _, _, _, _ = model( (np.array([batch_imgs]), np.array([batch_metas]), np.array([batch_bboxes]), np.array([batch_labels])), training=True)