def main():
args = parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpuid
img_dir = args.img_dir
out_dir = args.out_dir
batch_size = args.batch_size
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
# TODO: support multiple images per gpu (only minor changes are needed)
file_list = common.load_filepaths(img_dir,
suffix=('.jpg', '.png', '.jpeg'),
recursive=True)
print(file_list[:10])
print('imgs: ', len(file_list))
#
print(file_list[0].replace(img_dir, ''))
dataset = FilesDataset(file_list, cfg.test_pipeline)
data_loader = build_dataloader(dataset,
imgs_per_gpu=batch_size,
workers_per_gpu=batch_size,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
model = reweight_cls(model, args.tau)
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, False, cfg)
else:
model = MMDistributedDataParallel(model.cuda())
outputs = multi_gpu_test(model, data_loader, args.tmpdir)
print(type(outputs))
print(len(outputs))
print(len(outputs[0]))
exit()
# save outputs
for file_path, mmdet_ret in zip(file_list, outputs):
save_name = file_path.replace(img_dir, '')
save_path = os.path.join(out_dir, save_name)
common.makedirs(os.path.dirname(save_path))
save_in_tao_format(mmdet_ret, save_path)
def main():
args = parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpuid
img_dir = args.img_dir
out_dir = args.out_dir
batch_size = args.batch_size
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
if args.img_dir != '':
file_list = common.load_filepaths(args.img_dir,
suffix=('.jpg', '.png', '.jpeg'),
recursive=True)
elif args.img_list != '':
file_list = parse_testfile(args.img_list)
else:
raise "Both img_dir and img_list is empty."
dataset = FilesDataset(file_list, cfg.test_pipeline)
data_loader = build_dataloader(dataset,
imgs_per_gpu=batch_size,
workers_per_gpu=batch_size,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
model = reweight_cls(model, args.tau).cuda()
model = MMDataParallel(model, device_ids=[0])
model.eval()
count = 0
for i, data in enumerate(data_loader):
with torch.no_grad():
# bbox_results, segm_results
results = model(return_loss=False, rescale=True, **data)
# batch
#for result in results:
# file_path = file_list[count]
# save_name = file_path.replace('/home/songbai.xb/workspace/projects/TAO/data/TAO/frames/val/', '')
# save_path = os.path.join(out_dir, save_name)
# common.makedirs(os.path.dirname(save_path))
# save_in_tao_format(result, save_path)
# count += 1
file_path = file_list[i]
save_name = file_path.replace(
'/home/songbai.xb/workspace/projects/TAO/data/TAO/frames/val/', '')
save_name = save_name.replace('.jpg', '.pkl')
save_path = os.path.join(out_dir, save_name)
common.makedirs(os.path.dirname(save_path))
save_in_tao_format(results[0], save_path)
def main():
args = parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpuid
out_dir = args.out_dir
batch_size = args.batch_size
split = args.split
cfg = mmcv.Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
cfg.model.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
if args.img_file != '':
file_list = [args.img_file]
elif args.img_list != '':
file_list = parse_textfile(args.img_list)
else:
raise "Both img_file and img_list is empty."
if args.img_root != '':
file_list = [os.path.join(args.img_root, file) for file in file_list]
dataset = FilesDataset(file_list, cfg.test_pipeline)
data_loader = build_dataloader(
dataset,
imgs_per_gpu=batch_size,
workers_per_gpu=batch_size,
dist=distributed,
shuffle=False)
# build the model and load checkpoint
model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
model = reweight_cls(model, args.tau).cuda()
model = MMDataParallel(model, device_ids=[0])
model.eval()
count = 0
for i, data in enumerate(data_loader):
with torch.no_grad():
# bbox_results, segm_results
results = model(return_loss=False, rescale=True, **data)
file_path = file_list[i]
save_name = file_path.replace(f'{args.img_root}', '')
save_name = save_name.replace('.jpg', '.pkl').replace('.jpeg', '.pkl')
save_path = os.path.join(out_dir, save_name)
mmcv.mkdir_or_exist(os.path.dirname(save_path))
save_in_tao_format(results[0], save_path)