def _main_():
if tf.test.gpu_device_name():
print('>>>>> USING GPU: Default GPU Device: {}'.format(
tf.test.gpu_device_name()))
else:
print(">>>>> Please install GPU version of TF")
with open(CONFIG_FILE) as config_buffer:
config = json.loads(config_buffer.read())
################################
# Load data info
################################
print('>>>>> Loading the annotation data')
train_data_infos = parse_input_data(
image_folder=Path(config['train']['train_images_folder']),
annotation_folder=Path(config['train']['train_annotations_folder']),
annotation_extension=config['train']['annotations_format_extension'],
image_extension=config['train']['image_format_extension'])
train_dataset, validation_dataset = train_test_split(
train_data_infos,
test_size=config['train']['validation_dataset_ratio'])
################################
# Make and train model
################################
print('>>>>> Creating model')
yolo = YOLO(input_size=tuple(config['model']['input_size']),
grid_size=int(config['model']['grid_size']),
bbox_count=int(config['model']['bboxes_per_grid_cell']),
classes=config['model']['class_names'],
lambda_coord=config['model']['lambda_coord'],
lambda_noobj=config['model']['lambda_noobj'],
bbox_params=config['model']['bbox_params'])
print('>>>>> Starting the training process')
yolo.train_gen(training_infos=train_dataset,
validation_infos=validation_dataset,
save_model_path=config['train']['model_path'],
batch_size=config['train']['batch_size'],
nb_epochs=config['train']['nb_epochs'],
learning_rate=config['train']['learning_rate'],
use_pretrained_model=bool(
config['train']['use_pretrained_model']),
model_name=config['train']['model_name'],
steps_per_epoch=config['train']['steps_per_epoch'])
def _main_():
with open(CONFIG_FILE) as config_buffer:
config = json.loads(config_buffer.read())
################################
# Load data info
################################
train_data_infos = parse_input_data(
image_folder=Path(config['train']['train_images_folder']),
annotation_folder=Path(config['train']['train_annotations_folder']),
annotation_extension=config['train']['annotations_format_extension'],
image_extension=config['train']['image_format_extension'])
validation_data_infos = parse_input_data(
image_folder=Path(config['train']['validation_images_folder']),
annotation_folder=Path(
config['train']['validation_annotations_folder']),
annotation_extension=config['train']['annotations_format_extension'],
image_extension=config['train']['image_format_extension'])
################################
# Make and train model
################################
yolo = YOLO(input_size=tuple(config['model']['input_size']),
grid_size=int(config['model']['grid_size']),
bbox_count=int(config['model']['bboxes_per_grid_cell']),
classes=config['model']['class_names'],
lambda_coord=config['model']['lambda_coord'],
lambda_noobj=config['model']['lambda_noobj'],
bbox_params=config['model']['bbox_params'])
yolo.train_gen(training_infos=train_data_infos,
validation_infos=validation_data_infos,
save_weights_path=config['train']['trained_weights_path'],
batch_size=config['train']['batch_size'],
nb_epochs=config['train']['nb_epochs'],
learning_rate=config['train']['learning_rate'])
