import torch import torch.nn as nn import torch.nn.functional as F from TriplaneGaussian.utils.base import BaseModule from TriplaneGaussian.utils.typing import * from dataclasses import dataclass, field from pytorch3d.renderer import ( AlphaCompositor, NormWeightedCompositor, PointsRasterizationSettings, PointsRasterizer, PointsRenderer) from pytorch3d.renderer.cameras import CamerasBase from pytorch3d.structures import Pointclouds from pytorch3d.utils.camera_conversions import cameras_from_opencv_projection from .utils import fps_subsample from einops import rearrange from .utils import MLP_CONV from .SPD import SPD from .SPD_crossattn import SPD_crossattn from .SPD_pp import SPD_pp SPD_BLOCK = { 'SPD': SPD, 'SPD_crossattn': SPD_crossattn, 'SPD_PP': SPD_pp, } def points_projection(points: Float[Tensor, "B Np 3"], c2ws: Float[Tensor, "B 4 4"], intrinsics: Float[Tensor, "B 3 3"], local_features: Float[Tensor, "B C H W"], raster_point_radius: float = 0.0075, # point size raster_points_per_pixel: int = 1, # a single point per pixel, for now bin_size: int = 0): """ points: (B, Np, 3) """ B, C, H, W = local_features.shape device = local_features.device raster_settings = PointsRasterizationSettings( image_size=(H, W), radius=raster_point_radius, points_per_pixel=raster_points_per_pixel, bin_size=bin_size, ) Np = points.shape[1] c2ws = c2ws.transpose(0, 1).flatten(0, 1) intrinsics = intrinsics.transpose(0, 1).flatten(0, 1) R = raster_settings.points_per_pixel w2cs = torch.inverse(c2ws) image_size = torch.as_tensor([H, W]).view( 1, 2).expand(w2cs.shape[0], -1).to(device) cameras = cameras_from_opencv_projection( w2cs[:, :3, :3], w2cs[:, :3, 3], intrinsics, image_size) rasterize = PointsRasterizer( cameras=cameras, raster_settings=raster_settings) fragments = rasterize(Pointclouds(points)) fragments_idx: Tensor = fragments.idx.long() visible_pixels = (fragments_idx > -1) # (B, H, W, R) points_to_visible_pixels = fragments_idx[visible_pixels] # Reshape local features to (B, H, W, R, C) local_features = local_features.permute( 0, 2, 3, 1).unsqueeze(-2).expand(-1, -1, -1, R, -1) # (B, H, W, R, C) # Get local features corresponding to visible points local_features_proj = torch.zeros(B * Np, C, device=device) local_features_proj[points_to_visible_pixels] = local_features[visible_pixels] local_features_proj = local_features_proj.reshape(B, Np, C) return local_features_proj class Decoder(nn.Module): def __init__(self, input_channels=1152, dim_feat=512, num_p0=512, radius=1, bounding=True, up_factors=None, SPD_type='SPD', token_type='image_token' ): super(Decoder, self).__init__() # self.decoder_coarse = SeedGenerator(dim_feat=dim_feat, num_pc=num_p0) if up_factors is None: up_factors = [1] else: up_factors = up_factors uppers = [] self.num_p0 = num_p0 self.mlp_feat_cond = MLP_CONV(in_channel=input_channels, layer_dims=[dim_feat*2, dim_feat]) for i, factor in enumerate(up_factors): uppers.append( SPD_BLOCK[SPD_type](dim_feat=dim_feat, up_factor=factor, i=i, bounding=bounding, radius=radius)) self.uppers = nn.ModuleList(uppers) self.token_type = token_type def calculate_pcl_token(self, pcl_token, up_factor): up_token = F.interpolate(pcl_token, scale_factor=up_factor, mode='nearest') return up_token def calculate_image_token(self, pcd, input_image_tokens, batch): """ Args: """ batch_size, n_input_views = batch["rgb_cond"].shape[:2] h_cond, w_cond = batch["rgb_cond"].shape[2:4] input_image_tokens = rearrange( input_image_tokens, '(B Nv) C Nt -> B (Nv Nt) C', Nv=n_input_views) local_features = input_image_tokens[:, 1:].reshape( batch_size, h_cond // 14, w_cond // 14, -1).permute(0, 3, 1, 2) local_features = F.interpolate(local_features, size=( h_cond, w_cond), mode='bilinear', align_corners=False) batch['c2w_cond'][..., :3, 1:3] *= -1 local_features_proj = points_projection( pcd, batch['c2w_cond'], batch['intrinsic_cond'], local_features, ) local_features_proj = local_features_proj.permute(0, 2, 1).contiguous() return local_features_proj def forward(self, x): """ Args: points: Tensor, (b, num_p0, 3) feat_cond: Tensor, (b, dim_feat) dinov2: 325x768 # partial_coarse: Tensor, (b, n_coarse, 3) """ points = x['points'] if self.token_type == 'pcl_token': feat_cond = x['pcl_token'] elif self.token_type == 'image_token': feat_cond = x['input_image_tokens'] feat_cond = self.mlp_feat_cond(feat_cond) arr_pcd = [] feat_prev = None pcd = torch.permute(points, (0, 2, 1)).contiguous() pcl_up_scale = 1 for upper in self.uppers: if self.token_type == 'pcl_token': up_cond = self.calculate_pcl_token( feat_cond, pcl_up_scale) pcl_up_scale *= upper.up_factor elif self.token_type == 'image_token': up_cond = self.calculate_image_token(points, feat_cond, x) pcd, feat_prev = upper(pcd, up_cond, feat_prev) points = torch.permute(pcd, (0, 2, 1)).contiguous() arr_pcd.append(points) return arr_pcd class SnowflakeModelSPDPP(BaseModule): """ apply PC^2 / PCL token to decoder """ @dataclass class Config(BaseModule.Config): input_channels: int = 1152 dim_feat: int = 128 num_p0: int = 512 radius: float = 1 bounding: bool = True use_fps: bool = True up_factors: List[int] = field(default_factory=lambda: [2, 2]) image_full_token_cond: bool = False SPD_type: str = 'SPD_PP' token_type: str = 'pcl_token' cfg: Config def configure(self) -> None: super().configure() self.decoder = Decoder(input_channels=self.cfg.input_channels, dim_feat=self.cfg.dim_feat, num_p0=self.cfg.num_p0, radius=self.cfg.radius, up_factors=self.cfg.up_factors, bounding=self.cfg.bounding, SPD_type=self.cfg.SPD_type, token_type=self.cfg.token_type ) def forward(self, x): results = self.decoder(x) return results