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)
