def main():
args = parse_args()
print(args)
assert args.out or args.json_out, \
('Please specify at least one operation to save the results) '
'with the argument "--out" or "--json_out"')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
if args.json_out is not None and args.json_out.endswith('.json'):
args.json_out = args.json_out[:-5]
cfg = Config.fromfile(args.config)
cfg.model.pretrained = None
distributed = False
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
dataset = build_dataset(cfg.data.test)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
# dummy data to init network
img = tf.random.uniform(shape=[1333, 1333, 3], dtype=tf.float32)
img_meta = tf.constant(
[465., 640., 3., 800., 1101., 3., 1333., 1333., 3., 1.7204301, 0.],
dtype=tf.float32)
_ = model((tf.expand_dims(img, axis=0), tf.expand_dims(img_meta, axis=0)),
training=False)
load_checkpoint(model, args.checkpoint)
model.CLASSES = dataset.CLASSES
if not distributed:
outputs = single_gpu_test(model, dataset)
else:
raise NotImplementedError
rank, _, _, _ = get_dist_info()
if args.out and rank == 0:
print('\nwriting results to {}'.format(args.out))
fileio.dump(outputs, args.out)
eval_types = args.eval
if eval_types:
print('Starting evaluate {}'.format(' and '.join(eval_types)))
if eval_types == ['proposal_fast']:
result_file = args.out
coco_eval(result_file, eval_types, dataset.coco)
else:
if not isinstance(outputs[0], dict):
result_files = results2json(dataset, outputs, args.out)
coco_eval(result_files, eval_types, dataset.coco)
else:
for name in outputs[0]:
print('\nEvaluating {}'.format(name))
outputs_ = [out[name] for out in outputs]
result_file = args.out + '.{}'.format(name)
result_files = results2json(dataset, outputs_,
result_file)
coco_eval(result_files, eval_types, dataset.coco)
# Save predictions in the COCO json format
if args.json_out and rank == 0:
if not isinstance(outputs[0], dict):
results2json(dataset, outputs, args.json_out)
else:
for name in outputs[0]:
outputs_ = [out[name] for out in outputs]
result_file = args.json_out + '.{}'.format(name)
results2json(dataset, outputs_, result_file)
def build_dataloader(dataset,
imgs_per_gpu,
workers_per_gpu=1, # unused
num_gpus=0,
dist=True,
shuffle=True,
**kwargs):
"""Build a TF Dataset pipeline that returns padded batches.
In distributed training, each GPU/process has a dataloader.
In non-distributed training, there is only one dataloader for all GPUs.
Args:
dataset (Dataset): A dataset.
imgs_per_gpu (int): Number of images on each GPU, i.e., batch size of
each GPU.
workers_per_gpu (int): How many subprocesses to use for data loading
for each GPU. - TODO: unused
num_gpus (int): Number of GPUs. Only used in non-distributed training - TODO
dist (bool): Distributed training/test or not. Default: True.
shuffle (bool): Whether to shuffle the data at every epoch.
Default: True.
kwargs: any keyword argument to be used to initialize DataLoader
Returns:
tf.data.Dataset: A TF dataset pipeline.
"""
batch_size = imgs_per_gpu
if dist:
rank, local_rank, size, local_size = get_dist_info()
if dataset.train:
generator = data_generator.DataGenerator(dataset, index=rank, num_gpus=size, shuffle=shuffle)
else:
generator = data_generator.DataGenerator(dataset, index=rank, num_gpus=local_size, shuffle=False) # evaluation on node 0 workers
else:
generator = data_generator.DataGenerator(dataset, shuffle=False)
if dataset.train:
tf_dataset = tf.data.Dataset.from_generator(
generator, (tf.float32, tf.float32, tf.float32, tf.int32))
tf_dataset = tf_dataset.map(lambda *args: args, num_parallel_calls=tf.data.experimental.AUTOTUNE)
tf_dataset = tf_dataset.prefetch(tf.data.experimental.AUTOTUNE)
tf_dataset = tf_dataset.padded_batch(
batch_size,
padded_shapes=(
tf.TensorShape([None, None, 3]), # image padded to largest in batch
tf.TensorShape([11]), # image meta - no padding
tf.TensorShape([100, 4]), # bounding boxes, padded to longest
tf.TensorShape([100]) # labels, padded to longest
),
padding_values=(0.0, 0.0, 0.0, -1))
tf_dataset = tf_dataset.prefetch(tf.data.experimental.AUTOTUNE)
return tf_dataset, generator.num_examples // batch_size
else:
tf_dataset = tf.data.Dataset.from_generator(
generator, (tf.float32, tf.float32))
tf_dataset = tf_dataset.padded_batch(
batch_size,
padded_shapes=(
tf.TensorShape([None, None,
3]), # image padded to largest in batch
tf.TensorShape([11]), # image meta - no padding
),
padding_values=(0.0, 0.0))
tf_dataset = tf_dataset.repeat()
return tf_dataset, generator.num_examples // batch_size