def main():
    args = parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = str((args.data_index % 2))
    assert (args.out or args.show or args.json_out), 'Please specify at least one operation (save or show the results) with the argument "--out" or "--show" 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 = mmcv.Config.fromfile(args.config)
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    cfg.model.pretrained = None
    cfg.data.test.test_mode = True
    if (args.launcher == 'none'):
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
    dataset = build_dataset(cfg.data.test)
    data_loader = build_dataloader(dataset, imgs_per_gpu=1, workers_per_gpu=0, dist=distributed, shuffle=False)
    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')
    if ('CLASSES' in checkpoint['meta']):
        model.CLASSES = checkpoint['meta']['CLASSES']
    else:
        model.CLASSES = dataset.CLASSES
    model = reweight_cls(model, args.tau)
    if (not distributed):
        model = MMDataParallel(model, device_ids=[0])
        (outputs, logits) = single_gpu_test(model, data_loader, args.show, cfg, args.data_index)
    else:
        model = MMDistributedDataParallel(model.cuda())
        outputs = multi_gpu_test(model, data_loader, args.tmpdir)
    (bboxes_mat, labels_mat, logits_mat, proposal_num) = logits_process(logits)
    gt_list = []
    results_per_image = []
    for (i, data) in enumerate(data_loader):
        if (i >= TEMP_DATASET_SIZE):
            break
        print(i)
        img_id = dataset.img_infos[i]['id']
        gt = dataset.get_ann_info(i)
        gt_dict = dict()
        gt_dict['id'] = img_id
        gt_dict['bboxes'] = gt['bboxes']
        gt_dict['labels'] = gt['labels']
        gt_list.append(gt_dict)
        results = filter_logits_by_gt(bboxes_mat[i], logits_mat[i], gt_list[i], proposal_num[i], i)
        results_per_image.append(results)
    with open(f'dragon_bboxes_logits_map24.p', 'wb') as outfile:
        pickle.dump(results_per_image, outfile)
    print('saved')
    (rank, _) = get_dist_info()
    if (args.out and (rank == 0)):
        print('\nwriting results to {}'.format(args.out))
        mmcv.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
                lvis_eval(result_file, eval_types, dataset.lvis)
            elif (not isinstance(outputs[0], dict)):
                result_files = results2json(dataset, outputs, args.out, args.data_index)
                lvis_eval(result_files, eval_types, dataset.lvis, max_dets=300)
            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)
                    lvis_eval(result_files, eval_types, dataset.lvis)
    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 main():
    args = parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = str((args.data_index % 2))
    assert (args.out or args.show or args.json_out), 'Please specify at least one operation (save or show the results) with the argument "--out" or "--show" 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 = mmcv.Config.fromfile(args.config)
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    cfg.model.pretrained = None
    cfg.data.test.test_mode = True
    if (args.launcher == 'none'):
        distributed = False
    else:
        distributed = True
        init_dist(args.launcher, **cfg.dist_params)
    dataset = build_dataset(cfg.data.test)
    data_loader = build_dataloader(dataset, imgs_per_gpu=1, workers_per_gpu=0, dist=distributed, shuffle=False)
    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')
    if ('CLASSES' in checkpoint['meta']):
        model.CLASSES = checkpoint['meta']['CLASSES']
    else:
        model.CLASSES = dataset.CLASSES
    model = reweight_cls(model, args.tau)
    if (not distributed):
        model = MMDataParallel(model, device_ids=[0])
        (outputs, logits) = single_gpu_test(model, data_loader, args.show, cfg, args.data_index)
    else:
        model = MMDistributedDataParallel(model.cuda())
        outputs = multi_gpu_test(model, data_loader, args.tmpdir)
    (bboxes_mat, labels_mat, logits_mat, proposal_num) = logits_process(logits)
    gt_list = []
    results_per_image = []
    for (i, data) in enumerate(data_loader):
        if (i >= TEMP_DATASET_SIZE):
            break
        print(i)
        img_id = dataset.img_infos[i]['id']
        gt = dataset.get_ann_info(i)
        gt_dict = dict()
        gt_dict['id'] = img_id
        gt_dict['bboxes'] = gt['bboxes']
        gt_dict['labels'] = gt['labels']
        gt_list.append(gt_dict)
        results = filter_logits_by_gt(bboxes_mat[i], logits_mat[i], gt_list[i], proposal_num[i], i)
        results_per_image.append(results)
    with open(f'dragon_bboxes_logits_map24.p', 'wb') as outfile:
        pickle.dump(results_per_image, outfile)
    print('saved')
    (rank, _) = get_dist_info()
    if (args.out and (rank == 0)):
        print('\nwriting results to {}'.format(args.out))
        mmcv.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
                lvis_eval(result_file, eval_types, dataset.lvis)
            elif (not isinstance(outputs[0], dict)):
                result_files = results2json(dataset, outputs, args.out, args.data_index)
                lvis_eval(result_files, eval_types, dataset.lvis, max_dets=300)
            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)
                    lvis_eval(result_files, eval_types, dataset.lvis)
    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)