def predict(self, path=None, config=None):
# restore model
model = Model(config)
ckpt = tf.train.Checkpoint(model=model)
ckpt_manager = tf.train.CheckpointManager(ckpt,
directory=config.MODEL_DIR,
max_to_keep=5)
if ckpt_manager.latest_checkpoint:
ckpt.restore(ckpt_manager.latest_checkpoint)
print(
f'Latest checkpoint restored!!\n\tModel path is {ckpt_manager.latest_checkpoint}'
)
idx = np.random.choice(range(len(self.train_dataset)))
print(f"Predicting image id: {idx}")
img, img_meta, bboxes, labels, masks, global_mask = self.train_dataset[
idx]
res = self.predict_step([img, img_meta], model)
print("res", res)
ori_img = utils.get_original_image(img, img_meta, config.MEAN_PIXEL)
# visualize.display_instances(ori_img, res['rois'], res['class_ids'],
# self.train_dataset.get_categories(), res['masks'], scores=res['scores'])
imshow = show.visualize_boxes(
image=ori_img.astype(np.uint8),
boxes=res['rois'].astype(np.int32),
labels=res['class_ids'],
scores=res['scores'],
class_labels=self.train_dataset.get_categories(),
#masks=res['masks'].transpose([2,0,1])
)
plt.title(f"{idx}")
plt.imshow(imshow)
plt.show()
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
with strategy.scope():
def distributed_train_step(inputs):
per_replica_losses = strategy.experimental_run_v2(train_step,
args=(inputs, ))
return strategy.reduce(tf.distribute.ReduceOp.SUM,
per_replica_losses,
axis=None)
optimizer = tf.keras.optimizers.SGD(1e-3, momentum=0.9, nesterov=True)
epochs = 12
for epoch in range(epochs):
loss_history = []
for (batch, inputs) in enumerate(train_tf_dataset):
batch_loss = distributed_train_step(inputs)
loss_history.append(batch_loss.numpy())
print('epoch:', epoch, ' batch:', batch, ' loss:',
np.mean(loss_history))
with strategy.scope():
from detection.datasets.utils import get_original_image
ori_img = get_original_image(img, img_meta, img_mean)
proposals = model.simple_test_rpn(img, img_meta)
res = model.simple_test_bboxes(img, img_meta, proposals)
visualize.display_instances(ori_img,
res['rois'],
res['class_ids'],
train_dataset.get_categories(),
scores=res['scores'])
Note that the first item 'bg' means background.
'''
return ['bg'] + [
self.coco.loadCats(i)[0]["name"] for i in self.cat2label.keys()
]
if __name__ == "__main__":
config = Config()
img_mean = (123.675, 116.28, 103.53)
# img_std = (58.395, 57.12, 57.375)
img_std = (1., 1., 1.)
train_dataset = CocoDataSet('../../COCO2017/',
'train',
flip_ratio=0.5,
pad_mode='fixed',
config=config,
debug=True)
# for i in range(1000):
img, img_meta, bboxes, labels, masks, global_mask = train_dataset[1]
rgb_img = np.round(img + img_mean)
ori_img = utils.get_original_image(img, img_meta, img_mean)
imshow = show.visualize_boxes(image=rgb_img.astype(np.uint8),
boxes=bboxes,
labels=labels,
scores=np.ones_like(labels),
class_labels=train_dataset.get_categories(),
masks=masks.transpose([2, 0, 1]))
plt.imshow(imshow)
plt.show()
# visualize.display_instances(rgb_img, bboxes, labels, train_dataset.get_categories(), masks=masks)