Example #1
0
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)
Example #3
0
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,
Example #4
0
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,
Example #5
0
    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
    #