diff --git a/pcdet/models/detectors/detector3d_template.py b/pcdet/models/detectors/detector3d_template.py new file mode 100644 index 0000000..91e44bd --- /dev/null +++ b/pcdet/models/detectors/detector3d_template.py @@ -0,0 +1,415 @@ +import os + +import torch +import torch.nn as nn +import numpy as np +from ...ops.iou3d_nms import iou3d_nms_utils +from ...utils.spconv_utils import find_all_spconv_keys +from .. import backbones_2d, backbones_3d, dense_heads, roi_heads +from ..backbones_2d import map_to_bev +from ..backbones_3d import pfe, vfe +from ..model_utils import model_nms_utils + + +class Detector3DTemplate(nn.Module): + def __init__(self, model_cfg, num_class, dataset): + super().__init__() + self.model_cfg = model_cfg + self.num_class = num_class + self.dataset = dataset + self.class_names = dataset.class_names + self.register_buffer('global_step', torch.LongTensor(1).zero_()) + + self.module_topology = [ + 'vfe', 'backbone_3d', 'map_to_bev_module', 'pfe', + 'backbone_2d', 'dense_head', 'point_head', 'roi_head' + ] + + @property + def mode(self): + return 'TRAIN' if self.training else 'TEST' + + def update_global_step(self): + self.global_step += 1 + + def build_networks(self): + model_info_dict = { + 'module_list': [], + 'num_rawpoint_features': self.dataset.point_feature_encoder.num_point_features, + 'num_point_features': self.dataset.point_feature_encoder.num_point_features, + 'grid_size': self.dataset.grid_size, + 'point_cloud_range': self.dataset.point_cloud_range, + 'voxel_size': self.dataset.voxel_size, + 'depth_downsample_factor': self.dataset.depth_downsample_factor + } + for module_name in self.module_topology: + module, model_info_dict = getattr(self, 'build_%s' % module_name)( + model_info_dict=model_info_dict + ) + self.add_module(module_name, module) + return model_info_dict['module_list'] + + def build_vfe(self, model_info_dict): + if self.model_cfg.get('VFE', None) is None: + return None, model_info_dict + + vfe_module = vfe.__all__[self.model_cfg.VFE.NAME]( + model_cfg=self.model_cfg.VFE, + num_point_features=model_info_dict['num_rawpoint_features'], + point_cloud_range=model_info_dict['point_cloud_range'], + voxel_size=model_info_dict['voxel_size'], + grid_size=model_info_dict['grid_size'], + depth_downsample_factor=model_info_dict['depth_downsample_factor'] + ) + model_info_dict['num_point_features'] = vfe_module.get_output_feature_dim() + model_info_dict['module_list'].append(vfe_module) + return vfe_module, model_info_dict + + def build_backbone_3d(self, model_info_dict): + if self.model_cfg.get('BACKBONE_3D', None) is None: + return None, model_info_dict + + backbone_3d_module = backbones_3d.__all__[self.model_cfg.BACKBONE_3D.NAME]( + model_cfg=self.model_cfg.BACKBONE_3D, + input_channels=model_info_dict['num_point_features'], + grid_size=model_info_dict['grid_size'], + voxel_size=model_info_dict['voxel_size'], + point_cloud_range=model_info_dict['point_cloud_range'] + ) + model_info_dict['module_list'].append(backbone_3d_module) + model_info_dict['num_point_features'] = backbone_3d_module.num_point_features + model_info_dict['backbone_channels'] = backbone_3d_module.backbone_channels \ + if hasattr(backbone_3d_module, 'backbone_channels') else None + return backbone_3d_module, model_info_dict + + def build_map_to_bev_module(self, model_info_dict): + if self.model_cfg.get('MAP_TO_BEV', None) is None: + return None, model_info_dict + + map_to_bev_module = map_to_bev.__all__[self.model_cfg.MAP_TO_BEV.NAME]( + model_cfg=self.model_cfg.MAP_TO_BEV, + grid_size=model_info_dict['grid_size'] + ) + model_info_dict['module_list'].append(map_to_bev_module) + model_info_dict['num_bev_features'] = map_to_bev_module.num_bev_features + return map_to_bev_module, model_info_dict + + def build_backbone_2d(self, model_info_dict): + if self.model_cfg.get('BACKBONE_2D', None) is None: + return None, model_info_dict + + backbone_2d_module = backbones_2d.__all__[self.model_cfg.BACKBONE_2D.NAME]( + model_cfg=self.model_cfg.BACKBONE_2D, + input_channels=model_info_dict.get('num_bev_features', None) + ) + model_info_dict['module_list'].append(backbone_2d_module) + model_info_dict['num_bev_features'] = backbone_2d_module.num_bev_features + return backbone_2d_module, model_info_dict + + def build_pfe(self, model_info_dict): + if self.model_cfg.get('PFE', None) is None: + return None, model_info_dict + + pfe_module = pfe.__all__[self.model_cfg.PFE.NAME]( + model_cfg=self.model_cfg.PFE, + voxel_size=model_info_dict['voxel_size'], + point_cloud_range=model_info_dict['point_cloud_range'], + num_bev_features=model_info_dict['num_bev_features'], + num_rawpoint_features=model_info_dict['num_rawpoint_features'] + ) + model_info_dict['module_list'].append(pfe_module) + model_info_dict['num_point_features'] = pfe_module.num_point_features + model_info_dict['num_point_features_before_fusion'] = pfe_module.num_point_features_before_fusion + return pfe_module, model_info_dict + + def build_dense_head(self, model_info_dict): + if self.model_cfg.get('DENSE_HEAD', None) is None: + return None, model_info_dict + dense_head_module = dense_heads.__all__[self.model_cfg.DENSE_HEAD.NAME]( + model_cfg=self.model_cfg.DENSE_HEAD, + input_channels=model_info_dict['num_bev_features'] if 'num_bev_features' in model_info_dict else self.model_cfg.DENSE_HEAD.INPUT_FEATURES, + num_class=self.num_class if not self.model_cfg.DENSE_HEAD.CLASS_AGNOSTIC else 1, + class_names=self.class_names, + grid_size=model_info_dict['grid_size'], + point_cloud_range=model_info_dict['point_cloud_range'], + predict_boxes_when_training=self.model_cfg.get('ROI_HEAD', False), + voxel_size=model_info_dict.get('voxel_size', False) + ) + model_info_dict['module_list'].append(dense_head_module) + return dense_head_module, model_info_dict + + def build_point_head(self, model_info_dict): + if self.model_cfg.get('POINT_HEAD', None) is None: + return None, model_info_dict + + if self.model_cfg.POINT_HEAD.get('USE_POINT_FEATURES_BEFORE_FUSION', False): + num_point_features = model_info_dict['num_point_features_before_fusion'] + else: + num_point_features = model_info_dict['num_point_features'] + + point_head_module = dense_heads.__all__[self.model_cfg.POINT_HEAD.NAME]( + model_cfg=self.model_cfg.POINT_HEAD, + input_channels=num_point_features, + num_class=self.num_class if not self.model_cfg.POINT_HEAD.CLASS_AGNOSTIC else 1, + predict_boxes_when_training=self.model_cfg.get('ROI_HEAD', False) + ) + + model_info_dict['module_list'].append(point_head_module) + return point_head_module, model_info_dict + + def build_roi_head(self, model_info_dict): + if self.model_cfg.get('ROI_HEAD', None) is None: + return None, model_info_dict + point_head_module = roi_heads.__all__[self.model_cfg.ROI_HEAD.NAME]( + model_cfg=self.model_cfg.ROI_HEAD, + input_channels=model_info_dict['num_point_features'], + backbone_channels= model_info_dict.get('backbone_channels', None), + point_cloud_range=model_info_dict['point_cloud_range'], + voxel_size=model_info_dict['voxel_size'], + num_class=self.num_class if not self.model_cfg.ROI_HEAD.CLASS_AGNOSTIC else 1, + ) + + model_info_dict['module_list'].append(point_head_module) + return point_head_module, model_info_dict + + def forward(self, **kwargs): + raise NotImplementedError + + def post_processing(self, batch_dict): + """ + Args: + batch_dict: + batch_size: + batch_cls_preds: (B, num_boxes, num_classes | 1) or (N1+N2+..., num_classes | 1) + or [(B, num_boxes, num_class1), (B, num_boxes, num_class2) ...] + multihead_label_mapping: [(num_class1), (num_class2), ...] + batch_box_preds: (B, num_boxes, 7+C) or (N1+N2+..., 7+C) + cls_preds_normalized: indicate whether batch_cls_preds is normalized + batch_index: optional (N1+N2+...) + has_class_labels: True/False + roi_labels: (B, num_rois) 1 .. num_classes + batch_pred_labels: (B, num_boxes, 1) + Returns: + + """ + post_process_cfg = self.model_cfg.POST_PROCESSING + batch_size = batch_dict['batch_size'] + recall_dict = {} + pred_dicts = [] + for index in range(batch_size): + if batch_dict.get('batch_index', None) is not None: + assert batch_dict['batch_box_preds'].shape.__len__() == 2 + batch_mask = (batch_dict['batch_index'] == index) + else: + assert batch_dict['batch_box_preds'].shape.__len__() == 3 + batch_mask = index + + box_preds = batch_dict['batch_box_preds'][batch_mask] + src_box_preds = box_preds + + if not isinstance(batch_dict['batch_cls_preds'], list): + cls_preds = batch_dict['batch_cls_preds'][batch_mask] + + src_cls_preds = cls_preds + assert cls_preds.shape[1] in [1, self.num_class] + + if not batch_dict['cls_preds_normalized']: + cls_preds = torch.sigmoid(cls_preds) + else: + cls_preds = [x[batch_mask] for x in batch_dict['batch_cls_preds']] + src_cls_preds = cls_preds + if not batch_dict['cls_preds_normalized']: + cls_preds = [torch.sigmoid(x) for x in cls_preds] + + if post_process_cfg.NMS_CONFIG.MULTI_CLASSES_NMS: + if not isinstance(cls_preds, list): + cls_preds = [cls_preds] + multihead_label_mapping = [torch.arange(1, self.num_class, device=cls_preds[0].device)] + else: + multihead_label_mapping = batch_dict['multihead_label_mapping'] + + cur_start_idx = 0 + pred_scores, pred_labels, pred_boxes = [], [], [] + for cur_cls_preds, cur_label_mapping in zip(cls_preds, multihead_label_mapping): + assert cur_cls_preds.shape[1] == len(cur_label_mapping) + cur_box_preds = box_preds[cur_start_idx: cur_start_idx + cur_cls_preds.shape[0]] + cur_pred_scores, cur_pred_labels, cur_pred_boxes = model_nms_utils.multi_classes_nms( + cls_scores=cur_cls_preds, box_preds=cur_box_preds, + nms_config=post_process_cfg.NMS_CONFIG, + score_thresh=post_process_cfg.SCORE_THRESH + ) + cur_pred_labels = cur_label_mapping[cur_pred_labels] + pred_scores.append(cur_pred_scores) + pred_labels.append(cur_pred_labels) + pred_boxes.append(cur_pred_boxes) + cur_start_idx += cur_cls_preds.shape[0] + + final_scores = torch.cat(pred_scores, dim=0) + final_labels = torch.cat(pred_labels, dim=0) + final_boxes = torch.cat(pred_boxes, dim=0) + else: + cls_preds, label_preds = torch.max(cls_preds, dim=-1) + if batch_dict.get('has_class_labels', False): + label_key = 'roi_labels' if 'roi_labels' in batch_dict else 'batch_pred_labels' + label_preds = batch_dict[label_key][index] + else: + label_preds = label_preds + 1 + selected, selected_scores = model_nms_utils.class_agnostic_nms( + box_scores=cls_preds, box_preds=box_preds, + nms_config=post_process_cfg.NMS_CONFIG, + score_thresh=post_process_cfg.SCORE_THRESH + ) + + if post_process_cfg.OUTPUT_RAW_SCORE: + max_cls_preds, _ = torch.max(src_cls_preds, dim=-1) + selected_scores = max_cls_preds[selected] + + final_scores = selected_scores + final_labels = label_preds[selected] + final_boxes = box_preds[selected] + + recall_dict = self.generate_recall_record( + box_preds=final_boxes if 'rois' not in batch_dict else src_box_preds, + recall_dict=recall_dict, batch_index=index, data_dict=batch_dict, + thresh_list=post_process_cfg.RECALL_THRESH_LIST + ) + + record_dict = { + 'pred_boxes': final_boxes, + 'pred_scores': final_scores, + 'pred_labels': final_labels + } + pred_dicts.append(record_dict) + + return pred_dicts, recall_dict + + @staticmethod + def generate_recall_record(box_preds, recall_dict, batch_index, data_dict=None, thresh_list=None): + if 'gt_boxes' not in data_dict: + return recall_dict + + rois = data_dict['rois'][batch_index] if 'rois' in data_dict else None + gt_boxes = data_dict['gt_boxes'][batch_index] + + if recall_dict.__len__() == 0: + recall_dict = {'gt': 0} + for cur_thresh in thresh_list: + recall_dict['roi_%s' % (str(cur_thresh))] = 0 + recall_dict['rcnn_%s' % (str(cur_thresh))] = 0 + + cur_gt = gt_boxes + k = cur_gt.__len__() - 1 + while k >= 0 and cur_gt[k].sum() == 0: + k -= 1 + cur_gt = cur_gt[:k + 1] + + if cur_gt.shape[0] > 0: + if box_preds.shape[0] > 0: + iou3d_rcnn = iou3d_nms_utils.boxes_iou3d_gpu(box_preds[:, 0:7], cur_gt[:, 0:7]) + else: + iou3d_rcnn = torch.zeros((0, cur_gt.shape[0])) + + if rois is not None: + iou3d_roi = iou3d_nms_utils.boxes_iou3d_gpu(rois[:, 0:7], cur_gt[:, 0:7]) + + for cur_thresh in thresh_list: + if iou3d_rcnn.shape[0] == 0: + recall_dict['rcnn_%s' % str(cur_thresh)] += 0 + else: + rcnn_recalled = (iou3d_rcnn.max(dim=0)[0] > cur_thresh).sum().item() + recall_dict['rcnn_%s' % str(cur_thresh)] += rcnn_recalled + if rois is not None: + roi_recalled = (iou3d_roi.max(dim=0)[0] > cur_thresh).sum().item() + recall_dict['roi_%s' % str(cur_thresh)] += roi_recalled + + recall_dict['gt'] += cur_gt.shape[0] + else: + gt_iou = box_preds.new_zeros(box_preds.shape[0]) + return recall_dict + + def _load_state_dict(self, model_state_disk, *, strict=True): + state_dict = self.state_dict() # local cache of state_dict + + spconv_keys = find_all_spconv_keys(self) + + update_model_state = {} + for key, val in model_state_disk.items(): + if key in spconv_keys and key in state_dict and state_dict[key].shape != val.shape: + # with different spconv versions, we need to adapt weight shapes for spconv blocks + # adapt spconv weights from version 1.x to version 2.x if you used weights from spconv 1.x + + val_native = val.transpose(-1, -2) # (k1, k2, k3, c_in, c_out) to (k1, k2, k3, c_out, c_in) + if val_native.shape == state_dict[key].shape: + val = val_native.contiguous() + else: + assert val.shape.__len__() == 5, 'currently only spconv 3D is supported' + val_implicit = val.permute(4, 0, 1, 2, 3) # (k1, k2, k3, c_in, c_out) to (c_out, k1, k2, k3, c_in) + if val_implicit.shape == state_dict[key].shape: + val = val_implicit.contiguous() + + if key in state_dict and state_dict[key].shape == val.shape: + update_model_state[key] = val + # logger.info('Update weight %s: %s' % (key, str(val.shape))) + + if strict: + self.load_state_dict(update_model_state) + else: + state_dict.update(update_model_state) + self.load_state_dict(state_dict) + return state_dict, update_model_state + + def load_params_from_file(self, filename, logger, to_cpu=False, pre_trained_path=None): + if not os.path.isfile(filename): + raise FileNotFoundError + + logger.info('==> Loading parameters from checkpoint %s to %s' % (filename, 'CPU' if to_cpu else 'GPU')) + loc_type = torch.device('cpu') if to_cpu else None + checkpoint = torch.load(filename, map_location=loc_type) + model_state_disk = checkpoint['model_state'] + if not pre_trained_path is None: + pretrain_checkpoint = torch.load(pre_trained_path, map_location=loc_type) + pretrain_model_state_disk = pretrain_checkpoint['model_state'] + model_state_disk.update(pretrain_model_state_disk) + + version = checkpoint.get("version", None) + if version is not None: + logger.info('==> Checkpoint trained from version: %s' % version) + + state_dict, update_model_state = self._load_state_dict(model_state_disk, strict=False) + + for key in state_dict: + if key not in update_model_state: + logger.info('Not updated weight %s: %s' % (key, str(state_dict[key].shape))) + + logger.info('==> Done (loaded %d/%d)' % (len(update_model_state), len(state_dict))) + + def load_params_with_optimizer(self, filename, to_cpu=False, optimizer=None, logger=None): + if not os.path.isfile(filename): + raise FileNotFoundError + + logger.info('==> Loading parameters from checkpoint %s to %s' % (filename, 'CPU' if to_cpu else 'GPU')) + loc_type = torch.device('cpu') if to_cpu else None + checkpoint = torch.load(filename, map_location=loc_type) + epoch = checkpoint.get('epoch', -1) + it = checkpoint.get('it', 0.0) + + self._load_state_dict(checkpoint['model_state'], strict=True) + + if optimizer is not None: + if 'optimizer_state' in checkpoint and checkpoint['optimizer_state'] is not None: + logger.info('==> Loading optimizer parameters from checkpoint %s to %s' + % (filename, 'CPU' if to_cpu else 'GPU')) + optimizer.load_state_dict(checkpoint['optimizer_state']) + else: + assert filename[-4] == '.', filename + src_file, ext = filename[:-4], filename[-3:] + optimizer_filename = '%s_optim.%s' % (src_file, ext) + if os.path.exists(optimizer_filename): + optimizer_ckpt = torch.load(optimizer_filename, map_location=loc_type) + optimizer.load_state_dict(optimizer_ckpt['optimizer_state']) + + if 'version' in checkpoint: + print('==> Checkpoint trained from version: %s' % checkpoint['version']) + logger.info('==> Done') + + return it, epoch