188 lines
6.9 KiB
Python
188 lines
6.9 KiB
Python
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 |