Lift, Splat, Shoot Encoding Images From Arbitrary Camera Rigs by Implicitly Unprojecting to 3D

class BaseTransform(nn.Module):

    def __init__(...):
        ...
        self.frustum = self.create_frustum()
        ...

    @force_fp32()
    def create_frustum(self):
        iH, iW = self.image_size
        fH, fW = self.feature_size

        ds = (
            torch.arange(*self.dbound, dtype=torch.float)  # 按起、终点及步长生成深度序列
            .view(-1, 1, 1)  # 在图像宽、高方向插入维度
            .expand(-1, fH, fW)  # 在宽、高维度广播数据（不复制仅引用）
        )
        D, _, _ = ds.shape

        xs = (
            torch.linspace(0, iW - 1, fW, dtype=torch.float)
            .view(1, 1, fW)
            .expand(D, fH, fW)
        )
        ys = (
            torch.linspace(0, iH - 1, fH, dtype=torch.float)
            .view(1, fH, 1)
            .expand(D, fH, fW)
        )

        frustum = torch.stack((xs, ys, ds), -1)  # 堆叠为 D x H x W x 3
        return nn.Parameter(frustum, requires_grad=False)

    @force_fp32()
    def get_geometry(
        self,
        camera2lidar_rots,
        camera2lidar_trans,
        intrins,
        post_rots,
        post_trans,
        **kwargs,
    ):
        B, N, _ = camera2lidar_trans.shape

        # undo post-transformation
        ## undo translation
        ## 此处存在维度广播：(D, H, W, 3) - (B, N, 1, 1, 1, 3) -> (B, N, D, H, W, 3)
        points = self.frustum - post_trans.view(B, N, 1, 1, 1, 3)
        ## undo rotation
        ## 使用 unsqueeze 在最后插入新维度以右对齐，再做矩阵乘法
        ## (B, N, 1, 1, 1, 3, 3) x (B, N, D, H, W, 3, 1) -> (B, N, D, H, W, 3, 1)
        points = (
            torch.inverse(post_rots)
            .view(B, N, 1, 1, 1, 3, 3)
            .matmul(points.unsqueeze(-1))
        )
        # cam_to_lidar
        # formula reference: https://review.yirami.xyz/academic/camera_calibration/#4-%E5%9D%90%E6%A0%87%E7%B3%BB%E8%BD%AC%E6%8D%A2
        points = torch.cat(
            (
                points[:, :, :, :, :, :2] * points[:, :, :, :, :, 2:3],
                points[:, :, :, :, :, 2:3],
            ),
            5,
        )  # (B, N, D, H, W, 3, 1)
        combine = camera2lidar_rots.matmul(torch.inverse(intrins))  # (B, N, 3, 3) x (B, N, 3, 3)
        points = combine.view(B, N, 1, 1, 1, 3, 3).matmul(points).squeeze(-1)  # (B, N, 1, 1, 1, 3, 3) x (B, N, D, H, W, 3, 1) -> (B, N, D, H, W, 3)
        points += camera2lidar_trans.view(B, N, 1, 1, 1, 3)  # (B, N, D, H, W, 3)

        # apply lidar augmentation transform
        ## rotation
        if "extra_rots" in kwargs:
            extra_rots = kwargs["extra_rots"]
            points = (
                extra_rots.view(B, 1, 1, 1, 1, 3, 3)
                .repeat(1, N, 1, 1, 1, 1, 1)
                .matmul(points.unsqueeze(-1))
                .squeeze(-1)
            )
        ## translation
        if "extra_trans" in kwargs:
            extra_trans = kwargs["extra_trans"]
            points += extra_trans.view(B, 1, 1, 1, 1, 3).repeat(1, N, 1, 1, 1, 1)

        return points  # (B, N, D, H, W, 3)

    @force_fp32()
    def forward(...):
        ...
        # rots = camera2ego[..., :3, :3]
        # trans = camera2ego[..., :3, 3]
        intrins = camera_intrinsics[..., :3, :3]
        post_rots = img_aug_matrix[..., :3, :3]
        post_trans = img_aug_matrix[..., :3, 3]
        # lidar2ego_rots = lidar2ego[..., :3, :3]
        # lidar2ego_trans = lidar2ego[..., :3, 3]
        camera2lidar_rots = camera2lidar[..., :3, :3]
        camera2lidar_trans = camera2lidar[..., :3, 3]

        extra_rots = lidar_aug_matrix[..., :3, :3]
        extra_trans = lidar_aug_matrix[..., :3, 3]

        geom = self.get_geometry(
            camera2lidar_rots,
            camera2lidar_trans,
            intrins,
            post_rots,
            post_trans,
            extra_rots=extra_rots,
            extra_trans=extra_trans,
        )  # (B, N, D, H, W, 3)

