def main(args): DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args['flow_control']['use_cuda'] else "cpu") # eval_path = pathlib.Path(args.eval_dir) # eval_path.mkdir(exist_ok=True) if not os.path.exists(args['flow_control']['eval_dir']): os.mkdir(args['flow_control']['eval_dir']) timer = Timer() class_names = [ name.strip() for name in open(args['flow_control']['label_file']).readlines() ] _net = args['flow_control']['net'] _dataset_type = args['flow_control']['dataset_type'] if _dataset_type == "voc": raise NotImplementedError("Not implement error") dataset = VOCDataset(args['flow_control']['dataset'], is_test=True) elif _dataset_type == 'open_images': raise NotImplementedError("Not implement error") dataset = OpenImagesDataset(args['flow_control']['dataset'], dataset_type="test") elif _dataset_type == "coco": # dataset = CocoDetection("/home/wenyen4desh/datasets/coco/test2017","/home/wenyen4desh/datasets/annotations/image_info_test2017.json") #dataset = CocoDetection("../../dataset/datasets/coco/val2017","../../dataset/datasets/coco/annotations/instances_val2017.json") # dataset = CocoDetection("/home/wenyen4desh/datasets/coco/train2017","/home/wenyen4desh/datasets/coco/annotations/instances_train2017.json") dataset = CocoDetection(args['Datasets']['coco']['val_image_path'], args['Datasets']['coco']['val_anno_path']) elif _dataset_type == "ecp": dataset = EuroCity_Dataset(args['Datasets']['ecp']['val_image_path'], args['Datasets']['ecp']['val_anno_path']) true_case_stat, all_gb_boxes, all_difficult_cases = group_annotation_by_class( dataset) if _net == 'vgg16-ssd': net = create_vgg_ssd(len(class_names), is_test=True) elif _net == 'mb1-ssd': net = create_mobilenetv1_ssd(len(class_names), is_test=True) elif _net == 'mb1-ssd-lite': net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True) elif _net == 'sq-ssd-lite': net = create_squeezenet_ssd_lite(len(class_names), is_test=True) elif _net == 'mb2-ssd-lite': net = create_mobilenetv2_ssd_lite( len(class_names), width_mult=args['flow_control']['mb2_width_mult'], is_test=True) else: logging.fatal( "The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite." ) parser.print_help(sys.stderr) sys.exit(1) #train_transform = MatchPrior(config.priors, config.center_variance, # config.size_variance, 0.5) #test_transform = TestTransform(config.image_size, config.image_mean, config.image_std) import pdb pdb.set_trace() ############################## automatically validation ############################################ timer.start("Load Model") net.load(args['flow_control']['trained_model']) net = net.to(DEVICE) print('It took {} seconds to load the model.'.format( timer.end("Load Model"))) _nms_method = args['flow_control']['nms_method'] if _net == 'vgg16-ssd': predictor = create_vgg_ssd_predictor(net, nms_method=_nms_method, device=DEVICE) elif _net == 'mb1-ssd': predictor = create_mobilenetv1_ssd_predictor(net, nms_method=_nms_method, device=DEVICE) elif _net == 'mb1-ssd-lite': predictor = create_mobilenetv1_ssd_lite_predictor( net, nms_method=_nms_method, device=DEVICE) elif _net == 'sq-ssd-lite': predictor = create_squeezenet_ssd_lite_predictor( net, nms_method=_nms_method, device=DEVICE) elif _net == 'mb2-ssd-lite': predictor = create_mobilenetv2_ssd_lite_predictor( net, nms_method=_nms_method, device=DEVICE) else: logging.fatal( "The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite." ) parser.print_help(sys.stderr) sys.exit(1) results = [] # Predict Bounding Box for i in range(len(dataset)): print("process image {}", i) timer.start("Load Image") image = dataset.get_image(i) print("Load Image: {:4f} seconds.".format(timer.end("Load Image"))) timer.start("Predict") boxes, labels, probs = predictor.predict(image) print("Prediction: {:4f} seconds.".format(timer.end("Predict"))) indexes = torch.ones(labels.size(0), 1, dtype=torch.float32) * i results.append( torch.cat( [ indexes.reshape(-1, 1), labels.reshape(-1, 1).float(), probs.reshape(-1, 1), boxes + 1.0 # matlab's indexes start from 1 ], dim=1)) results = torch.cat(results) # Write the result to file for class_index, class_name in enumerate(class_names): if class_index == 0: continue # ignore background file_name = "det_test_{}.txt".format(class_name) prediction_path = os.path.join(args['flow_control']['eval_dir'], file_name) with open(prediction_path, "w") as f: sub = results[results[:, 1] == class_index, :] for i in range(sub.size(0)): prob_box = sub[i, 2:].numpy() image_id, _ = dataset.get_annotation(int(sub[i, 0])) f.write( str(image_id) + " " + " ".join([str(v) for v in prob_box]) + "\n") # image_id = dataset.ids[int(sub[i, 0])] # print(str(image_id) + " " + " ".join([str(v) for v in prob_box]), file=f) aps = [] prcs = [] recalls = [] print("\n\nAverage Precision Per-class:") for class_index, class_name in enumerate(class_names): if class_index == 0: continue file_name = "det_test_{}.txt".format(class_name) prediction_path = os.path.join(args['flow_control']['eval_dir'], file_name) # [email protected] evaluation method ap, precision, recall = compute_average_precision_per_class( args, true_case_stat[class_index], all_gb_boxes[class_index], all_difficult_cases[class_index], prediction_path, args['flow_control']['iou_threshold'], args['flow_control']['use_2007_metric']) # # COCO eval # ap, precision, recall = coco_ap_per_class( # true_case_stat[class_index], # all_gb_boxes[class_index], # all_difficult_cases[class_index], # prediction_path, # args.use_2007_metric # ) aps.append(ap) prcs.append(precision) recalls.append(recall) print("{}: {}".format(class_name, ap)) print("\nAverage Precision Across All Classes:{}".format( sum(aps[0:5]) / len(aps[0:5]))) print("\nAverage Precision :{}".format(sum(prcs[0:5]) / len(prcs[0:5]))) print("\nAverage Recall :{}".format(sum(recalls[0:5]) / len(recalls[0:5])))
def dataset_loading(args, config): train_transform = TrainAugmentation(config.image_size, config.image_mean, config.image_std) target_transform = MatchPrior(config.priors, config.center_variance, config.size_variance, 0.5) test_transform = TestTransform(config.image_size, config.image_mean, config.image_std) logging.info("Prepare training datasets.") dataset_name = args['flow_control']['dataset_type'] if dataset_name == 'voc': dataset = VOCDataset("", transform=train_transform, target_transform=target_transform) label_txt_name = "voc-model-labels.txt" elif dataset_name == 'open_images': dataset = OpenImagesDataset(dataset_path, transform=train_transform, target_transform=target_transform, dataset_type="train", balance_data=args.balance_data) label_txt_name = "open-images-model-labels.txt" elif dataset_name == 'coco': dataset = CocoDetection(args["Datasets"]["coco"]["train_image_path"], args["Datasets"]["coco"]["train_anno_path"], transform=train_transform, target_transform=target_transform) label_txt_name = "open-images-model-labels.txt" elif dataset_name == 'ecp': #dataset = EuroCity_Dataset( args["Datasets"]["ecp"]["train_image_path"], # args["Datasets"]["ecp"]["train_anno_path"], # transform=train_transform, target_transform=target_transform) dataset = ECP_table_comm(args["Datasets"]["ecp"]["train_image_path"], args["Datasets"]["ecp"]["train_anno_path"], transform=train_transform, target_transform=target_transform) dataset.Active_mode() if len(dataset) == 0: raise ValueError("Doesn't exist any file") label_txt_name = "open-images-model-labels.txt" elif dataset_name == "ecp-random": dataset = ECP_subsample_dataset( args["Datasets"]["ecp"]["train_image_path"], args["Datasets"]["ecp"]["train_anno_path"], transform=train_transform, target_transform=target_transform, _sampling_mode="random", ratio=args['flow_control']['dataset_ratio']) label_txt_name = "open-images-model-labels.txt" elif dataset_name == "ecp-centroid": dataset = ECP_subsample_dataset( args["Datasets"]["ecp"]["train_image_path"], args["Datasets"]["ecp"]["train_anno_path"], transform=train_transform, target_transform=target_transform, _sampling_mode="centroid", ratio=args['flow_control']['dataset_ratio']) label_txt_name = "open-images-model-labels.txt" elif dataset_name in ["virat", "VIRAT"]: # dataset = VIRAT_Loader(args["Datasets"]["virat"]["train_image_path"], # args["Datasets"]["virat"]["train_anno_path"], # transform=train_transform, target_transform=target_transform) dataset = VIRAT_Dataset( args["Datasets"]["virat_seq"]["train_image_path"], args["Datasets"]["virat_seq"]["train_anno_path"], transform=train_transform, target_transform=target_transform, downpurning_ratio=0.2) #0.2 # dataset = VIRAT_table_comm(args["Datasets"]["virat_seq"]["train_image_path"], # args["Datasets"]["virat_seq"]["train_anno_path"], # transform=train_transform, target_transform=target_transform) label_txt_name = "virat_labels.txt" else: raise ValueError( "Dataset type {} is not supported.".format(dataset_name)) label_file = os.path.join(args["flow_control"]["checkpoint_folder"], label_txt_name) if os.path.exists(label_file): store_labels(label_file, dataset.class_names) logging.info(dataset) num_classes = len(dataset.class_names) train_dataset = dataset logging.info("Stored labels into file {}.".format(label_file)) logging.info("Train dataset size: {}".format(len(train_dataset))) logging.debug("================= train_loader ===================") logging.debug("DataLoader batchsize : ", args["flow_control"]["batch_size"]) # if dataset_name == "virat": # indicies = np.arange(args["flow_control"]["batch_size"]) # train_loader = DataLoader(train_dataset, args["flow_control"]["batch_size"], # num_workers=args["flow_control"]["num_workers"], # shuffle=False, sampler=SubsetRandomSampler(indicies)) # else: train_loader = DataLoader(train_dataset, args["flow_control"]["batch_size"], num_workers=args["flow_control"]["num_workers"], shuffle=True) logging.info("Prepare Validation datasets.") if dataset_name == "voc": raise NotImplementedError("Doesn't modify") val_dataset = VOCDataset("", transform=test_transform, target_transform=target_transform, is_test=True) elif dataset_name == 'open_images': val_dataset = OpenImagesDataset(dataset_path, transform=test_transform, target_transform=target_transform, dataset_type="test") logging.info(val_dataset) elif dataset_name == "coco": val_dataset = CocoDetection(args["Datasets"]["coco"]["val_image_path"], args["Datasets"]["coco"]["val_anno_path"], transform=test_transform, target_transform=target_transform) logging.info(val_dataset) elif dataset_name in ["ecp", "ecp-random", "ecp-centroid"]: val_dataset = EuroCity_Dataset( args["Datasets"]["ecp"]["val_image_path"], args["Datasets"]["ecp"]["val_anno_path"], transform=test_transform, target_transform=target_transform) # elif dataset_name = "ecp-random": # val_dataset = ECP_subsample_dataset( args["Datasets"]["ecp"]["val_image_path"], # args["Datasets"]["ecp"]["val_anno_path"], # transform=test_transform, target_transform=target_transform, _sampling_mode = "random", ratio = 0.1) # elif dataset_name = "ecp-centroid": # val_dataset = ECP_subsample_dataset( args["Datasets"]["ecp"]["val_image_path"], # args["Datasets"]["ecp"]["val_anno_path"], # transform=test_transform, target_transform=target_transform, _sampling_mode = "centroid", ratio = 0.1) elif dataset_name in ["virat", "VIRAT"]: # val_dataset = VIRAT_Loader(args["Datasets"]["virat"]["test_image_path"], # args["Datasets"]["virat"]["test_anno_path"], # transform=train_transform, target_transform=target_transform) val_dataset = VIRAT_Dataset( args["Datasets"]["virat_seq"]["train_image_path"], args["Datasets"]["virat_seq"]["train_anno_path"], transform=train_transform, target_transform=target_transform, downpurning_ratio=0.2 * 3. / 9.) logging.info("validation dataset size: {}".format(len(val_dataset))) # if dataset_name == "virat": # indicies = np.arange(args["flow_control"]["batch_size"]) # val_loader = DataLoader(train_dataset, args["flow_control"]["batch_size"], # num_workers=args["flow_control"]["num_workers"], # shuffle=False, sampler=SubsetRandomSampler(indicies)) # else: # val_loader = DataLoader(val_dataset, args["flow_control"]["batch_size"], # num_workers=args["flow_control"]["num_workers"], # shuffle=False) val_loader = DataLoader(val_dataset, args["flow_control"]["batch_size"], num_workers=args["flow_control"]["num_workers"], shuffle=False) logging.info("Build network.") return train_loader, val_loader, num_classes
def main(args): DEVICE = torch.device( "cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu") #DEVICE = torch.device("cpu") if args.use_cuda and torch.cuda.is_available(): torch.backends.cudnn.benchmark = True logging.info("Use Cuda.") timer = Timer() logging.info(args) if args.net == 'vgg16-ssd': create_net = create_vgg_ssd config = vgg_ssd_config elif args.net == 'mb1-ssd': create_net = create_mobilenetv1_ssd config = mobilenetv1_ssd_config elif args.net == 'mb1-ssd-lite': create_net = create_mobilenetv1_ssd_lite config = mobilenetv1_ssd_config elif args.net == 'sq-ssd-lite': create_net = create_squeezenet_ssd_lite config = squeezenet_ssd_config elif args.net == 'mb2-ssd-lite': create_net = lambda num: create_mobilenetv2_ssd_lite( num, width_mult=args.mb2_width_mult) config = mobilenetv1_ssd_config else: logging.fatal("The net type is wrong.") parser.print_help(sys.stderr) sys.exit(1) train_transform = TrainAugmentation(config.image_size, config.image_mean, config.image_std) target_transform = MatchPrior(config.priors, config.center_variance, config.size_variance, 0.5) test_transform = TestTransform(config.image_size, config.image_mean, config.image_std) logging.info("Prepare training datasets.") datasets = [] for dataset_path in args.datasets: if args.dataset_type == 'voc': dataset = VOCDataset(dataset_path, transform=train_transform, target_transform=target_transform) label_file = os.path.join(args.checkpoint_folder, "voc-model-labels.txt") store_labels(label_file, dataset.class_names) num_classes = len(dataset.class_names) elif args.dataset_type == 'open_images': dataset = OpenImagesDataset(dataset_path, transform=train_transform, target_transform=target_transform, dataset_type="train", balance_data=args.balance_data) label_file = os.path.join(args.checkpoint_folder, "open-images-model-labels.txt") store_labels(label_file, dataset.class_names) logging.info(dataset) num_classes = len(dataset.class_names) elif args.dataset_type == 'coco': # root, annFile, transform=None, target_transform=None, transforms=None) # dataset_type="train", balance_data=args.balance_data) dataset = CocoDetection( "/home/wenyen4desh/datasets/coco/train2017", "/home/wenyen4desh/datasets/coco/annotations/instances_train2017.json", transform=train_transform, target_transform=target_transform) label_file = os.path.join(args.checkpoint_folder, "open-images-model-labels.txt") store_labels(label_file, dataset.class_names) logging.info(dataset) num_classes = len(dataset.class_names) # raise ValueError("Dataset type {} yet implement.".format(args.dataset_type)) else: raise ValueError("Dataset type {} is not supported.".format( args.dataset_type)) datasets.append(dataset) logging.info("Stored labels into file {}.".format(label_file)) train_dataset = ConcatDataset(datasets) logging.info("Train dataset size: {}".format(len(train_dataset))) train_loader = DataLoader(train_dataset, args.batch_size, num_workers=args.num_workers, shuffle=True) logging.info("Prepare Validation datasets.") if args.dataset_type == "voc": val_dataset = VOCDataset(args.validation_dataset, transform=test_transform, target_transform=target_transform, is_test=True) elif args.dataset_type == 'open_images': val_dataset = OpenImagesDataset(dataset_path, transform=test_transform, target_transform=target_transform, dataset_type="test") logging.info(val_dataset) elif args.dataset_type == "coco": val_dataset = CocoDetection( "/home/wenyen4desh/datasets/coco/val2017", "/home/wenyen4desh/datasets/coco/annotations/instances_val2017.json", transform=test_transform, target_transform=target_transform) logging.info(val_dataset) logging.info("validation dataset size: {}".format(len(val_dataset))) val_loader = DataLoader(val_dataset, args.batch_size, num_workers=args.num_workers, shuffle=False) logging.info("Build network.") net = create_net(num_classes) min_loss = -10000.0 last_epoch = -1 base_net_lr = args.base_net_lr if args.base_net_lr is not None else args.lr extra_layers_lr = args.extra_layers_lr if args.extra_layers_lr is not None else args.lr if args.freeze_base_net: logging.info("Freeze base net.") freeze_net_layers(net.base_net) params = itertools.chain(net.source_layer_add_ons.parameters(), net.extras.parameters(), net.regression_headers.parameters(), net.classification_headers.parameters()) params = [{ 'params': itertools.chain(net.source_layer_add_ons.parameters(), net.extras.parameters()), 'lr': extra_layers_lr }, { 'params': itertools.chain(net.regression_headers.parameters(), net.classification_headers.parameters()) }] elif args.freeze_net: freeze_net_layers(net.base_net) freeze_net_layers(net.source_layer_add_ons) freeze_net_layers(net.extras) params = itertools.chain(net.regression_headers.parameters(), net.classification_headers.parameters()) logging.info("Freeze all the layers except prediction heads.") else: params = [{ 'params': net.base_net.parameters(), 'lr': base_net_lr }, { 'params': itertools.chain(net.source_layer_add_ons.parameters(), net.extras.parameters()), 'lr': extra_layers_lr }, { 'params': itertools.chain(net.regression_headers.parameters(), net.classification_headers.parameters()) }] timer.start("Load Model") if args.resume: logging.info("Resume from the model {}".format(args.resume)) net.load(args.resume) elif args.base_net: logging.info("Init from base net {}".format(args.base_net)) net.init_from_base_net(args.base_net) elif args.pretrained_ssd: logging.info("Init from pretrained ssd {}".format(args.pretrained_ssd)) net.init_from_pretrained_ssd(args.pretrained_ssd) logging.info('Took {:.2f} seconds to load the model.'.format( timer.end("Load Model"))) net.to(DEVICE) criterion = MultiboxLoss(config.priors, iou_threshold=0.5, neg_pos_ratio=3, center_variance=0.1, size_variance=0.2, device=DEVICE) optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) logging.info("Learning rate: {}, Base net learning rate: {}, ".format( args.lr, base_net_lr) + "Extra Layers learning rate: {}.".format(extra_layers_lr)) if args.scheduler == 'multi-step': logging.info("Uses MultiStepLR scheduler.") milestones = [int(v.strip()) for v in args.milestones.split(",")] scheduler = MultiStepLR(optimizer, milestones=milestones, gamma=0.1, last_epoch=last_epoch) elif args.scheduler == 'cosine': logging.info("Uses CosineAnnealingLR scheduler.") scheduler = CosineAnnealingLR(optimizer, args.t_max, last_epoch=last_epoch) else: logging.fatal("Unsupported Scheduler: {}.".format(args.scheduler)) parser.print_help(sys.stderr) sys.exit(1) logging.info("Start training from epoch {}.".format(last_epoch + 1)) for epoch in range(last_epoch + 1, args.num_epochs): scheduler.step() train(train_loader, net, criterion, optimizer, device=DEVICE, debug_steps=args.debug_steps, epoch=epoch) if epoch % args.validation_epochs == 0 or epoch == args.num_epochs - 1: val_loss, val_regression_loss, val_classification_loss = test( val_loader, net, criterion, DEVICE) logging.info("Epoch: {}, ".format(epoch) + "Validation Loss: {:.4f}, ".format(val_loss) + "Validation Regression Loss {:.4f}, ".format( val_regression_loss) + "Validation Classification Loss: {:.4f}".format( val_classification_loss)) model_path = os.path.join( args.checkpoint_folder, "{}-Epoch-{}-Loss-{}.pth".format(args.net, epoch, val_loss)) net.save(model_path) logging.info("Saved model {}".format(model_path))
if __name__ == '__main__': # eval_path = pathlib.Path(args.eval_dir) # eval_path.mkdir(exist_ok=True) if not os.path.exists(args.eval_dir): os.mkdir(args.eval_dir) timer = Timer() class_names = [name.strip() for name in open(args.label_file).readlines()] if args.dataset_type == "voc": dataset = VOCDataset(args.dataset, is_test=True) elif args.dataset_type == 'open_images': dataset = OpenImagesDataset(args.dataset, dataset_type="test") elif args.dataset_type == "coco": # dataset = CocoDetection("/home/wenyen4desh/datasets/coco/test2017","/home/wenyen4desh/datasets/annotations/image_info_test2017.json") dataset = CocoDetection( "../../dataset/datasets/coco/val2017", "../../dataset/datasets/coco/annotations/instances_val2017.json") # dataset = CocoDetection("/home/wenyen4desh/datasets/coco/train2017","/home/wenyen4desh/datasets/coco/annotations/instances_train2017.json") true_case_stat, all_gb_boxes, all_difficult_cases = group_annotation_by_class( dataset) if args.net == 'vgg16-ssd': net = create_vgg_ssd(len(class_names), is_test=True) elif args.net == 'mb1-ssd': net = create_mobilenetv1_ssd(len(class_names), is_test=True) elif args.net == 'mb1-ssd-lite': net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True) elif args.net == 'sq-ssd-lite': net = create_squeezenet_ssd_lite(len(class_names), is_test=True) elif args.net == 'mb2-ssd-lite': net = create_mobilenetv2_ssd_lite(len(class_names),
def dataset_loading(args, config): train_transform = TrainAugmentation(config.image_size, config.image_mean, config.image_std) target_transform = MatchPrior(config.priors, config.center_variance, config.size_variance, 0.5) test_transform = TestTransform(config.image_size, config.image_mean, config.image_std) test_normal_transform = TestTransform(config.image_size, 0, 1) logging.info("Prepare training datasets.") dataset_name = args['flow_control']['dataset_type'] if dataset_name == 'voc': dataset = VOCDataset(dataset_path, transform=train_transform, target_transform=target_transform) label_txt_name = "voc-model-labels.txt" elif dataset_name == 'open_images': dataset = OpenImagesDataset(dataset_path, transform=train_transform, target_transform=target_transform, dataset_type="train", balance_data=args.balance_data) label_txt_name = "open-images-model-labels.txt" elif dataset_name == 'coco': dataset = CocoDetection(args["Datasets"]["coco"]["train_image_path"], args["Datasets"]["coco"]["train_anno_path"], transform=test_normal_transform, target_transform=target_transform) # target_transform=target_transform) # transform=train_transform, target_transform=target_transform) label_txt_name = "open-images-model-labels.txt" elif dataset_name == "ecp": dataset = EuroCity_Dataset(args["Datasets"]["ecp"]["train_image_path"], args["Datasets"]["ecp"]["train_anno_path"], transform=test_normal_transform, target_transform=target_transform) label_txt_name = "open-images-model-labels.txt" elif dataset_name == "ecp-random": dataset = ECP_subsample_dataset( args["Datasets"]["ecp"]["train_image_path"], args["Datasets"]["ecp"]["train_anno_path"], transform=test_normal_transform, target_transform=target_transform, _sampling_mode="random", ratio=0.1) label_txt_name = "open-images-model-labels.txt" elif dataset_name == "ecp-centroid": dataset = ECP_subsample_dataset( args["Datasets"]["ecp"]["train_image_path"], args["Datasets"]["ecp"]["train_anno_path"], transform=test_normal_transform, target_transform=target_transform, _sampling_mode="centroid", ratio=0.1) label_txt_name = "open-images-model-labels.txt" else: raise ValueError( "Dataset type {} is not supported.".format(dataset_name)) label_file = os.path.join(args["flow_control"]["checkpoint_folder"], label_txt_name) if os.path.exists(label_file): store_labels(label_file, dataset.class_names) logging.info(dataset) num_classes = len(dataset.class_names) train_dataset = dataset logging.info("Stored labels into file {}.".format(label_file)) logging.info("Train dataset size: {}".format(len(train_dataset))) train_loader = DataLoader(train_dataset, args["flow_control"]["batch_size"], num_workers=args["flow_control"]["num_workers"], shuffle=True) logging.info("Prepare Validation datasets.") if dataset_name == "voc": val_dataset = VOCDataset(args.validation_dataset, transform=test_transform, target_transform=target_transform, is_test=True) elif dataset_name == 'open_images': val_dataset = OpenImagesDataset(dataset_path, transform=test_transform, target_transform=target_transform, dataset_type="test") logging.info(val_dataset) elif dataset_name == "coco": val_dataset = CocoDetection(args["Datasets"]["coco"]["val_image_path"], args["Datasets"]["coco"]["val_anno_path"], transform=test_transform, target_transform=target_transform) logging.info(val_dataset) elif dataset_name == "ecp": val_dataset = EuroCity_Dataset( args["Datasets"]["ecp"]["val_image_path"], args["Datasets"]["ecp"]["val_anno_path"], transform=test_transform, target_transform=target_transform) logging.info(val_dataset) elif dataset_name == "ecp-random": val_dataset = ECP_subsample_dataset( args["Datasets"]["ecp"]["val_image_path"], args["Datasets"]["ecp"]["val_anno_path"], transform=test_transform, target_transform=target_transform, _sampling_mode="random", ratio=0.1) elif dataset_name == "ecp-centroid": val_dataset = ECP_subsample_dataset( args["Datasets"]["ecp"]["val_image_path"], args["Datasets"]["ecp"]["val_anno_path"], transform=test_transform, target_transform=target_transform, _sampling_mode="centroid", ratio=0.1) logging.info("validation dataset size: {}".format(len(val_dataset))) val_loader = DataLoader(val_dataset, args["flow_control"]["batch_size"], num_workers=args["flow_control"]["num_workers"], shuffle=False) logging.info("Build network.") return train_loader, val_loader, num_classes