if __name__ == "__main__":
args = io_utils.handle_args()
if args.smoke_test:
ray.init(num_cpus=2)
else:
ray.init(address=args.address)
if args.backbone == "mobilenet_v2":
from models.ssd_mobilenet_v2 import get_model, init_model
else:
from models.ssd_vgg16 import get_model, init_model
ssd_log_path = io_utils.get_log_path(args.backbone)
ssd_model_path = io_utils.get_model_path(args.backbone)
hyper_params = train_utils.get_hyper_params(args.backbone)
_, info = data_utils.get_dataset("voc/2007", "train+validation")
_, voc_2012_info = data_utils.get_dataset("voc/2012", "train+validation")
voc_2012_total_items = data_utils.get_total_item_size(
voc_2012_info, "train+validation")
train_total_items = data_utils.get_total_item_size(info,
"train+validation")
val_total_items = data_utils.get_total_item_size(info, "test")
if args.with_voc12:
train_total_items += voc_2012_total_items
labels = data_utils.get_labels(info)
labels = ["bg"] + labels
hyper_params["total_labels"] = len(labels)
img_paths = data_utils.get_custom_imgs(custom_image_path)
total_items = len(img_paths)
test_data = tf.data.Dataset.from_generator(
lambda: data_utils.custom_data_generator(img_paths, img_size, img_size
), data_types, data_shapes)
else:
test_data = test_data.map(
lambda x: data_utils.preprocessing(x, img_size, img_size))
test_data = test_data.padded_batch(batch_size,
padded_shapes=data_shapes,
padding_values=padding_values)
rpn_model, _ = get_model(hyper_params)
frcnn_model_path = io_utils.get_model_path("faster_rcnn", backbone)
rpn_model_path = io_utils.get_model_path("rpn", backbone)
model_path = frcnn_model_path if load_weights_from_frcnn else rpn_model_path
rpn_model.load_weights(model_path, by_name=True)
anchors = bbox_utils.generate_anchors(hyper_params)
for image_data in test_data:
imgs, _, _ = image_data
rpn_bbox_deltas, rpn_labels = rpn_model.predict_on_batch(imgs)
#
rpn_bbox_deltas = tf.reshape(rpn_bbox_deltas, (batch_size, -1, 4))
rpn_labels = tf.reshape(rpn_labels, (batch_size, -1))
#
rpn_bbox_deltas *= hyper_params["variances"]
rpn_bboxes = bbox_utils.get_bboxes_from_deltas(anchors, rpn_bbox_deltas)
train_data = train_data.padded_batch(batch_size,
padded_shapes=data_shapes,
padding_values=padding_values)
val_data = val_data.padded_batch(batch_size,
padded_shapes=data_shapes,
padding_values=padding_values)
anchors = bbox_utils.generate_anchors(hyper_params)
rpn_train_feed = train_utils.rpn_generator(train_data, anchors, hyper_params)
rpn_val_feed = train_utils.rpn_generator(val_data, anchors, hyper_params)
rpn_model, _ = get_model(hyper_params)
rpn_model.compile(optimizer=tf.optimizers.Adam(learning_rate=1e-5),
loss=[train_utils.reg_loss, train_utils.rpn_cls_loss])
# Load weights
rpn_model_path = io_utils.get_model_path("rpn", backbone)
if load_weights:
rpn_model.load_weights(rpn_model_path)
checkpoint_callback = ModelCheckpoint(rpn_model_path,
monitor="val_loss",
save_best_only=True,
save_weights_only=True)
step_size_train = train_utils.get_step_size(train_total_items, batch_size)
step_size_val = train_utils.get_step_size(val_total_items, batch_size)
rpn_model.fit(rpn_train_feed,
steps_per_epoch=step_size_train,
validation_data=rpn_val_feed,
validation_steps=step_size_val,
padding_values = data_utils.get_padding_values()
train_data = train_data.shuffle(batch_size * 4).padded_batch(
batch_size, padded_shapes=data_shapes, padding_values=padding_values)
val_data = val_data.padded_batch(batch_size,
padded_shapes=data_shapes,
padding_values=padding_values)
# Setup training model (ssd+vgg) and loss function (location + confidence)
ssd_model = get_model(hyper_params)
ssd_custom_losses = CustomLoss(hyper_params["neg_pos_ratio"],
hyper_params["loc_loss_alpha"])
ssd_model.compile(
optimizer=Adam(learning_rate=1e-3),
loss=[ssd_custom_losses.loc_loss_fn, ssd_custom_losses.conf_loss_fn])
init_model(ssd_model)
ssd_model_path = io_utils.get_model_path(backbone)
if load_weights:
ssd_model.load_weights(ssd_model_path)
ssd_log_path = io_utils.get_log_path(backbone)
# We calculate prior boxes for one time and use it for all operations because of the all images are the same sizes
prior_boxes = bbox_utils.generate_prior_boxes(
hyper_params["feature_map_shapes"], hyper_params["aspect_ratios"])
ssd_train_feed = train_utils.generator(train_data, prior_boxes, hyper_params)
ssd_val_feed = train_utils.generator(val_data, prior_boxes, hyper_params)
checkpoint_callback = ModelCheckpoint(ssd_model_path,
monitor="val_loss",
save_best_only=True,
save_weights_only=True)
tensorboard_callback = TensorBoard(log_dir=ssd_log_path)
learning_rate_callback = LearningRateScheduler(train_utils.scheduler,
test_data = tf.data.Dataset.from_generator(
lambda: data_utils.custom_data_generator(img_paths, img_size, img_size
), data_types, data_shapes)
else:
test_split = "train[80%:]"
test_data, info = data_utils.get_dataset("the300w_lp", test_split)
total_items = data_utils.get_total_item_size(info, test_split)
test_data = test_data.map(
lambda x: data_utils.preprocessing(x, img_size, img_size))
#
test_data = test_data.padded_batch(batch_size,
padded_shapes=data_shapes,
padding_values=padding_values)
model = blazeface.get_model(hyper_params)
model_path = io_utils.get_model_path()
model.load_weights(model_path)
prior_boxes = bbox_utils.generate_prior_boxes(
hyper_params["feature_map_shapes"], hyper_params["aspect_ratios"])
variances = hyper_params["variances"]
total_landmarks = hyper_params["total_landmarks"]
landmark_variances = total_landmarks * variances[0:2]
variances += landmark_variances
for image_data in test_data:
img, _, _ = image_data
pred_deltas, pred_scores = model.predict_on_batch(img)
pred_deltas *= variances
#
