def test_indexnet(): model_cfg, _, test_cfg = _get_model_cfg( 'indexnet/indexnet_mobv2_1x16_78k_comp1k.py') model_cfg['pretrained'] = None # test indexnet inference with torch.no_grad(): indexnet = build_model(model_cfg, train_cfg=None, test_cfg=test_cfg) indexnet.eval() input_test = _demo_input_test((64, 64)) output_test = indexnet(**input_test, test_mode=True) assert isinstance(output_test['pred_alpha'], np.ndarray) assert output_test['pred_alpha'].shape == (64, 64) assert_dict_keys_equal(output_test['eval_result'], ['SAD', 'MSE', 'GRAD', 'CONN']) # test inference with gpu if torch.cuda.is_available(): indexnet = build_model( model_cfg, train_cfg=None, test_cfg=test_cfg).cuda() indexnet.eval() input_test = _demo_input_test((64, 64), cuda=True) output_test = indexnet(**input_test, test_mode=True) assert isinstance(output_test['pred_alpha'], np.ndarray) assert output_test['pred_alpha'].shape == (64, 64) assert_dict_keys_equal(output_test['eval_result'], ['SAD', 'MSE', 'GRAD', 'CONN']) # test forward train though we do not guarantee the training for present model_cfg.loss_alpha = None model_cfg.loss_comp = dict(type='L1CompositionLoss') indexnet = build_model( model_cfg, train_cfg=mmcv.ConfigDict(train_backbone=True), test_cfg=test_cfg) input_train = _demo_input_train((64, 64), batch_size=2) output_train = indexnet(**input_train) assert output_train['num_samples'] == 2 assert_dict_keys_equal(output_train['losses'], ['loss_comp']) if torch.cuda.is_available(): model_cfg.loss_alpha = dict(type='L1Loss') model_cfg.loss_comp = None indexnet = build_model( model_cfg, train_cfg=mmcv.ConfigDict(train_backbone=True), test_cfg=test_cfg).cuda() input_train = _demo_input_train((64, 64), batch_size=2, cuda=True) output_train = indexnet(**input_train) assert output_train['num_samples'] == 2 assert_dict_keys_equal(output_train['losses'], ['loss_alpha']) # test forward_dummy indexnet.cpu().eval() inputs = torch.ones((1, 4, 32, 32)) indexnet.forward_dummy(inputs)
def models_train(test_status=-10, data_type="fabric", const_weights=None): # train all models root = '../configs/{}/'.format(data_type) paths = glob.glob(os.path.join(root, 'defectnet/baseline_model.py')) paths.sort() for cfg_path in paths: m = BatchTrain(cfg_path=cfg_path, data_mode='test', train_sleep_time=0, test_sleep_time=test_status) m.common_train() # test stacking tricks paths = glob.glob(os.path.join(root, 'defectnet/*_mst_*_k*.py')) paths.sort() for cfg_path in paths: m = BatchTrain(cfg_path=cfg_path, data_mode='test', train_sleep_time=0, test_sleep_time=test_status) params = { 'test_cfg': { 'rcnn': mmcv.ConfigDict(score_thr=0.05, nms=dict(type='soft_nms', iou_thr=0.5), max_per_img=100) }, 'uid': "stacking tricks" } m.common_train(**params) # test different softnms thresholds iou_thrs = np.linspace(0.1, 0.9, 9) for iou_thr in iou_thrs: softnms_model = BatchTrain( cfg_path='../configs/{}/baseline_model.py'.format(data_type), data_mode='test', train_sleep_time=0, test_sleep_time=test_status) params = { 'test_cfg': { 'rcnn': mmcv.ConfigDict(score_thr=0.05, nms=dict(type='soft_nms', iou_thr=iou_thr), max_per_img=100) }, 'uid': iou_thr } softnms_model.common_train(**params)
def multi_scale_train(img_scale=None, multiscale_mode='range', ratio_range=None, keep_ratio=True): cfg_dir = '../config_alcohol/cascade_rcnn_r50_fpn_1x' cfg_names = [DATA_NAME + '.py', ] cfg = mmcv.Config.fromfile(os.path.join(cfg_dir, cfg_names[0])) cfg.train_pipeline[2] = mmcv.ConfigDict( type='Resize', img_scale=img_scale, ratio_range=ratio_range, multiscale_mode=multiscale_mode, keep_ratio=keep_ratio) sx = int(np.mean([v[0] for v in img_scale])) sy = int(np.mean([v[1] for v in img_scale])) cfg.test_pipeline[1]['img_scale'] = [(sx, sy)] cfg.data['train']['pipeline'] = cfg.train_pipeline cfg.data['val']['pipeline'] = cfg.test_pipeline cfg.data['test']['pipeline'] = cfg.test_pipeline cfg.data['imgs_per_gpu'] = 2 cfg.optimizer['lr'] = cfg.optimizer['lr'] / 8 * (cfg.data['imgs_per_gpu'] / 2) cfg.cfg_name = DATA_NAME + '_baseline' cfg.uid = 'img_scale={},multiscale_mode={},ratio_range={}' \ .format(str(img_scale), str(multiscale_mode), str(ratio_range)) cfg.work_dir = os.path.join(cfg.work_dir, cfg.cfg_name, cfg.cfg_name + ',' + cfg.uid) cfg.resume_from = os.path.join(cfg.work_dir, 'latest.pth') if not os.path.exists(cfg.resume_from): cfg.resume_from = None cfgs = [cfg] batch_train(cfgs, sleep_time=0 * 60 * 2) from batch_test import batch_test save_path = os.path.join(cfg_dir, DATA_NAME + '_test.txt') batch_test(cfgs, save_path, 60 * 2, mode=DATA_MODE)
def models_train(test_status=10, data_type="fabric", const_weights=None): # train all models root = '../configs/{}/'.format(data_type) paths = glob.glob(os.path.join(root, 'baseline_model_*.py')) paths.sort() for cfg_path in paths: m = BatchTrain(cfg_path=cfg_path, data_mode='test', train_sleep_time=0, test_sleep_time=test_status) m.common_train() # test stacking tricks paths = glob.glob(os.path.join(root, '*_mst_*k*.py')) paths.sort() for cfg_path in paths: m = BatchTrain(cfg_path=cfg_path, data_mode='test', train_sleep_time=0, test_sleep_time=test_status) params = { 'test_cfg': { 'rcnn': mmcv.ConfigDict(score_thr=0.05, nms=dict(type='soft_nms', iou_thr=0.5), max_per_img=100) }, 'uid': "stacking tricks" }
def test_object_sample(): db_sampler = mmcv.ConfigDict({ 'data_root': './tests/data/kitti/', 'info_path': './tests/data/kitti/kitti_dbinfos_train.pkl', 'rate': 1.0, 'prepare': { 'filter_by_difficulty': [-1], 'filter_by_min_points': { 'Pedestrian': 10 } }, 'classes': ['Pedestrian', 'Cyclist', 'Car'], 'sample_groups': { 'Pedestrian': 6 } }) object_sample = ObjectSample(db_sampler) points = np.fromfile( './tests/data/kitti/training/velodyne_reduced/000000.bin', np.float32).reshape(-1, 4) annos = mmcv.load('./tests/data/kitti/kitti_infos_train.pkl') info = annos[0] rect = info['calib']['R0_rect'].astype(np.float32) Trv2c = info['calib']['Tr_velo_to_cam'].astype(np.float32) annos = info['annos'] loc = annos['location'] dims = annos['dimensions'] rots = annos['rotation_y'] gt_names = annos['name'] gt_bboxes_3d = np.concatenate([loc, dims, rots[..., np.newaxis]], axis=1).astype(np.float32) gt_bboxes_3d = CameraInstance3DBoxes(gt_bboxes_3d).convert_to( Box3DMode.LIDAR, np.linalg.inv(rect @ Trv2c)) CLASSES = ('car', 'pedestrian', 'cyclist') gt_labels = [] for cat in gt_names: if cat in CLASSES: gt_labels.append(CLASSES.index(cat)) else: gt_labels.append(-1) input_dict = dict( points=points, gt_bboxes_3d=gt_bboxes_3d, gt_labels_3d=gt_labels) input_dict = object_sample(input_dict) points = input_dict['points'] gt_bboxes_3d = input_dict['gt_bboxes_3d'] gt_labels_3d = input_dict['gt_labels_3d'] expected_gt_bboxes_3d = torch.tensor( [[8.7314, -1.8559, -1.5997, 0.4800, 1.2000, 1.8900, 0.0100], [8.7314, -1.8559, -1.5997, 0.4800, 1.2000, 1.8900, 0.0100]]) expected_gt_labels_3d = np.array([-1, 0]) repr_str = repr(object_sample) expected_repr_str = 'ObjectSample' assert repr_str == expected_repr_str assert points.shape == (1177, 4) assert torch.allclose(gt_bboxes_3d.tensor, expected_gt_bboxes_3d) assert np.all(gt_labels_3d == expected_gt_labels_3d)
def test_base_mattor(): backbone = dict( type='SimpleEncoderDecoder', encoder=dict(type='VGG16', in_channels=4), decoder=dict(type='PlainDecoder')) refiner = dict(type='PlainRefiner') train_cfg = mmcv.ConfigDict(train_backbone=True, train_refiner=True) test_cfg = mmcv.ConfigDict( refine=True, metrics=['SAD', 'MSE', 'GRAD', 'CONN']) with pytest.raises(KeyError): # metrics should be specified in test_cfg BaseMattor( backbone, refiner, train_cfg.copy(), test_cfg=mmcv.ConfigDict(refine=True)) with pytest.raises(KeyError): # supported metric should be one of {'SAD', 'MSE'} BaseMattor( backbone, refiner, train_cfg.copy(), test_cfg=mmcv.ConfigDict( refine=True, metrics=['UnsupportedMetric'])) with pytest.raises(TypeError): # metrics must be None or a list of str BaseMattor( backbone, refiner, train_cfg.copy(), test_cfg=mmcv.ConfigDict(refine=True, metrics='SAD')) # build mattor without refiner mattor = BaseMattor( backbone, refiner=None, train_cfg=None, test_cfg=test_cfg.copy()) assert not mattor.with_refiner # only train the refiner, this will freeze the backbone mattor = BaseMattor( backbone, refiner, train_cfg=mmcv.ConfigDict(train_backbone=False, train_refiner=True), test_cfg=test_cfg.copy()) assert not mattor.train_cfg.train_backbone assert mattor.train_cfg.train_refiner assert mattor.test_cfg.refine # only train the backbone while the refiner is used for inference but not # trained, this behavior is allowed currently but will cause a warning. mattor = BaseMattor( backbone, refiner, train_cfg=mmcv.ConfigDict(train_backbone=True, train_refiner=False), test_cfg=test_cfg.copy()) assert mattor.train_cfg.train_backbone assert not mattor.train_cfg.train_refiner assert mattor.test_cfg.refine
def test_object_sample(): import pickle db_sampler = mmcv.ConfigDict({ 'data_root': './tests/data/kitti/', 'info_path': './tests/data/kitti/kitti_dbinfos_train.pkl', 'rate': 1.0, 'prepare': { 'filter_by_difficulty': [-1], 'filter_by_min_points': { 'Pedestrian': 10 } }, 'classes': ['Pedestrian', 'Cyclist', 'Car'], 'sample_groups': { 'Pedestrian': 6 } }) with open('./tests/data/kitti/kitti_dbinfos_train.pkl', 'rb') as f: db_infos = pickle.load(f) np.random.seed(0) object_sample = ObjectSample(db_sampler) points = np.fromfile( './tests/data/kitti/training/velodyne_reduced/000000.bin', np.float32).reshape(-1, 4) annos = mmcv.load('./tests/data/kitti/kitti_infos_train.pkl') info = annos[0] annos = info['annos'] gt_names = annos['name'] gt_bboxes_3d = db_infos['Pedestrian'][0]['box3d_lidar'] gt_bboxes_3d = LiDARInstance3DBoxes([gt_bboxes_3d]) CLASSES = ('Car', 'Pedestrian', 'Cyclist') gt_labels = [] for cat in gt_names: if cat in CLASSES: gt_labels.append(CLASSES.index(cat)) else: gt_labels.append(-1) input_dict = dict(points=points, gt_bboxes_3d=gt_bboxes_3d, gt_labels_3d=gt_labels) input_dict = object_sample(input_dict) points = input_dict['points'] gt_bboxes_3d = input_dict['gt_bboxes_3d'] gt_labels_3d = input_dict['gt_labels_3d'] repr_str = repr(object_sample) expected_repr_str = 'ObjectSample sample_2d=False, ' \ 'data_root=./tests/data/kitti/, ' \ 'info_path=./tests/data/kitti/kitti' \ '_dbinfos_train.pkl, rate=1.0, ' \ 'prepare={\'filter_by_difficulty\': [-1], ' \ '\'filter_by_min_points\': {\'Pedestrian\': 10}}, ' \ 'classes=[\'Pedestrian\', \'Cyclist\', \'Car\'], ' \ 'sample_groups={\'Pedestrian\': 6}' assert repr_str == expected_repr_str assert points.shape == (1177, 4) assert gt_bboxes_3d.tensor.shape == (2, 7) assert np.all(gt_labels_3d == [1, 0])
def joint_train(): cfg_dir = '../config_alcohol/cascade_rcnn_r50_fpn_1x' cfg_names = [DATA_NAME + '.py', ] cfg = mmcv.Config.fromfile(os.path.join(cfg_dir, cfg_names[0])) # 0.810 (img_scale = [1920, 1080]) ==> (img_scale = [(1920, 1080), (1333, 800)], multiscale_mode='range') img_scale = [(1920, 1080), (1333, 800)] cfg.train_pipeline[2] = mmcv.ConfigDict( type='Resize', img_scale=img_scale, ratio_range=None, multiscale_mode='range', keep_ratio=True) sx = int(max([v[0] for v in img_scale])) sy = int(max([v[1] for v in img_scale])) cfg.test_pipeline[1]['img_scale'] = [(sx, sy)] cfg.data['train'] = cfg.train_pipeline cfg.data['val'] = cfg.test_pipeline cfg.data['test'] = cfg.test_pipeline # 0.819 cfg.model['backbone']['dcn'] = dict( # 在最后三个block加入可变形卷积 modulated=False, deformable_groups=1, fallback_on_stride=False) cfg.model['backbone']['stage_with_dcn'] = (False, True, True, True) # 0.822 # global context cfg.model['bbox_roi_extractor']['global_context'] = True # # 0.746 # from tricks.data_cluster import anchor_cluster # anchor_ratios = anchor_cluster(cfg.data['train']['ann_file'], n=6) # cfg.model['rpn_head']['anchor_ratios'] = list(anchor_ratios) # 0.??? # focal loss for rcnn # for head in cfg.model['bbox_head']: # head['loss_cls'] = dict(type='FocalLoss', use_sigmoid=True, loss_weight=1.0) cfg.data['imgs_per_gpu'] = 2 cfg.optimizer['lr'] = cfg.optimizer['lr'] / 8 * (cfg.data['imgs_per_gpu'] / 2) cfg.cfg_name = DATA_NAME + '_baseline' cfg.uid = 'mode=joint_train+multiscale_mode=range' cfg.work_dir = os.path.join(cfg.work_dir, cfg.cfg_name, cfg.cfg_name + ',' + cfg.uid) cfg.resume_from = os.path.join(cfg.work_dir, 'latest.pth') if not os.path.exists(cfg.resume_from): cfg.resume_from = None cfgs = [cfg] batch_train(cfgs, sleep_time=0 * 60 * 2) from batch_test import batch_test save_path = os.path.join(cfg_dir, DATA_NAME + '_test.txt') batch_test(cfgs, save_path, 60 * 2, mode=DATA_MODE) from batch_train import batch_infer batch_infer(cfgs)
def test_merge_aug_bboxes_3d(): if not torch.cuda.is_available(): pytest.skip('test requires GPU and torch+cuda') img_meta_0 = dict( pcd_horizontal_flip=False, pcd_vertical_flip=True, pcd_scale_factor=1.0) img_meta_1 = dict( pcd_horizontal_flip=True, pcd_vertical_flip=False, pcd_scale_factor=1.0) img_meta_2 = dict( pcd_horizontal_flip=False, pcd_vertical_flip=False, pcd_scale_factor=0.5) img_metas = [[img_meta_0], [img_meta_1], [img_meta_2]] boxes_3d = DepthInstance3DBoxes( torch.tensor( [[1.0473, 4.1687, -1.2317, 2.3021, 1.8876, 1.9696, 1.6956], [2.5831, 4.8117, -1.2733, 0.5852, 0.8832, 0.9733, 1.6500], [-1.0864, 1.9045, -1.2000, 0.7128, 1.5631, 2.1045, 0.1022]], device='cuda')) labels_3d = torch.tensor([0, 7, 6]) scores_3d = torch.tensor([0.5, 1.0, 1.0]) aug_result = dict( boxes_3d=boxes_3d, labels_3d=labels_3d, scores_3d=scores_3d) aug_results = [aug_result, aug_result, aug_result] test_cfg = mmcv.ConfigDict( use_rotate_nms=True, nms_across_levels=False, nms_thr=0.01, score_thr=0.1, min_bbox_size=0, nms_pre=100, max_num=50) results = merge_aug_bboxes_3d(aug_results, img_metas, test_cfg) expected_boxes_3d = torch.tensor( [[-1.0864, -1.9045, -1.2000, 0.7128, 1.5631, 2.1045, -0.1022], [1.0864, 1.9045, -1.2000, 0.7128, 1.5631, 2.1045, 3.0394], [-2.1728, 3.8090, -2.4000, 1.4256, 3.1262, 4.2090, 0.1022], [2.5831, -4.8117, -1.2733, 0.5852, 0.8832, 0.9733, -1.6500], [-2.5831, 4.8117, -1.2733, 0.5852, 0.8832, 0.9733, 1.4916], [5.1662, 9.6234, -2.5466, 1.1704, 1.7664, 1.9466, 1.6500], [1.0473, -4.1687, -1.2317, 2.3021, 1.8876, 1.9696, -1.6956], [-1.0473, 4.1687, -1.2317, 2.3021, 1.8876, 1.9696, 1.4460], [2.0946, 8.3374, -2.4634, 4.6042, 3.7752, 3.9392, 1.6956]]) expected_scores_3d = torch.tensor([ 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.5000, 0.5000, 0.5000 ]) expected_labels_3d = torch.tensor([6, 6, 6, 7, 7, 7, 0, 0, 0]) assert torch.allclose(results['boxes_3d'].tensor, expected_boxes_3d) assert torch.allclose(results['scores_3d'], expected_scores_3d) assert torch.all(results['labels_3d'] == expected_labels_3d)
def init_detector(config, checkpoint=None, device='cuda:0'): """Initialize a detector from config file. Args: config (str or :obj:`mmcv.Config`): Config file path or the config object. checkpoint (str, optional): Checkpoint path. If left as None, the model will not load any weights. Returns: nn.Module: The constructed detector. """ if isinstance(config, str): config = mmcv.Config.fromfile(config) elif isinstance(config.model, dict): config.model = mmcv.ConfigDict(config.model) config.test_cfg = mmcv.ConfigDict(config.test_cfg) elif not isinstance(config, mmcv.Config): raise TypeError('config must be a filename or Config object, ' f'but got {type(config)}') config.model.pretrained = None model = build_detector(config.model, test_cfg=config.test_cfg) if checkpoint is not None: map_loc = 'cpu' if device == 'cpu' else None checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc) if 'CLASSES' in checkpoint['meta']: model.CLASSES = checkpoint['meta']['CLASSES'] else: warnings.simplefilter('once') warnings.warn('Class names are not saved in the checkpoint\'s ' 'meta data, use COCO classes by default.') model.CLASSES = get_classes('coco') model.cfg = config # save the config in the model for convenience model.to(device) model.eval() return model
def __init__(self, mc_cfg): cfg = mc_cfg # Create absolute paths for reid cfg.feature_extractor.reid_strong_extractor.reid_strong_baseline_config = os.path.join( cfg.repository_root, cfg.feature_extractor.reid_strong_extractor. reid_strong_baseline_config) cfg.feature_extractor.reid_strong_extractor.checkpoint_file = os.path.join( cfg.repository_root, cfg.feature_extractor.reid_strong_extractor.checkpoint_file) cfg.feature_extractor.abd_net_extractor.load_weights = os.path.join( cfg.repository_root, cfg.feature_extractor.abd_net_extractor.load_weights) feature_extractor_cfg = mmcv.ConfigDict(cfg.feature_extractor) if feature_extractor_cfg.feature_extractor_name == "reid_strong_extractor": self.feature_extractor = Reid_strong_extractor( feature_extractor_cfg) if feature_extractor_cfg.feature_extractor_name == "abd_net_extractor": self.feature_extractor = Abd_net_extractor( feature_extractor_cfg.abd_net_extractor)
def test_object_sample(): db_sampler = mmcv.ConfigDict({ 'data_root': './tests/data/kitti/', 'info_path': './tests/data/kitti/kitti_dbinfos_train.pkl', 'rate': 1.0, 'prepare': { 'filter_by_difficulty': [-1], 'filter_by_min_points': { 'Pedestrian': 10 } }, 'classes': ['Pedestrian', 'Cyclist', 'Car'], 'sample_groups': { 'Pedestrian': 6 } }) np.random.seed(0) object_sample = ObjectSample(db_sampler) points = np.fromfile( './tests/data/kitti/training/velodyne_reduced/000000.bin', np.float32).reshape(-1, 4) annos = mmcv.load('./tests/data/kitti/kitti_infos_train.pkl') info = annos[0] rect = info['calib']['R0_rect'].astype(np.float32) Trv2c = info['calib']['Tr_velo_to_cam'].astype(np.float32) annos = info['annos'] loc = annos['location'] dims = annos['dimensions'] rots = annos['rotation_y'] gt_names = annos['name'] gt_bboxes_3d = np.concatenate([loc, dims, rots[..., np.newaxis]], axis=1).astype(np.float32) gt_bboxes_3d = CameraInstance3DBoxes(gt_bboxes_3d).convert_to( Box3DMode.LIDAR, np.linalg.inv(rect @ Trv2c)) CLASSES = ('Pedestrian', 'Cyclist', 'Car') gt_labels = [] for cat in gt_names: if cat in CLASSES: gt_labels.append(CLASSES.index(cat)) else: gt_labels.append(-1) gt_labels = np.array(gt_labels, dtype=np.long) points = LiDARPoints(points, points_dim=4) input_dict = dict(points=points, gt_bboxes_3d=gt_bboxes_3d, gt_labels_3d=gt_labels) input_dict = object_sample(input_dict) points = input_dict['points'] gt_bboxes_3d = input_dict['gt_bboxes_3d'] gt_labels_3d = input_dict['gt_labels_3d'] repr_str = repr(object_sample) expected_repr_str = 'ObjectSample sample_2d=False, ' \ 'data_root=./tests/data/kitti/, ' \ 'info_path=./tests/data/kitti/kitti' \ '_dbinfos_train.pkl, rate=1.0, ' \ 'prepare={\'filter_by_difficulty\': [-1], ' \ '\'filter_by_min_points\': {\'Pedestrian\': 10}}, ' \ 'classes=[\'Pedestrian\', \'Cyclist\', \'Car\'], ' \ 'sample_groups={\'Pedestrian\': 6}' assert repr_str == expected_repr_str assert points.tensor.numpy().shape == (800, 4) assert gt_bboxes_3d.tensor.shape == (1, 7) assert np.all(gt_labels_3d == [0])
def test_ssn_test(): test_cfg = mmcv.ConfigDict( dict(ssn=dict(sampler=dict(test_interval=6, batch_size=16), evaluater=dict(top_k=2000, nms=0.2, softmax_before_filter=True, cls_score_dict=None, cls_top_k=2)))) base_model_cfg = dict(type='SSN', backbone=dict(type='ResNet', pretrained=None, depth=18, norm_eval=True), spatial_type='avg', dropout_ratio=0.8, cls_head=dict(type='SSNHead', dropout_ratio=0., in_channels=512, num_classes=20, consensus=dict(type='STPPTest', stpp_stage=(1, 1, 1)), use_regression=True), test_cfg=test_cfg) maxpool_model_cfg = copy.deepcopy(base_model_cfg) maxpool_model_cfg['spatial_type'] = 'max' non_regression_cfg = copy.deepcopy(base_model_cfg) non_regression_cfg['cls_head']['use_regression'] = False non_regression_cfg['cls_head']['consensus']['use_regression'] = False tuple_stage_cfg = copy.deepcopy(base_model_cfg) tuple_stage_cfg['cls_head']['consensus']['stpp_stage'] = (1, (1, 2), 1) str_stage_cfg = copy.deepcopy(base_model_cfg) str_stage_cfg['cls_head']['consensus']['stpp_stage'] = ('error', ) imgs = torch.rand(1, 8, 3, 224, 224) relative_proposal_list = torch.Tensor([[[0.2500, 0.6250], [0.3750, 0.7500]]]) scale_factor_list = torch.Tensor([[[1.0000, 1.0000], [1.0000, 0.2661]]]) proposal_tick_list = torch.LongTensor([[[1, 2, 5, 7], [20, 30, 60, 80]]]) reg_norm_consts = torch.Tensor([[[-0.0603, 0.0325], [0.0752, 0.1596]]]) localizer_ssn = build_localizer(base_model_cfg) localizer_ssn_maxpool = build_localizer(maxpool_model_cfg) localizer_ssn_non_regression = build_localizer(non_regression_cfg) localizer_ssn_tuple_stage_cfg = build_localizer(tuple_stage_cfg) with pytest.raises(ValueError): build_localizer(str_stage_cfg) if torch.cuda.is_available(): localizer_ssn = localizer_ssn.cuda() localizer_ssn_maxpool = localizer_ssn_maxpool.cuda() localizer_ssn_non_regression = localizer_ssn_non_regression.cuda() localizer_ssn_tuple_stage_cfg = localizer_ssn_tuple_stage_cfg.cuda() imgs = imgs.cuda() relative_proposal_list = relative_proposal_list.cuda() scale_factor_list = scale_factor_list.cuda() proposal_tick_list = proposal_tick_list.cuda() reg_norm_consts = reg_norm_consts.cuda() with torch.no_grad(): # Test normal case localizer_ssn(imgs, relative_proposal_list=relative_proposal_list, scale_factor_list=scale_factor_list, proposal_tick_list=proposal_tick_list, reg_norm_consts=reg_norm_consts, return_loss=False) # Test SSN model with max spatial pooling localizer_ssn_maxpool(imgs, relative_proposal_list=relative_proposal_list, scale_factor_list=scale_factor_list, proposal_tick_list=proposal_tick_list, reg_norm_consts=reg_norm_consts, return_loss=False) # Test SSN model without regression localizer_ssn_non_regression( imgs, relative_proposal_list=relative_proposal_list, scale_factor_list=scale_factor_list, proposal_tick_list=proposal_tick_list, reg_norm_consts=reg_norm_consts, return_loss=False) # Test SSN model with tuple stage cfg. localizer_ssn_tuple_stage_cfg( imgs, relative_proposal_list=relative_proposal_list, scale_factor_list=scale_factor_list, proposal_tick_list=proposal_tick_list, reg_norm_consts=reg_norm_consts, return_loss=False)
def test_ssn_train(): train_cfg = mmcv.ConfigDict( dict(ssn=dict(assigner=dict(positive_iou_threshold=0.7, background_iou_threshold=0.01, incomplete_iou_threshold=0.3, background_coverage_threshold=0.02, incomplete_overlap_threshold=0.01), sampler=dict(num_per_video=8, positive_ratio=1, background_ratio=1, incomplete_ratio=6, add_gt_as_proposals=True), loss_weight=dict(comp_loss_weight=0.1, reg_loss_weight=0.1), debug=False))) base_model_cfg = dict(type='SSN', backbone=dict(type='ResNet', pretrained=None, depth=18, norm_eval=True), spatial_type='avg', dropout_ratio=0.8, loss_cls=dict(type='SSNLoss'), cls_head=dict(type='SSNHead', dropout_ratio=0., in_channels=512, num_classes=20, consensus=dict(type='STPPTrain', stpp_stage=(1, 1, 1), num_segments_list=(2, 5, 2)), use_regression=True), train_cfg=train_cfg) dropout_cfg = copy.deepcopy(base_model_cfg) dropout_cfg['dropout_ratio'] = 0 dropout_cfg['cls_head']['dropout_ratio'] = 0.5 non_regression_cfg = copy.deepcopy(base_model_cfg) non_regression_cfg['cls_head']['use_regression'] = False imgs = torch.rand(1, 8, 9, 3, 224, 224) proposal_scale_factor = torch.Tensor([[[1.0345, 1.0345], [1.0028, 0.0028], [1.0013, 1.0013], [1.0008, 1.0008], [0.3357, 1.0006], [1.0006, 1.0006], [0.0818, 1.0005], [1.0030, 1.0030]]]) proposal_type = torch.Tensor([[0, 1, 1, 1, 1, 1, 1, 2]]) proposal_labels = torch.LongTensor([[8, 8, 8, 8, 8, 8, 8, 0]]) reg_targets = torch.Tensor([[[0.2929, 0.2694], [0.0000, 0.0000], [0.0000, 0.0000], [0.0000, 0.0000], [0.0000, 0.0000], [0.0000, 0.0000], [0.0000, 0.0000], [0.0000, 0.0000]]]) localizer_ssn = build_localizer(base_model_cfg) localizer_ssn_dropout = build_localizer(dropout_cfg) localizer_ssn_non_regression = build_localizer(non_regression_cfg) if torch.cuda.is_available(): localizer_ssn = localizer_ssn.cuda() localizer_ssn_dropout = localizer_ssn_dropout.cuda() localizer_ssn_non_regression = localizer_ssn_non_regression.cuda() imgs = imgs.cuda() proposal_scale_factor = proposal_scale_factor.cuda() proposal_type = proposal_type.cuda() proposal_labels = proposal_labels.cuda() reg_targets = reg_targets.cuda() # Train normal case losses = localizer_ssn(imgs, proposal_scale_factor=proposal_scale_factor, proposal_type=proposal_type, proposal_labels=proposal_labels, reg_targets=reg_targets) assert isinstance(losses, dict) # Train SSN without dropout in model, with dropout in head losses = localizer_ssn_dropout(imgs, proposal_scale_factor=proposal_scale_factor, proposal_type=proposal_type, proposal_labels=proposal_labels, reg_targets=reg_targets) assert isinstance(losses, dict) # Train SSN model without regression losses = localizer_ssn_non_regression( imgs, proposal_scale_factor=proposal_scale_factor, proposal_type=proposal_type, proposal_labels=proposal_labels, reg_targets=reg_targets) assert isinstance(losses, dict)
def test_edvr_model(): model_cfg = dict( type='EDVR', generator=dict( type='EDVRNet', in_channels=3, out_channels=3, mid_channels=8, num_frames=5, deform_groups=2, num_blocks_extraction=1, num_blocks_reconstruction=1, center_frame_idx=2, with_tsa=False), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='sum'), ) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'EDVR' assert isinstance(restorer.generator, EDVRNet) assert isinstance(restorer.pixel_loss, L1Loss) # prepare data inputs = torch.rand(1, 5, 3, 8, 8) targets = torch.rand(1, 3, 32, 32) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() data_batch = {'lq': inputs.cuda(), 'gt': targets.cuda()} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())) } # train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 32, 32) # with TSA model_cfg['generator']['with_tsa'] = True with pytest.raises(KeyError): # In TSA mode, train_cfg must contain "tsa_iter" train_cfg = dict(other_conent='xxx') restorer = build_model( model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() outputs = restorer.train_step(data_batch, optimizer) train_cfg = None restorer = build_model( model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() outputs = restorer.train_step(data_batch, optimizer) train_cfg = mmcv.ConfigDict(tsa_iter=1) restorer = build_model( model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())) } # train without updating tsa module outputs = restorer.train_step(data_batch, optimizer) # train with updating tsa module outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 32, 32) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['lq']) assert torch.is_tensor(output) assert output.size() == (1, 3, 32, 32) # forward_test with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 32, 32) with torch.no_grad(): outputs = restorer(inputs.cuda(), test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 32, 32) # test with metric and save image if torch.cuda.is_available(): train_cfg = mmcv.ConfigDict(tsa_iter=1) test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'lq': inputs.cuda(), 'gt': targets.cuda(), 'meta': [{ 'gt_path': 'fake_path/fake_name.png', 'key': '000/00000000' }] } restorer = build_model( model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() with pytest.raises(AssertionError): # evaluation with metrics must have gt images restorer(lq=inputs.cuda(), test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer( **data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer( **data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) with pytest.raises(ValueError): # iteration should be number or None restorer( **data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100')
def test_tdan_model(): model_cfg = dict( type='TDAN', generator=dict(type='TDANNet', in_channels=3, mid_channels=64, out_channels=3, num_blocks_before_align=5, num_blocks_after_align=10), pixel_loss=dict(type='MSELoss', loss_weight=1.0, reduction='sum'), lq_pixel_loss=dict(type='MSELoss', loss_weight=1.0, reduction='sum'), ) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'TDAN' assert isinstance(restorer.generator, TDANNet) assert isinstance(restorer.pixel_loss, MSELoss) # prepare data inputs = torch.rand(1, 5, 3, 8, 8) targets = torch.rand(1, 3, 32, 32) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() data_batch = {'lq': inputs.cuda(), 'gt': targets.cuda()} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())) } # train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 32, 32) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['lq']) assert isinstance(output, tuple) assert torch.is_tensor(output[0]) assert output[0].size() == (1, 3, 32, 32) assert torch.is_tensor(output[1]) assert output[1].size() == (1, 5, 3, 8, 8) # forward_test with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 32, 32) with torch.no_grad(): outputs = restorer(inputs.cuda(), test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 32, 32) # test with metric and save image if torch.cuda.is_available(): train_cfg = mmcv.ConfigDict(tsa_iter=1) test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'lq': inputs.cuda(), 'gt': targets.cuda(), 'meta': [{ 'gt_path': 'fake_path/fake_name.png', 'key': '000/00000000' }] } restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() with pytest.raises(AssertionError): # evaluation with metrics must have gt images restorer(lq=inputs.cuda(), test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) with pytest.raises(ValueError): # iteration should be number or None restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100')
def test_basicvsr_model(): model_cfg = dict( type='BasicVSR', generator=dict(type='BasicVSRNet', mid_channels=64, num_blocks=30, spynet_pretrained=None), pixel_loss=dict(type='MSELoss', loss_weight=1.0, reduction='sum'), ) train_cfg = dict(fix_iter=1) train_cfg = mmcv.Config(train_cfg) test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'BasicVSR' assert isinstance(restorer.generator, BasicVSRNet) assert isinstance(restorer.pixel_loss, MSELoss) # prepare data inputs = torch.rand(1, 5, 3, 64, 64) targets = torch.rand(1, 5, 3, 256, 256) if torch.cuda.is_available(): inputs = inputs.cuda() targets = targets.cuda() restorer = restorer.cuda() # prepare data and optimizer data_batch = {'lq': inputs, 'gt': targets} optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict(optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())) } # train_step (wihout updating spynet) outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 5, 3, 256, 256) # train with spynet updated outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 5, 3, 256, 256) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['lq']) assert torch.is_tensor(output) assert output.size() == (1, 5, 3, 256, 256) # forward_test with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 5, 3, 256, 256) with torch.no_grad(): outputs = restorer(inputs, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 5, 3, 256, 256) # test with metric and save image train_cfg = mmcv.ConfigDict(fix_iter=1) test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'lq': inputs, 'gt': targets, 'meta': [{ 'gt_path': 'fake_path/fake_name.png', 'key': '000' }] } restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) if torch.cuda.is_available(): restorer = restorer.cuda() with pytest.raises(AssertionError): # evaluation with metrics must have gt images restorer(lq=inputs, test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) with pytest.raises(ValueError): # iteration should be number or None restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100')
def main(): args = parse_args() if args.out is not None and not args.out.endswith(('.pkl', '.pickle')): raise ValueError('The output file must be a pkl file.') 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) dataset_type = 'OIDSegDataset' data_root = 'gs://oid2019/data/' img_norm_cfg = dict(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True) dataset = get_dataset( dict(type=dataset_type, ann_file='/home/bo_liu/' + args.ann_file, img_prefix=data_root + ('val/' if args.ann_file == 'seg_val_2844_ann.pkl' else 'OD_test/'), img_scale=(1333, 800), img_norm_cfg=img_norm_cfg, size_divisor=32, flip_ratio=0, with_mask=True, with_label=False, test_mode=True)) data_loader = build_dataloader(dataset, imgs_per_gpu=1, workers_per_gpu=cfg.data.workers_per_gpu, dist=distributed, shuffle=False) # build the model and load checkpoint test_cfg = mmcv.ConfigDict( dict( rpn=dict(nms_across_levels=False, nms_pre=1000, nms_post=1000, max_num=1000, nms_thr=0.7, min_bbox_size=0), rcnn=dict( score_thr=args.thres, # score_thr=0.0, nms=dict(type=args.nms_type, iou_thr=0.5), max_per_img=args.max_per_img, mask_thr_binary=0.5), keep_all_stages=False)) model = build_detector(cfg.model, train_cfg=None, test_cfg=test_cfg) checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu') # old versions did not save class info in checkpoints, this walkaround is # for backward compatibility if 'CLASSES' in checkpoint['meta']: model.CLASSES = checkpoint['meta']['CLASSES'] else: model.CLASSES = dataset.CLASSES if not distributed: model = MMDataParallel(model, device_ids=[0]) outputs = single_gpu_test(model, data_loader, args.show) else: model = MMDistributedDataParallel(model.cuda()) outputs = multi_gpu_test(model, data_loader, args.tmpdir) rank, _ = get_dist_info() if args.out and rank == 0: print('writing 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 coco_eval(result_file, eval_types, dataset.coco) else: if not isinstance(outputs[0], dict): result_file = args.out + '.json' results2json(dataset, outputs, result_file) coco_eval(result_file, eval_types, dataset.coco) else: for name in outputs[0]: print('Evaluating {}'.format(name)) outputs_ = [out[name] for out in outputs] result_file = args.out + '.{}.json'.format(name) results2json(dataset, outputs_, result_file) coco_eval(result_file, eval_types, dataset.coco)