<reponame>BB88Lee/mmdetection3d
import torch
from mmcv.cnn import build_norm_layer, build_conv_layer
from mmcv.runner import force_fp32
from torch import nn

#from mmdet3d.ops import DynamicScatter
import torch_scatter

#from .. import builder
#from ..registry import VOXEL_ENCODERS
#from .utils import VFELayer, get_paddings_indicator


import ipdb

#<EMAIL>()
class Cylinder_VFE(nn.Module):
    """Dynamic Voxel feature encoder used in DV-SECOND.

    It encodes features of voxels and their points. It could also fuse
    image feature into voxel features in a point-wise manner.
    The number of points inside the voxel varies.

    Args:
        in_channels (int): Input channels of VFE. Defaults to 4.
        feat_channels (list(int)): Channels of features in VFE.
        with_distance (bool): Whether to use the L2 distance of points to the
            origin point. Default False.
        with_cluster_center (bool): Whether to use the distance to cluster
            center of points inside a voxel. Default to False.
        with_voxel_center (bool): Whether to use the distance to center of
            voxel for each points inside a voxel. Default to False.
        voxel_size (tuple[float]): Size of a single voxel. Default to
            (0.2, 0.2, 4).
        point_cloud_range (tuple[float]): The range of points or voxels.
            Default to (0, -40, -3, 70.4, 40, 1).
        norm_cfg (dict): Config dict of normalization layers.
        mode (str): The mode when pooling features of points inside a voxel.
            Available options include 'max' and 'avg'. Default to 'max'.
        fusion_layer (dict | None): The config dict of fusion layer used in
            multi-modal detectors. Default to None.
        return_point_feats (bool): Whether to return the features of each
            points. Default to False.
    """

    def __init__(self,
                 in_channels=4,
                 feat_channels=[],
                 with_distance=False,
                 with_cluster_center=False,
                 with_voxel_center=False,
                 voxel_size=(0.2, 0.2, 4),
                 point_cloud_range=(0, -40, -3, 70.4, 40, 1),
                 conv_cfg=dict(type='Conv1d'),
                 norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
                 mode='max',
                 fusion_layer=None,
                 return_point_feats=False,
                 model_mode='voxel',
                 compress_dim=None):
        super(Cylinder_VFE, self).__init__()
        assert mode in ['avg', 'max']
        assert len(feat_channels) > 0
        if with_cluster_center:
            in_channels += 3
        if with_voxel_center:
            in_channels += 3
        if with_distance:
            in_channels += 3
        self.in_channels = in_channels
        self._with_distance = with_distance
        self._with_cluster_center = with_cluster_center
        self._with_voxel_center = with_voxel_center
        self.return_point_feats = return_point_feats
        self.fp16_enabled = False

        # Need pillar (voxel) size and x/y offset in order to calculate offset
        self.vx = voxel_size[0]
        self.vy = voxel_size[1]
        self.vz = voxel_size[2]
        self.x_offset = self.vx / 2 + point_cloud_range[0]
        self.y_offset = self.vy / 2 + point_cloud_range[1]
        self.z_offset = self.vz / 2 + point_cloud_range[2]
        self.point_cloud_range = point_cloud_range

        '''
        feat_channels = [self.in_channels] + list(feat_channels)
        vfe_layers = []
        for i in range(len(feat_channels) - 1):
            in_filters = feat_channels[i]
            out_filters = feat_channels[i + 1]
            if i > 0:
                in_filters *= 2
            norm_name, norm_layer = build_norm_layer(norm_cfg, out_filters)
            vfe_layers.append(
                nn.Sequential(
                    nn.Linear(in_filters, out_filters, bias=False), norm_layer,
                    nn.ReLU(inplace=True)))
            self.vfe_layers = nn.ModuleList(vfe_layers)
        '''

        build_conv_layer(conv_cfg,
                         in_channels, 64, 1)

        if model_mode == 'voxel':
            self.feat_model = nn.Sequential(
                build_norm_layer(norm_cfg, in_channels)[1],

                build_conv_layer(conv_cfg, in_channels, 64, 1, bias=False),
                build_norm_layer(norm_cfg, 64)[1],
                nn.ReLU(),

                build_conv_layer(conv_cfg, 64, 128, 1, bias=False),
                build_norm_layer(norm_cfg, 128)[1],
                nn.ReLU(),

                build_conv_layer(conv_cfg, 128, 256, 1, bias=False),
                build_norm_layer(norm_cfg, 256)[1],
                nn.ReLU(),

                build_conv_layer(conv_cfg, 256, 256, 1)
            )
        else:
            self.feat_model = nn.Sequential(
                build_conv_layer(conv_cfg, in_channels, 64, 1, bias=False),
                build_norm_layer(norm_cfg, 64)[1],
                nn.ReLU(),
            )

        self.compress_dim = compress_dim
        if compress_dim is not None:
            self.feat_compression = nn.Sequential(
                nn.Linear(256, compress_dim),
                nn.ReLU())


    def map_voxel_center_to_point(self, pts_coors, voxel_mean, voxel_coors):
        """Map voxel features to its corresponding points.

        Args:
            pts_coors (torch.Tensor): Voxel coordinate of each point.
            voxel_mean (torch.Tensor): Voxel features to be mapped.
            voxel_coors (torch.Tensor): Coordinates of valid voxels

        Returns:
            torch.Tensor: Features or centers of each point.
        """
        # Step 1: scatter voxel into canvas
        # Calculate necessary things for canvas creation
        canvas_z = int(
            (self.point_cloud_range[5] - self.point_cloud_range[2]) / self.vz)
        canvas_y = int(
            (self.point_cloud_range[4] - self.point_cloud_range[1]) / self.vy)
        canvas_x = int(
            (self.point_cloud_range[3] - self.point_cloud_range[0]) / self.vx)
        # canvas_channel = voxel_mean.size(1)
        batch_size = pts_coors[-1, 0] + 1
        canvas_len = canvas_z * canvas_y * canvas_x * batch_size
        # Create the canvas for this sample
        canvas = voxel_mean.new_zeros(canvas_len, dtype=torch.long)
        # Only include non-empty pillars
        indices = (
            voxel_coors[:, 0] * canvas_z * canvas_y * canvas_x +
            voxel_coors[:, 1] * canvas_y * canvas_x +
            voxel_coors[:, 2] * canvas_x + voxel_coors[:, 3])
        # Scatter the blob back to the canvas
        canvas[indices.long()] = torch.arange(
            start=0, end=voxel_mean.size(0), device=voxel_mean.device)

        # Step 2: get voxel mean for each point
        voxel_index = (
            pts_coors[:, 0] * canvas_z * canvas_y * canvas_x +
            pts_coors[:, 1] * canvas_y * canvas_x +
            pts_coors[:, 2] * canvas_x + pts_coors[:, 3])
        voxel_inds = canvas[voxel_index.long()]
        center_per_point = voxel_mean[voxel_inds, ...]
        return center_per_point

    @force_fp32(out_fp16=True)
    def forward(self,
                features,
                coors,
                points=None,
                img_feats=None,
                img_metas=None):
        """Forward functions.

        Args:
            features (torch.Tensor): Features of voxels, shape is NxC.
            coors (torch.Tensor): Coordinates of voxels, shape is  Nx(1+NDim).
            points (list[torch.Tensor], optional): Raw points used to guide the
                multi-modality fusion. Defaults to None.
            img_feats (list[torch.Tensor], optional): Image fetures used for
                multi-modality fusion. Defaults to None.
            img_metas (dict, optional): [description]. Defaults to None.

        Returns:
            tuple: If `return_point_feats` is False, returns voxel features and
                its coordinates. If `return_point_feats` is True, returns
                feature of each points inside voxels.
        """

        features_ls = [features]
        # Find distance of x, y, and z from cluster center
        if self._with_cluster_center:
            voxel_mean, mean_coors = self.cluster_scatter(features, coors)
            points_mean = self.map_voxel_center_to_point(
                coors, voxel_mean, mean_coors)
            # TODO: maybe also do cluster for reflectivity
            f_cluster = features[:, :3] - points_mean[:, :3]
            features_ls.append(f_cluster)

        # Find distance of x, y, and z from pillar center
        if self._with_voxel_center:
            f_center = features.new_zeros(size=(features.size(0), 3))
            f_center[:, 0] = features[:, 0] - (
                coors[:, 3].type_as(features) * self.vx + self.x_offset)
            f_center[:, 1] = features[:, 1] - (
                coors[:, 2].type_as(features) * self.vy + self.y_offset)
            f_center[:, 2] = features[:, 2] - (
                coors[:, 1].type_as(features) * self.vz + self.z_offset)
            features_ls.append(f_center)

        if self._with_distance:
            points_dist = torch.norm(features[:, :3], 2, 1, keepdim=True)
            features_ls.append(points_dist)

        # Combine together feature decorations
        features = torch.cat(features_ls, dim=-1)

        # trick for sync-bn 1D
        features = features.t().unsqueeze(0)

        point_feats = self.feat_model(features)

        point_feats = point_feats.squeeze(0).t()
        #voxel_feats, voxel_coors = self.vfe_scatter(point_feats, coors)

        # unique xy grid index
        unq, unq_inv, unq_cnt = torch.unique(coors, return_inverse=True, return_counts=True, dim=0)
        voxel_coors = unq
        voxel_coors_inv = unq_inv

        voxel_feats = torch_scatter.scatter_max(point_feats, unq_inv, dim=0)[0]

        if self.compress_dim is not None:
            voxel_feats = self.feat_compression(voxel_feats)

        return voxel_feats, voxel_coors, voxel_coors_inv
