def main(args=None):
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
if args.hard_score_rate:
hard_score_rate = float(args.hard_score_rate.lower())
else:
hard_score_rate = 0.5
print("hard_score_rate={}".format(hard_score_rate))
# make sure keras is the minimum required version
check_keras_version()
# optionally choose specific GPU
use_cpu = False
if args.gpu:
gpu_num = args.gpu
else:
gpu_num = str(0)
if use_cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(666)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num
##newly added line
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num = str(0)
keras.backend.tensorflow_backend.set_session(get_session())
# make save path if it doesn't exist
if args.save_path is not None and not os.path.exists(args.save_path):
os.makedirs(args.save_path)
# create the generator
generator = create_generator(args)
# load the model
print('Loading model, this may take a second...')
model = models.load_model(os.path.join(root_dir(), args.model),
backbone_name=args.backbone,
convert=args.convert_model,
nms=False)
# start prediction
while True:
predict(generator,
model,
score_threshold=args.score_threshold,
save_path=os.path.join(root_dir(), 'res_images_iou'),
hard_score_rate=hard_score_rate)
value = input("Please enter y to continue\n")
if value != 'y':
break
def main(args=None):
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
if DEBUG_MODE:
args.steps = 10
# create object that stores backbone information
backbone = models.backbone(args.backbone)
# make sure keras is the minimum required version
check_keras_version()
# optionally choose specific GPU
use_cpu = False
if args.gpu:
gpu_num = args.gpu
else:
gpu_num = str(0)
if use_cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(666)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num
keras.backend.tensorflow_backend.set_session(get_session())
# Weights and logs saves in a new locations
stmp = time.strftime("%c").replace(" ", "_")
args.snapshot_path = os.path.join(args.snapshot_path, stmp)
args.tensorboard_dir = os.path.join(args.tensorboard_dir, stmp)
print("Weights will be saved in {}".format(args.snapshot_path))
print("Logs will be saved in {}".format(args.tensorboard_dir))
create_folder(args.snapshot_path)
create_folder(args.tensorboard_dir)
# create the generators
train_generator, validation_generator = create_generators(args)
print('train_size:{},val_size:{}'.format(train_generator.size(),
validation_generator.size()))
# create the model
if args.snapshot is not None:
print('Loading model, this may take a second...')
model = models.load_model(args.snapshot, backbone_name=args.backbone)
training_model = model
prediction_model = retinanet_bbox(model=model)
else:
weights = os.path.join(os.path.join(root_dir(), args.weights))
# default to imagenet if nothing else is specified
if weights is None and args.imagenet_weights:
weights = backbone.download_imagenet()
print('Creating model, this may take a second...')
model, training_model, prediction_model = create_models(
backbone_retinanet=backbone.retinanet,
num_classes=train_generator.num_classes(),
weights=weights,
multi_gpu=args.multi_gpu,
freeze_backbone=args.freeze_backbone)
# print model summary
# print(model.summary())
# this lets the generator compute backbone layer shapes using the actual backbone model
if 'vgg' in args.backbone or 'densenet' in args.backbone:
compute_anchor_targets = functools.partial(
anchor_targets_bbox, shapes_callback=make_shapes_callback(model))
train_generator.compute_anchor_targets = compute_anchor_targets
if validation_generator is not None:
validation_generator.compute_anchor_targets = compute_anchor_targets
# create the callbacks
callbacks = create_callbacks(
model,
training_model,
prediction_model,
validation_generator,
args,
)
# start training
training_model.fit_generator(generator=train_generator,
steps_per_epoch=args.steps,
epochs=args.epochs,
verbose=1,
callbacks=callbacks,
validation_data=validation_generator,
validation_steps=validation_generator.size())
def main(image_folder_path: str, detection_save_folder: str, model_wights_path: str,
classes=None,
hard_score_rate = 0.5,
backbone = 'resnet50',
gpu = None,
score_threshold = 0.1,
iou_threshold = 0.5,
image_min_side = 800,
image_max_side = 1333,
args = None):
# parse arguments, to fill args with parameters initialized from defaults
if args is None:
args = ["--model", model_wights_path, 'csv']
args = parse_args(args)
if args.hard_score_rate is None:
args.hard_score_rate = hard_score_rate
# hard_score_rate = float(args.hard_score_rate.lower())
if args.model is None:
args.model = model_wights_path
if args.base_dir is None:
args.base_dir = image_folder_path
if args.backbone is None:
args.backbone = backbone
if args.gpu is None:
args.gpu = gpu
if args.score_threshold is None:
args.score_threshold = float(score_threshold)
if args.iou_threshold is None:
args.iou_threshold = float(iou_threshold)
if args.save_path is None:
args.save_path = detection_save_folder
if args.image_max_side is None:
args.image_max_side = int(image_max_side)
if args.image_min_side is None:
args.image_min_side = int(image_min_side)
# make sure keras is the minimum required version
check_keras_version()
# optionally choose specific GPU
use_cpu = False
if args.gpu:
gpu_num = args.gpu
else:
gpu_num = str(0)
if use_cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(666)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num
keras.backend.tensorflow_backend.set_session(get_session())
# make save path if it doesn't exist
if args.save_path is not None and not os.path.exists(args.save_path):
os.makedirs(args.save_path)
# create the generator
generator = create_generator(args, args.base_dir)
# load the model
print('Loading model, this may take a second...')
model = models.load_model(os.path.join(root_dir(), args.model), backbone_name=args.backbone, convert=args.convert_model, nms=False)
# start prediction
detection_csv_results_file_path = predict(
generator,
model,
score_threshold=args.score_threshold,
# save_path=os.path.join(root_dir(), 'res_images_iou'),
image_save_path = os.path.join(args.save_path, "detection_results_images"),
results_save_path = args.save_path,
hard_score_rate=hard_score_rate)
return detection_csv_results_file_path