def get_model(config, args):
model_config = config["model"]
if model_config["name"] == "ssd_vgg16":
with open(args.label_maps, "r") as label_map_file:
label_maps = [i.strip("\n") for i in label_map_file.readlines()]
return SSD_VGG16(
config=config,
label_maps=label_maps,
is_training=True
)
elif model_config["name"] == "ssd_mobilenetv1":
with open(args.label_maps, "r") as label_map_file:
label_maps = [i.strip("\n") for i in label_map_file.readlines()]
return SSD_MOBILENET(
config=config,
label_maps=label_maps,
is_training=True
)
elif model_config["name"] == "ssd_mobilenetv2":
with open(args.label_maps, "r") as label_map_file:
label_maps = [i.strip("\n") for i in label_map_file.readlines()]
return SSD_MOBILENETV2(
config=config,
label_maps=label_maps,
is_training=True
)
elif model_config["name"] == "tbpp_vgg16":
return TBPP_VGG16(
config=config,
is_training=True
)
elif model_config["name"] == "qssd_vgg16":
with open(args.label_maps, "r") as label_map_file:
label_maps = [i.strip("\n") for i in label_map_file.readlines()]
model = QSSD_VGG16(
config=config,
label_maps=label_maps,
is_training=True
)
return model
elif model_config["name"] == "qssd_mobilenetv2":
with open(args.label_maps, "r") as label_map_file:
label_maps = [i.strip("\n") for i in label_map_file.readlines()]
model = QSSD_MOBILENETV2(
config=config,
label_maps=label_maps,
is_training=True
)
return model
else:
print(
f"model with name ${model_config['name']} has not been implemented yet")
exit()
def ssd_mobilenetv2(config, args):
assert args.label_maps is not None, "please specify a label map file"
assert os.path.exists(args.label_maps), "label_maps file does not exist"
with open(args.label_maps, "r") as file:
label_maps = [line.strip("\n") for line in file.readlines()]
model = SSD_MOBILENETV2(
config,
label_maps,
is_training=False,
num_predictions=args.num_predictions
)
process_input_fn = mobilenet_v2.preprocess_input
return model, process_input_fn, label_maps
def inference_ssd_mobilenetv2(config, args):
""""""
assert args.label_maps is not None, "please specify a label map file"
assert os.path.exists(args.label_maps), "label_maps file does not exist"
with open(args.label_maps, "r") as file:
label_maps = [line.strip("\n") for line in file.readlines()]
model = SSD_MOBILENETV2(
config,
label_maps,
is_training=False,
num_predictions=args.num_predictions)
process_input_fn = mobilenet_v2.preprocess_input
image, bboxes, classes = ssd_utils.read_sample(
image_path=args.input_image,
label_path=args.label_file,
)
return model, label_maps, process_input_fn, np.uint8(image), bboxes, classes
def train_ssd_mobilenetv2(config, args):
""""""
assert args.label_maps is not None, "please specify a label maps file for this model"
assert os.path.exists(args.label_maps), "label_maps file does not exist"
with open(args.label_maps, "r") as file:
label_maps = [line.strip("\n") for line in file.readlines()]
assert args.training_split is not None, "please specify a training split file for this model"
assert os.path.exists(
args.training_split), "training_split file does not exist"
training_samples = data_utils.get_samples_from_split(
split_file=args.training_split,
images_dir=args.images_dir,
labels_dir=args.labels_dir)
assert args.batch_size <= len(
training_samples
), "batch_size less than or equal to len(training_samples)"
training_config = config["training"]
model = SSD_MOBILENETV2(config=config)
loss = SSD_LOSS(
alpha=training_config["alpha"],
min_negative_boxes=training_config["min_negative_boxes"],
negative_boxes_ratio=training_config["negative_boxes_ratio"])
optimizer = SGD(lr=args.learning_rate,
momentum=0.9,
decay=0.0005,
nesterov=False)
generator = SSD_DATA_GENERATOR(
samples=training_samples,
label_maps=label_maps,
config=config,
shuffle=args.shuffle,
batch_size=args.batch_size,
augment=args.augment,
process_input_fn=mobilenet_v2.preprocess_input)
model.compile(optimizer=optimizer, loss=loss.compute)
return model, generator, optimizer, loss, training_samples
label_maps,
is_training=False,
num_predictions=args.num_predictions)
elif model_config["name"] == "ssd_mobilenetv1":
with open(args.label_maps, "r") as file:
label_maps = [line.strip("\n") for line in file.readlines()]
model = SSD_MOBILENET(config,
label_maps,
is_training=False,
num_predictions=args.num_predictions)
elif model_config["name"] == "ssd_mobilenetv2":
with open(args.label_maps, "r") as file:
label_maps = [line.strip("\n") for line in file.readlines()]
model = SSD_MOBILENETV2(config,
label_maps,
is_training=False,
num_predictions=args.num_predictions)
else:
print(
f"model with name ${model_config['name']} has not been implemented yet"
)
exit()
model.load_weights(args.weights)
config_file_name = os.path.basename(args.config)
config_file_name = config_file_name[:config_file_name.index(".")]
if args.output_type == "keras":
model.save(os.path.join(args.output_dir, f"{config_file_name}.h5"))
elif args.output_type == "tflite":
tflite_converter = tf.lite.TFLiteConverter.from_keras_model(model)
def ssd_mobilenetv2(config, args):
training_config = config["training"]
with open(args.label_maps, "r") as label_map_file:
label_maps = [i.strip("\n") for i in label_map_file.readlines()]
training_samples = data_utils.get_samples_from_split(
split_file=args.training_split,
images_dir=args.images_dir,
labels_dir=args.labels_dir)
if args.validation_split is not None:
validation_samples = data_utils.get_samples_from_split(
split_file=args.validation_split,
images_dir=args.images_dir,
labels_dir=args.labels_dir)
training_data_generator = SSD_DATA_GENERATOR(
samples=training_samples,
config=config,
label_maps=label_maps,
shuffle=args.shuffle,
batch_size=args.batch_size,
augment=args.augment,
process_input_fn=preprocess_input)
if args.validation_split is not None:
print("-- validation split specified")
validation_data_generator = SSD_DATA_GENERATOR(
samples=validation_samples,
config=config,
label_maps=label_maps,
shuffle=args.shuffle,
batch_size=args.batch_size,
augment=False,
process_input_fn=preprocess_input)
loss = SSD_LOSS(
alpha=training_config["alpha"],
min_negative_boxes=training_config["min_negative_boxes"],
negative_boxes_ratio=training_config["negative_boxes_ratio"])
model = SSD_MOBILENETV2(config=config,
label_maps=label_maps,
is_training=True)
optimizer = SGD(lr=args.learning_rate,
momentum=0.9,
decay=0.0005,
nesterov=False)
model.compile(optimizer=optimizer, loss=loss.compute)
if args.checkpoint is not None:
assert os.path.exists(args.checkpoint), "checkpoint does not exist"
model.load_weights(args.checkpoint, by_name=True)
model.fit(
x=training_data_generator,
validation_data=validation_data_generator
if args.validation_split is not None else None,
batch_size=args.batch_size,
validation_batch_size=args.batch_size,
epochs=args.epochs,
callbacks=[
ModelCheckpoint(filepath=os.path.join(
args.output_dir, "cp_{epoch:02d}_loss-{loss:.2f}.h5"
if args.validation_split is None else
"cp_{epoch:02d}_loss-{loss:.2f}_valloss-{val_loss:.2f}.h5"),
save_weights_only=True,
monitor='loss'
if args.validation_split is None else 'val_loss',
mode='min')
])
model.save_weights(os.path.join(args.output_dir, "model.h5"))
# )
# cv2.putText(
# image,
# coco.cats[category_id]["name"],
# (int(bbox[0]), int(bbox[1])),
# cv2.FONT_HERSHEY_SIMPLEX,
# 1,
# (0, 0, 255),
# 1
# )
# cv2.imshow('image', image)
# if cv2.waitKey(0) == ord('q'):
# cv2.destroyAllWindows()
# exit()
import json
from networks import SSD_MOBILENETV2
with open("configs/ssd320_mobilenetv2_coco2017-train.json", "r") as config_file:
config = json.load(config_file)
with open("/Users/socretlee/Downloads/fast-ai-coco/label_maps.txt", "r") as file:
label_maps = [line.strip("\n") for line in file.readlines()]
model = SSD_MOBILENETV2(
config,
label_maps=label_maps
)
model.summary()