def build_model_and_load_weights(phi, num_classes, score_threshold,
path_to_weights):
"""
Builds an EfficientPose model and init it with a given weight file
Args:
phi: EfficientPose scaling hyperparameter
num_classes: The number of classes
score_threshold: Minimum score threshold at which a prediction is not filtered out
path_to_weights: Path to the weight file
Returns:
efficientpose_prediction: The EfficientPose model
image_size: Integer image size used as the EfficientPose input resolution for the given phi
"""
print("\nBuilding model...\n")
_, efficientpose_prediction, _ = build_EfficientPose(
phi,
num_classes=num_classes,
num_anchors=9,
freeze_bn=True,
score_threshold=score_threshold,
num_rotation_parameters=3,
print_architecture=False)
print("\nDone!\n\nLoading weights...")
efficientpose_prediction.load_weights(path_to_weights, by_name=True)
print("Done!")
image_sizes = (512, 640, 768, 896, 1024, 1280, 1408)
image_size = image_sizes[phi]
return efficientpose_prediction, image_size
def main(args=None):
"""
Evaluate an EfficientPose model.
Args:
args: parseargs object containing configuration for the evaluation procedure.
"""
allow_gpu_growth_memory()
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
if args.validation_image_save_path:
os.makedirs(args.validation_image_save_path, exist_ok=True)
# create the generators
print("\nCreating the Generators...")
generator = create_generators(args)
print("Done!")
num_rotation_parameters = generator.get_num_rotation_parameters()
num_classes = generator.num_classes()
num_anchors = generator.num_anchors
# optionally choose specific GPU
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
print("\nBuilding the Model...")
_, prediction_model, _ = build_EfficientPose(
args.phi,
num_classes=num_classes,
num_anchors=num_anchors,
freeze_bn=True,
score_threshold=args.score_threshold,
num_rotation_parameters=num_rotation_parameters,
print_architecture=False)
print("Done!")
# load pretrained weights
print('Loading model, this may take a second...')
prediction_model.load_weights(args.weights, by_name=True)
print("\nDone!")
evaluate_model(prediction_model, generator,
args.validation_image_save_path, args.score_threshold)
def build_model(phi, model_path, generator):
"""
Builds an EfficientPose model and init it with a given weight file
Args:
phi: EfficientPose scaling hyperparameter
model_path: Path to the weight file
generator: Dataset generator
Returns:
model: EfficientPose model
"""
_, model, _ = build_EfficientPose(
phi,
num_classes=generator.num_classes(),
num_anchors=generator.num_anchors,
freeze_bn=True,
score_threshold=0.5,
num_rotation_parameters=generator.get_num_rotation_parameters(),
print_architecture=False)
model.load_weights(model_path, by_name=True)
return model
def main(args=None):
"""
Train an EfficientPose model.
Args:
args: parseargs object containing configuration for the training procedure.
"""
allow_gpu_growth_memory()
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
# create the generators
print("\nCreating the Generators...")
train_generator, validation_generator = create_generators(args)
print("Done!")
num_rotation_parameters = train_generator.get_num_rotation_parameters()
num_classes = train_generator.num_classes()
num_anchors = train_generator.num_anchors
# optionally choose specific GPU
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
print("\nBuilding the Model...")
model, prediction_model, all_layers = build_EfficientPose(
args.phi,
num_classes=num_classes,
num_anchors=num_anchors,
freeze_bn=not args.no_freeze_bn,
score_threshold=args.score_threshold,
num_rotation_parameters=num_rotation_parameters)
print("Done!")
# load pretrained weights
if args.weights:
if args.weights == 'imagenet':
model_name = 'efficientnet-b{}'.format(args.phi)
file_name = '{}_weights_tf_dim_ordering_tf_kernels_autoaugment_notop.h5'.format(
model_name)
file_hash = WEIGHTS_HASHES[model_name][1]
weights_path = keras.utils.get_file(file_name,
BASE_WEIGHTS_PATH + file_name,
cache_subdir='models',
file_hash=file_hash)
model.load_weights(weights_path, by_name=True)
else:
print('Loading model, this may take a second...')
custom_load_weights(filepath=args.weights,
layers=all_layers,
skip_mismatch=True)
print("\nDone!")
# freeze backbone layers
if args.freeze_backbone:
# 227, 329, 329, 374, 464, 566, 656
for i in range(1, [227, 329, 329, 374, 464, 566, 656][args.phi]):
model.layers[i].trainable = False
# compile model
model.compile(
optimizer=Adam(lr=args.lr, clipnorm=0.001),
loss={
'regression':
smooth_l1(),
'classification':
focal(),
'transformation':
transformation_loss(model_3d_points_np=train_generator.
get_all_3d_model_points_array_for_loss(),
num_rotation_parameter=num_rotation_parameters)
},
loss_weights={
'regression': 1.0,
'classification': 1.0,
'transformation': 0.02
})
# create the callbacks
callbacks = create_callbacks(
model,
prediction_model,
validation_generator,
args,
)
if not args.compute_val_loss:
validation_generator = None
elif args.compute_val_loss and validation_generator is None:
raise ValueError(
'When you have no validation data, you should not specify --compute-val-loss.'
)
# start training
return model.fit_generator(generator=train_generator,
steps_per_epoch=args.steps,
initial_epoch=0,
epochs=args.epochs,
verbose=1,
callbacks=callbacks,
workers=args.workers,
use_multiprocessing=args.multiprocessing,
max_queue_size=args.max_queue_size,
validation_data=validation_generator)
def run_eval(args):
"""
Evaluate an EfficientPose model.
Args:
args: parseargs object containing configuration for the evaluation procedure.
"""
allow_gpu_growth_memory()
if args.validation_image_save_path:
os.makedirs(args.validation_image_save_path, exist_ok = True)
# create the generators
print("\nCreating the Generators...")
generator = create_generators(args)
print("Done!")
num_rotation_parameters = generator.get_num_rotation_parameters()
num_classes = generator.num_classes()
num_anchors = generator.num_anchors
print("\nBuilding the Model...")
prediction_model = TfliteWrapper(args.lite_model, args.score_threshold)
nlub, _, _, lite_model = build_EfficientPose(args.phi,
num_classes = num_classes,
num_anchors = num_anchors,
freeze_bn = args.freeze_bn,
score_threshold = args.score_threshold,
num_rotation_parameters = num_rotation_parameters,
print_architecture = False,
lite = args.lite,
no_se = args.no_se)
# inp = tf.keras.layers.Input((512,512,3))
# inp2 = tf.keras.layers.Input((6,))
# lite_model = tfkeras.EfficientNetB0(input_tensor=inp)
# lite_model = tf.keras.Model(inputs=[inp, inp2], outputs=lite_model)
print("Done!")
# load pretrained weights
print('Loading model, this may take a second...')
load_weights_rec(lite_model, args.weights)
for x, _ in generator:
o1 = lite_model(x)
o2 = prediction_model.predict_on_batch(x)
# print(x[0])
# print(o1[1].shape, o2[1].shape)
# print(o1[1].numpy(), o2[1])
# print(np.allclose(o1[1], o2[1], rtol=0.1, atol=0.1))
a = o1[0]
b = o2[0]
a = a.numpy().flatten()
b = b.flatten()
print(a.shape)
print(b.shape)
for i in range(a.size):
print([a[i], b[i]])
if i > 100:
break
break