@VTRANSFORMS.register_module()
class LSSTransform(BaseTransform):

    def __init__(...):
        super().__init__(...)
        # 此处使用同一卷积网络将 D 和 C 合并输出，即
        # ... 分别取前 D 个通道为深度分布预测，后 C 个通道为特征预测，是为了
        # ... 让深度预测和特征预测进行共享和交互，使得
        # 1. 某个区域的特征有助于预测深度
        # 2. 深度信息反过来能改善特征表示
        # 此外，从使用上看，学习深度分布是为了**加权**特征，不是独立使用的估计
        self.depthnet = nn.Conv2d(in_channels, self.D + self.C, 1)
        ...

    @force_fp32()
    def get_cam_feats(self, x):
        B, N, C, fH, fW = x.shape

        # 将 B, N 看作一个维度，因为后续 nn.Conv2d() 的 `Tensor` 参数只接受 4 维
        #
        # nn.Conv2d() 在语义上是对输入 (B, C, H, W) 在
        # ... 1. 在空间维度 (H, W) 上应用相同的卷积核
        # ... 2. 在样本维度 (C, H, W) 上应用相同的参数
        # 此处将 B 和 N 看作一体，符合语义
        x = x.view(B * N, C, fH, fW)

        # nn.Module 中在 __call__() 中调用了 self.forward()
        x = self.depthnet(x)  # (B x N, D + C, H, W)
        # do softmax for depth prediction
        depth = x[:, : self.D].softmax(dim=1)  # (B x N, D, H, W)
        # calculate the weighted features
        x = depth.unsqueeze(1) * x[:, self.D : (self.D + self.C)].unsqueeze(2)  # (B x N, 1, D, H, W) x (B x N, C, 1, H, W) -> (B x N, C, D, H, W)

        x = x.view(B, N, self.C, self.D, fH, fW)  # (B, N, C, D, H, W)
        x = x.permute(0, 1, 3, 4, 5, 2)  # (B, N, D, H, W, C)
        return x

def gen_dx_bx(xbound, ybound, zbound):
    # delta
    dx = torch.Tensor([row[2] for row in [xbound, ybound, zbound]])
    # base ?: center of bound
    bx = torch.Tensor([row[0] + row[2] / 2.0 for row in [xbound, ybound, zbound]])
    # number
    nx = torch.LongTensor(
        [(row[1] - row[0]) / row[2] for row in [xbound, ybound, zbound]]
    )
    return dx, bx, nx

class QuickCumsum(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, geom_feats, ranks):
        # 在 Nprime 轴累积求和
        x = x.cumsum(0)
        # 标记各组最后一个元素的索引，并作为后续数据保留依据
        # 判断后一索引的 rank 与当前索引的 rank 相比是否跳变
        # ... 标记跳变处为 1 ，且最后一个值始终保留初始值 1
        kept = torch.ones(x.shape[0], device=x.device, dtype=torch.bool)
        kept[:-1] = ranks[1:] != ranks[:-1]

        # 按各组最后一个元素的索引取累积特征和相应坐标
        x, geom_feats = x[kept], geom_feats[kept]
        # 排除累积的前序组的特征
        x = torch.cat((x[:1], x[1:] - x[:-1]))

        # save kept for backward
        ctx.save_for_backward(kept)

        # no gradient for geom_feats
        ctx.mark_non_differentiable(geom_feats)

        return x, geom_feats

    @staticmethod
    def backward(ctx, gradx, gradgeom):  # PyTorch autograd 要求输入参数匹配 forward 的输出参数
        (kept,) = ctx.saved_tensors
        # 从跳变索引恢复出原始分组信息
        back = torch.cumsum(kept, 0)  # [0, 0, 1, 0, 1, 1] -> [0, 0, 1, 1, 2, 3]
        back[kept] -= 1  # [0, 0, 1, 1, 2, 3] -> [0, 0, 0, 1, 1, 2]

        # 因为 forward 中是组内求和，而 $\frac{\partial \text{sum}}{\partial x_i} = 1$
        # ... 所以 backward 中每个元素的梯度就是对应组的输出梯度
        val = gradx[back]

        return val, None, None  # 输出参数匹配 forward 的输入参数，不需要梯度的参数返回 None
        #      ↑x   ↑geom ↑ranks

