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