import torch
from mmcv.cnn import build_norm_layer, build_conv_layer
from mmcv.runner import force_fp32
from torch import nn

#from mmdet3d.ops import DynamicScatter
import torch_scatter

#from .. import builder
#from ..registry import VOXEL_ENCODERS
#from .utils import VFELayer, get_paddings_indicator


import ipdb

#<EMAIL>()
class Cylinder_VFE(nn.Module):
    """Dynamic Voxel feature encoder used in DV-SECOND.

    It encodes features of voxels and their points. It could also fuse
    image feature into voxel features in a point-wise manner.
    The number of points inside the voxel varies.

    Args:
        in_channels (int): Input channels of VFE. Defaults to 4.
        feat_channels (list(int)): Channels of features in VFE.
        with_distance (bool): Whether to use the L2 distance of points to the
            origin point. Default False.
        with_cluster_center (bool): Whether to use the distance to cluster
            center of points inside a voxel. Default to False.
        with_voxel_center (bool): Whether to use the distance to center of
            voxel for each points inside a voxel. Default to False.
        voxel_size (tuple[float]): Size of a single voxel. Default to
            (0.2, 0.2, 4).
        point_cloud_range (tuple[float]): The range of points or voxels.
            Default to (0, -40, -3, 70.4, 40, 1).
        norm_cfg (dict): Config dict of normalization layers.
        mode (str): The mode when pooling features of points inside a voxel.
            Available options include 'max' and 'avg'. Default to 'max'.
        fusion_layer (dict | None): The config dict of fusion layer used in
            multi-modal detectors. Default to None.
        return_point_feats (bool): Whether to return the features of each
            points. Default to False.
    """

    def __init__(self,
                 in_channels=4,
                 feat_channels=[],
                 with_distance=False,
                 with_cluster_center=False,
                 with_voxel_center=False,
                 voxel_size=(0.2, 0.2, 4),
                 point_cloud_range=(0, -40, -3, 70.4, 40, 1),
                 conv_cfg=dict(type='Conv1d'),
                 norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
                 mode='max',
                 fusion_layer=None,
                 return_point_feats=False,
                 model_mode='voxel',
                 compress_dim=None):
        super(Cylinder_VFE, self).__init__()
        assert mode in ['avg', 'max']
        assert len(feat_channels) > 0
        if with_cluster_center:
            in_channels += 3
        if with_voxel_center:
            in_channels += 3
        if with_distance:
            in_channels += 3
        self.in_channels = in_channels
        self._with_distance = with_distance
        self._with_cluster_center = with_cluster_center
        self._with_voxel_center = with_voxel_center
        self.return_point_feats = return_point_feats
        self.fp16_enabled = False

        # Need pillar (voxel) size and x/y offset in order to calculate offset
        self.vx = voxel_size[0]
        self.vy = voxel_size[1]
        self.vz = voxel_size[2]
        self.x_offset = self.vx / 2 + point_cloud_range[0]
        self.y_offset = self.vy / 2 + point_cloud_range[1]
        self.z_offset = self.vz / 2 + point_cloud_range[2]
        self.point_cloud_range = point_cloud_range

        '''
        feat_channels = [self.in_channels] + list(feat_channels)
        vfe_layers = []
        for i in range(len(feat_channels) - 1):
            in_filters = feat_channels[i]
            out_filters = feat_channels[i + 1]
            if i > 0:
                in_filters *= 2
            norm_name, norm_layer = build_norm_layer(norm_cfg, out_filters)
            vfe_layers.append(
                nn.Sequential(
                    nn.Linear(in_filters, out_filters, bias=False), norm_layer,
                    nn.ReLU(inplace=True)))
            self.vfe_layers = nn.ModuleList(vfe_layers)
        '''

        build_conv_layer(conv_cfg,
                         in_channels, 64, 1)

        if model_mode == 'voxel':
            self.feat_model = nn.Sequential(
                build_norm_layer(norm_cfg, in_channels)[1],

                build_conv_layer(conv_cfg, in_channels, 64, 1, bias=False),
                build_norm_layer(norm_cfg, 64)[1],
                nn.ReLU(),

                build_conv_layer(conv_cfg, 64, 128, 1, bias=False),
                build_norm_layer(norm_cfg, 128)[1],
                nn.ReLU(),

                build_conv_layer(conv_cfg, 128, 256, 1, bias=False),
                build_norm_layer(norm_cfg, 256)[1],
                nn.ReLU(),

                build_conv_layer(conv_cfg, 256, 256, 1)
            )
        else:
            self.feat_model = nn.Sequential(
                build_conv_layer(conv_cfg, in_channels, 64, 1, bias=False),
                build_norm_layer(norm_cfg, 64)[1],
                nn.ReLU(),
            )

        self.compress_dim = compress_dim
        if compress_dim is not None:
            self.feat_compression = nn.Sequential(
                nn.Linear(256, compress_dim),
                nn.ReLU())


    def map_voxel_center_to_point(self, pts_coors, voxel_mean, voxel_coors):
        """Map voxel features to its corresponding points.

        Args:
            pts_coors (torch.Tensor): Voxel coordinate of each point.
            voxel_mean (torch.Tensor): Voxel features to be mapped.
            voxel_coors (torch.Tensor): Coordinates of valid voxels

        Returns:
            torch.Tensor: Features or centers of each point.
        """
        # Step 1: scatter voxel into canvas
        # Calculate necessary things for canvas creation
        canvas_z = int(
            (self.point_cloud_range[5] - self.point_cloud_range[2]) / self.vz)
        canvas_y = int(
            (self.point_cloud_range[4] - self.point_cloud_range[1]) / self.vy)
        canvas_x = int(
            (self.point_cloud_range[3] - self.point_cloud_range[0]) / self.vx)
        # canvas_channel = voxel_mean.size(1)
        batch_size = pts_coors[-1, 0] + 1
        canvas_len = canvas_z * canvas_y * canvas_x * batch_size
        # Create the canvas for this sample
        canvas = voxel_mean.new_zeros(canvas_len, dtype=torch.long)
        # Only include non-empty pillars
        indices = (
            voxel_coors[:, 0] * canvas_z * canvas_y * canvas_x +
            voxel_coors[:, 1] * canvas_y * canvas_x +
            voxel_coors[:, 2] * canvas_x + voxel_coors[:, 3])
        # Scatter the blob back to the canvas
        canvas[indices.long()] = torch.arange(
            start=0, end=voxel_mean.size(0), device=voxel_mean.device)

        # Step 2: get voxel mean for each point
        voxel_index = (
            pts_coors[:, 0] * canvas_z * canvas_y * canvas_x +
            pts_coors[:, 1] * canvas_y * canvas_x +
            pts_coors[:, 2] * canvas_x + pts_coors[:, 3])
        voxel_inds = canvas[voxel_index.long()]
        center_per_point = voxel_mean[voxel_inds, ...]
        return center_per_point

    @force_fp32(out_fp16=True)
    def forward(self,
                features,
                coors,
                points=None,
                img_feats=None,
                img_metas=None):
        """Forward functions.

        Args:
            features (torch.Tensor): Features of voxels, shape is NxC.
            coors (torch.Tensor): Coordinates of voxels, shape is  Nx(1+NDim).
            points (list[torch.Tensor], optional): Raw points used to guide the
                multi-modality fusion. Defaults to None.
            img_feats (list[torch.Tensor], optional): Image fetures used for
                multi-modality fusion. Defaults to None.
            img_metas (dict, optional): [description]. Defaults to None.

        Returns:
            tuple: If `return_point_feats` is False, returns voxel features and
                its coordinates. If `return_point_feats` is True, returns
                feature of each points inside voxels.
        """

        features_ls = [features]
        # Find distance of x, y, and z from cluster center
        if self._with_cluster_center:
            voxel_mean, mean_coors = self.cluster_scatter(features, coors)
            points_mean = self.map_voxel_center_to_point(
                coors, voxel_mean, mean_coors)
            # TODO: maybe also do cluster for reflectivity
            f_cluster = features[:, :3] - points_mean[:, :3]
            features_ls.append(f_cluster)

        # Find distance of x, y, and z from pillar center
        if self._with_voxel_center:
            f_center = features.new_zeros(size=(features.size(0), 3))
            f_center[:, 0] = features[:, 0] - (
                coors[:, 3].type_as(features) * self.vx + self.x_offset)
            f_center[:, 1] = features[:, 1] - (
                coors[:, 2].type_as(features) * self.vy + self.y_offset)
            f_center[:, 2] = features[:, 2] - (
                coors[:, 1].type_as(features) * self.vz + self.z_offset)
            features_ls.append(f_center)

        if self._with_distance:
            points_dist = torch.norm(features[:, :3], 2, 1, keepdim=True)
            features_ls.append(points_dist)

        # Combine together feature decorations
        features = torch.cat(features_ls, dim=-1)

        # trick for sync-bn 1D
        features = features.t().unsqueeze(0)

        point_feats = self.feat_model(features)

        point_feats = point_feats.squeeze(0).t()
        #voxel_feats, voxel_coors = self.vfe_scatter(point_feats, coors)

        # unique xy grid index
        unq, unq_inv, unq_cnt = torch.unique(coors, return_inverse=True, return_counts=True, dim=0)
        voxel_coors = unq
        voxel_coors_inv = unq_inv

        voxel_feats = torch_scatter.scatter_max(point_feats, unq_inv, dim=0)[0]

        if self.compress_dim is not None:
            voxel_feats = self.feat_compression(voxel_feats)

        return voxel_feats, voxel_coors, voxel_coors_inv
