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