def main():
# dataset
dataset = TrashData()
dataset.load('data/dataset_surfrider_cleaned')
# model
model = faster_rcnn.FasterRCNN(num_classes=dataset.class_count)
optimizer = keras.optimizers.SGD(1e-3, momentum=0.9, nesterov=True)
for epoch in range(100):
loss_history = []
for (batch, inputs) in enumerate(train_tf_dataset):
batch_imgs, batch_metas, batch_bboxes, batch_labels = inputs
with tf.GradientTape() as tape:
rpn_class_loss, rpn_bbox_loss, rcnn_class_loss, rcnn_bbox_loss = model(
(batch_imgs, batch_metas, batch_bboxes, batch_labels),
training=True)
loss_value = rpn_class_loss + rpn_bbox_loss + rcnn_class_loss + rcnn_bbox_loss
grads = tape.gradient(loss_value, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
loss_history.append(loss_value.numpy())
if batch % 10 == 0:
print('epoch', epoch, batch, np.mean(loss_history))
train_datasets = ZiptrainDataset(
tf_record_path,
1,
96,
crop_size=crop_size,
roi_path='D:/datasets/bjod/roi_test/').prepare(True)
val_train = Zipvaluedata(
tf_record_path, crop_size=crop_size).prepare('train_api_97.record')
val_test = Zipvaluedata(
tf_record_path, crop_size=crop_size).prepare('val_api_19.record')
one_imgs, one_metas, one_bboxes, one_labels, _ = next(iter(val_train))
one_imgs = tf.expand_dims(tf.cast(one_imgs[0], tf.float32), axis=0)
one_metas = tf.expand_dims(tf.cast(one_metas[0], tf.float32), axis=0)
model = faster_rcnn.FasterRCNN(num_classes=2)
_ = model((one_imgs, one_metas), training=False)
model_ori = faster_rcnn.FasterRCNN(num_classes=81)
_ = model_ori((one_imgs, one_metas), training=False)
model_ori.load_weights(
'./weights/faster_rcnn_resnet101_fpn_coco2017_map35.h5',
by_name=True)
model.backbone.set_weights(model_ori.backbone.get_weights())
model.neck.set_weights(model_ori.neck.get_weights())
model.rpn_head.set_weights(model_ori.rpn_head.get_weights())
model.roi_align.set_weights(model_ori.roi_align.get_weights())
# print(cc)
model.summary()
def __init__(self,
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)
model.load_weights('weights/faster_rcnn.h5', by_name=True)
# use tf.data
global_batch_size = 2
with strategy.scope():
train_tf_dataset = tf.data.Dataset.from_generator(
train_generator, (tf.float32, tf.float32, tf.float32, tf.int32))
train_tf_dataset = train_tf_dataset.padded_batch(
flip_ratio=0.5,
pad_mode='fixed',
mean=img_mean,
std=img_std,
scale=(800, 1216))
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))
train_tf_dataset = train_tf_dataset.batch(batch_size).prefetch(100).shuffle(
100)
# train_tf_dataset = train_tf_dataset.padded_batch(
# batch_size, padded_shapes=([None, None, None], [None], [None, None], [None]))
model = faster_rcnn.FasterRCNN(num_classes=num_classes)
optimizer = keras.optimizers.SGD(1e-3, momentum=0.9, nesterov=True)
img, img_meta, bboxes, labels = train_dataset[6] # [N, 4], shape:[N]=data:[62]
rgb_img = np.round(img + img_mean)
ori_img = get_original_image(img, img_meta, img_mean)
# visualize.display_instances(rgb_img, bboxes, labels, train_dataset.get_categories())
batch_imgs = tf.convert_to_tensor(np.expand_dims(img, 0)) # [1, 1216, 1216, 3]
batch_metas = tf.convert_to_tensor(np.expand_dims(img_meta, 0)) # [1, 11]
# dummpy forward to build network variables
_ = model((batch_imgs, batch_metas), training=False)
proposals = model.simple_test_rpn(img, img_meta)
res = model.simple_test_bboxes(img, img_meta, proposals)