#from mmdet3d.ops import DynamicScatter #from .. import builder #from ..registry import VOXEL_ENCODERS #from .utils import VFELayer, get_paddings_indicator #<EMAIL>() Dynamic Voxel feature encoder used in DV-SECOND.      It encodes features of voxels and their points. It could also fuse     image feature into voxel features in a point-wise manner.     The number of points inside the voxel varies.      Args:         in_channels (int): Input channels of VFE. Defaults to 4.         feat_channels (list(int)): Channels of features in VFE.         with_distance (bool): Whether to use the L2 distance of points to the             origin point. Default False.         with_cluster_center (bool): Whether to use the distance to cluster             center of points inside a voxel. Default to False.         with_voxel_center (bool): Whether to use the distance to center of             voxel for each points inside a voxel. Default to False.         voxel_size (tuple[float]): Size of a single voxel. Default to             (0.2, 0.2, 4).         point_cloud_range (tuple[float]): The range of points or voxels.             Default to (0, -40, -3, 70.4, 40, 1).         norm_cfg (dict): Config dict of normalization layers.         mode (str): The mode when pooling features of points inside a voxel.             Available options include 'max' and 'avg'. Default to 'max'.         fusion_layer (dict | None): The config dict of fusion layer used in             multi-modal detectors. Default to None.         return_point_feats (bool): Whether to return the features of each             points. Default to False. # Need pillar (voxel) size and x/y offset in order to calculate offset feat_channels = [self.in_channels] + list(feat_channels)         vfe_layers = []         for i in range(len(feat_channels) - 1):             in_filters = feat_channels[i]             out_filters = feat_channels[i + 1]             if i > 0:                 in_filters *= 2             norm_name, norm_layer = build_norm_layer(norm_cfg, out_filters)             vfe_layers.append(                 nn.Sequential(                     nn.Linear(in_filters, out_filters, bias=False), norm_layer,                     nn.ReLU(inplace=True)))             self.vfe_layers = nn.ModuleList(vfe_layers) Map voxel features to its corresponding points.          Args:             pts_coors (torch.Tensor): Voxel coordinate of each point.             voxel_mean (torch.Tensor): Voxel features to be mapped.             voxel_coors (torch.Tensor): Coordinates of valid voxels          Returns:             torch.Tensor: Features or centers of each point. # Step 1: scatter voxel into canvas # Calculate necessary things for canvas creation # canvas_channel = voxel_mean.size(1) # Create the canvas for this sample # Only include non-empty pillars # Scatter the blob back to the canvas # Step 2: get voxel mean for each point Forward functions.          Args:             features (torch.Tensor): Features of voxels, shape is NxC.             coors (torch.Tensor): Coordinates of voxels, shape is  Nx(1+NDim).             points (list[torch.Tensor], optional): Raw points used to guide the                 multi-modality fusion. Defaults to None.             img_feats (list[torch.Tensor], optional): Image fetures used for                 multi-modality fusion. Defaults to None.             img_metas (dict, optional): [description]. Defaults to None.          Returns:             tuple: If `return_point_feats` is False, returns voxel features and                 its coordinates. If `return_point_feats` is True, returns                 feature of each points inside voxels. # Find distance of x, y, and z from cluster center # TODO: maybe also do cluster for reflectivity # Find distance of x, y, and z from pillar center # Combine together feature decorations # trick for sync-bn 1D #voxel_feats, voxel_coors = self.vfe_scatter(point_feats, coors) # unique xy grid index