def bev_pool(feats, coords, B, D, H, W):
    assert feats.shape[0] == coords.shape[0]

    # 将坐标（含批索引） $\in\mathbb{R}^4$ 映射到实数 $\mathbb{R}$ 空间, 
    # ... 且满足：1）同一网格的点映射后值相等 2）不同网格的点映射后值不等
    #
    # $x \in [0, H)$
    # $y \in [0, W)$
    # $z \in [0, D)$
    # $idx \in [0, B)$
    ranks = (
        coords[:, 0] * (W * D * B)  # x
        + coords[:, 1] * (D * B)  # y
        + coords[:, 2] * B  # z
        + coords[:, 3]  # batch idx
    )
    indices = ranks.argsort()
    feats, coords, ranks = feats[indices], coords[indices], ranks[indices]

    x, geom_feats = QuickCumsum.apply(feats, coords, ranks)

    # 创建空的特征网格数据
    C = x.shape[1]
    final = torch.zeros((B, C, D, H, W), device=x.device)
    # 将聚合后的特征按照索引赋到网格（高级索引赋值）
    # 此处切片作为左值是直接写回地址
    final[geom_feats[:, 3], :, geom_feats[:, 2], geom_feats[:, 0], geom_feats[:, 1]] = x

    return final  # (B, C, D, H, W)

class BaseTransform(nn.Module):

    def __init__(...):
        ...
        dx, bx, nx = gen_dx_bx(self.xbound, self.ybound, self.zbound)
        self.dx = nn.Parameter(dx, requires_grad=False)
        self.bx = nn.Parameter(bx, requires_grad=False)
        self.nx = nn.Parameter(nx, requires_grad=False)
        ...

    @force_fp32()
    def bev_pool(self, geom_feats, x):
        B, N, D, H, W, C = x.shape
        Nprime = B * N * D * H * W

        # flatten x
        x = x.reshape(Nprime, C)  # (B x N x D x H x W, C)

        # flatten indices
        ## 将 geom_feats 的三维坐标转到三维网格中得到索引，并使用 long() 将其长整化
        geom_feats = ((geom_feats - (self.bx - self.dx / 2.0)) / self.dx).long()
        geom_feats = geom_feats.view(Nprime, 3)  # (B x N x D x H x W, 3)
        ## 为各 Batch 的数据创建 Batch 索引
        batch_ix = torch.cat(
            [
                torch.full([Nprime // B, 1], ix, device=x.device, dtype=torch.long)
                for ix in range(B)
            ]
        )  # (B x N x D x H x W, 1)
        geom_feats = torch.cat((geom_feats, batch_ix), 1)  # (B x N x D x H x W, 4)

        # filter out points that are outside box
        kept = (
            (geom_feats[:, 0] >= 0)
            & (geom_feats[:, 0] < self.nx[0])
            & (geom_feats[:, 1] >= 0)
            & (geom_feats[:, 1] < self.nx[1])
            & (geom_feats[:, 2] >= 0)
            & (geom_feats[:, 2] < self.nx[2])
        )
        x = x[kept]
        geom_feats = geom_feats[kept]

        x = bev_pool(x, geom_feats, B, self.nx[2], self.nx[0], self.nx[1])  # (B, C, D, H, W)

        # collapse Z
        # 按 D 维展开，并在 C 维上拼接（同一深度的通道相邻）
        # 等价于 permute(0, 2, 1, 3, 4).contiguous().view(B, -1, H, W)
        # 等效于 view(B, -1, H, W) ，下游可自动适应任意通道排序
        final = torch.cat(x.unbind(dim=2), 1)

        return final

    @force_fp32()
    def forward(...):
        ...
        x = self.bev_pool(geom, x)
        ...