import logging import sys from dataclasses import dataclass, field from functools import partial from typing import List, Optional, Tuple import numpy as np import torch import torch.nn as nn from skimage import measure from torch.nn import functional as F from tqdm import tqdm from cube3d.model.autoencoder.embedder import PhaseModulatedFourierEmbedder from cube3d.model.autoencoder.grid import ( generate_dense_grid_points, marching_cubes_with_warp, ) from cube3d.model.autoencoder.spherical_vq import SphericalVectorQuantizer from cube3d.model.transformers.attention import ( EncoderCrossAttentionLayer, EncoderLayer, init_linear, init_tfixup, ) from cube3d.model.transformers.norm import LayerNorm def init_sort(x): """ Sorts the input tensor `x` based on its pairwise distances to the first element. This function computes the pairwise distances between all elements in `x` and the first element of `x`. It then sorts the elements of `x` in ascending order of their distances to the first element. Args: x (torch.Tensor): A 2D tensor where each row represents a data point. Returns: torch.Tensor: A tensor containing the rows of `x` sorted by their distances to the first row of `x`. """ distances = torch.cdist(x, x[:1]) _, indices = torch.sort(distances.squeeze(), dim=0) x = x[indices] return x class MLPEmbedder(nn.Module): def __init__(self, in_dim: int, embed_dim: int, bias: bool = True): super().__init__() self.in_layer = nn.Linear(in_dim, embed_dim, bias=bias) self.silu = nn.SiLU() self.out_layer = nn.Linear(embed_dim, embed_dim, bias=bias) self.apply(partial(init_linear, embed_dim=embed_dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.out_layer(self.silu(self.in_layer(x))) class OneDEncoder(nn.Module): def __init__( self, embedder, num_latents: int, point_feats: int, embed_point_feats: bool, width: int, num_heads: int, num_layers: int, with_cls_token: bool = False, cross_attention_levels: Optional[List[int]] = None, eps: float = 1e-6, ) -> None: """ Initializes the OneDEncoder model. Args: embedder: An embedding module that provides the input embedding functionality. num_latents (int): The number of latent variables. point_feats (int): The number of point features. embed_point_feats (bool): Whether to embed point features or not. width (int): The width of the embedding dimension. num_heads (int): The number of attention heads. num_layers (int): The number of encoder layers. with_cls_token (bool, optional): Whether to include a classification token like in Vision Transformers (ViT). Defaults to False. cross_attention_levels (Optional[List[int]], optional): The indices of layers where cross-attention is applied. Defaults to None. eps (float, optional): A small value added for numerical stability in normalization layers. Defaults to 1e-6. Returns: None """ super().__init__() self.embedder = embedder # add cls token like ViT self.with_cls_token = with_cls_token if self.with_cls_token: query = torch.empty((1 + num_latents, width)) else: query = torch.empty((num_latents, width)) # initialize then sort query to potentially get better ordering query.uniform_(-1.0, 1.0) query = init_sort(query) # set parameter self.query = nn.Parameter(query) self.embed_point_feats = embed_point_feats in_dim = ( self.embedder.out_dim * 2 if self.embed_point_feats else self.embedder.out_dim + point_feats ) self.feat_in = MLPEmbedder(in_dim, embed_dim=width) if cross_attention_levels is None: cross_attention_levels = [0] self.blocks = nn.ModuleList() for i in range(num_layers): if i in cross_attention_levels: self.blocks.append( EncoderCrossAttentionLayer( embed_dim=width, num_heads=num_heads, eps=eps, ) ) else: self.blocks.append( EncoderLayer(embed_dim=width, num_heads=num_heads, eps=eps) ) self.ln_f = LayerNorm(width, eps=eps) init_tfixup(self, num_layers) def _forward(self, h, data, attn_mask=None): """ Forward pass for the autoencoder model. Args: h (torch.Tensor): The input tensor to be processed, typically representing the hidden state or intermediate representation. data (torch.Tensor): The input data tensor to be transformed by the feature extraction layer and used in cross-attention layers. attn_mask (torch.Tensor, optional): An optional attention mask tensor to be used in attention layers for masking specific positions. Defaults to None. Returns: torch.Tensor: The output tensor after processing through the layers and applying final normalization. """ data = self.feat_in(data) for block in self.blocks: if isinstance(block, EncoderCrossAttentionLayer): h = block(h, data) else: h = block(h, attn_mask=attn_mask) h = self.ln_f(h) return h def forward( self, pts: torch.Tensor, feats: torch.Tensor ) -> Tuple[torch.Tensor, list[torch.Tensor]]: """ Forward pass of the 1D autoencoder model. Args: pts (torch.Tensor): Input tensor representing points with shape (batch_size, num_points, point_dim). feats (torch.Tensor): Input tensor representing features with shape (batch_size, num_points, feature_dim). Can be None if no features are provided. Returns: Tuple[torch.Tensor, list[torch.Tensor]]: - The output tensor after processing the input data. - A list of intermediate tensors (if applicable) generated during the forward pass. """ b = pts.shape[0] data = self.embedder(pts) if feats is not None: if self.embed_point_feats: feats = self.embedder(feats) data = torch.cat([data, feats], dim=-1) # prepare query and data h = self.query.unsqueeze(0).expand(b, -1, -1) return self._forward(h, data, attn_mask=None) class OneDBottleNeck(nn.Module): def __init__( self, block, ) -> None: """ Initializes the OneDBottleNeck class. Args: block: The building block or module used within the autoencoder. """ super().__init__() self.block = block def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, dict]: """ Forward pass of the OneDBottleNeck function. Args: h (torch.Tensor): Input tensor to the model. Returns: Tuple[torch.Tensor, dict]: A tuple containing: - The transformed tensor `z` after passing through the block (if applicable). - A dictionary `ret_dict` containing additional information: - "indices": Indices from the block output (if present). - "z_q": Quantized tensor from the block output (if present). """ z = h ret_dict = {} if self.block is not None: z, d = self.block(z) key_mappings = { "q": "indices", "z_q": "z_q", } for in_key, out_key in key_mappings.items(): if in_key in d: ret_dict[out_key] = d[in_key] return z, ret_dict class OneDDecoder(nn.Module): def __init__( self, num_latents: int, width: int, num_heads: int, num_layers: int, eps: float = 1e-6, ) -> None: """ Initializes the OneDDecoder class. Args: num_latents (int): The number of latent variables. width (int): The width of the embedding dimension. num_heads (int): The number of attention heads in each encoder layer. num_layers (int): The number of encoder layers. eps (float, optional): A small value added for numerical stability. Defaults to 1e-6. """ super().__init__() self.register_buffer("query", torch.empty([0, width]), persistent=False) self.positional_encodings = nn.Parameter( init_sort(F.normalize(torch.empty(num_latents, width).normal_())) ) self.blocks = nn.ModuleList( [ EncoderLayer(embed_dim=width, num_heads=num_heads, eps=eps) for _ in range(num_layers) ] ) init_tfixup(self, num_layers) def _forward(self, h): """ Applies a sequence of operations to the input tensor `h` using the blocks defined in the model. Args: h (torch.Tensor): The input tensor to be processed by the blocks. Returns: torch.Tensor: The output tensor after applying all blocks sequentially. """ for block in self.blocks: h = block(h) return h def forward(self, z): """ This method processes the input tensor `z` by padding it to a fixed length, adding positional encodings, and then passing it through the `_forward` method. Args: z (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor after processing through the autoencoder. Notes: - If the `query` attribute has a non-zero shape, the input tensor `z` is padded to match the required length using slices of `query`. - Positional encodings are added to the padded input tensor before passing it to the `_forward` method. """ # pad input to fixed length if self.query.shape[0] > 0: pad_len = self.query.shape[0] + 1 - z.shape[1] paddings = self.query[:pad_len, ...].unsqueeze(0).expand(z.shape[0], -1, -1) z = torch.cat([paddings, z], dim=1) h = z + self.positional_encodings[: z.shape[1], :].unsqueeze(0).expand( z.shape[0], -1, -1 ) return self._forward(h) class OneDOccupancyDecoder(nn.Module): def __init__( self, embedder, out_features: int, width: int, num_heads: int, eps=1e-6 ) -> None: """ Initializes the OneDOccupancyDecoder module. Args: embedder: An embedding module that provides input embeddings. out_features (int): The number of output features for the final linear layer. width (int): The width of the intermediate layers. num_heads (int): The number of attention heads for the cross-attention layer. eps (float, optional): A small value added for numerical stability in layer normalization. Defaults to 1e-6. """ super().__init__() self.embedder = embedder self.query_in = MLPEmbedder(self.embedder.out_dim, width) self.attn_out = EncoderCrossAttentionLayer(embed_dim=width, num_heads=num_heads) self.ln_f = LayerNorm(width, eps=eps) self.c_head = nn.Linear(width, out_features) def query(self, queries: torch.Tensor): """ Processes the input tensor through the embedder and query_in layers. Args: queries (torch.Tensor): A tensor containing the input data to be processed. Returns: torch.Tensor: The output tensor after being processed by the embedder and query_in layers. """ return self.query_in(self.embedder(queries)) def forward(self, queries: torch.Tensor, latents: torch.Tensor): """ Defines the forward pass of the model. Args: queries (torch.Tensor): Input tensor representing the queries. latents (torch.Tensor): Input tensor representing the latent representations. Returns: torch.Tensor: Output tensor after applying the query transformation, attention mechanism, and final processing layers. """ queries = self.query(queries) x = self.attn_out(queries, latents) x = self.c_head(self.ln_f(x)) return x class OneDAutoEncoder(nn.Module): @dataclass class Config: checkpoint_path: str = "" # network params num_encoder_latents: int = 256 num_decoder_latents: int = 256 embed_dim: int = 12 width: int = 768 num_heads: int = 12 out_dim: int = 1 eps: float = 1e-6 # grid features embedding num_freqs: int = 128 point_feats: int = 0 embed_point_feats: bool = False num_encoder_layers: int = 1 encoder_cross_attention_levels: list[int] = field(default_factory=list) num_decoder_layers: int = 23 encoder_with_cls_token: bool = True num_codes: int = 16384 def __init__(self, cfg: Config) -> None: """ Initializes the OneDAutoencoder model. Args: cfg (Config): Configuration object containing the parameters for the model. Attributes: cfg (Config): Stores the configuration object. embedder (PhaseModulatedFourierEmbedder): Embeds input data using phase-modulated Fourier features. encoder (OneDEncoder): Encodes the input data into latent representations. bottleneck (OneDBottleNeck): Bottleneck layer containing a spherical vector quantizer for dimensionality reduction. decoder (OneDDecoder): Decodes latent representations back into the original data space. occupancy_decoder (OneDOccupancyDecoder): Decodes occupancy information from latent representations. """ super().__init__() self.cfg = cfg self.embedder = PhaseModulatedFourierEmbedder( num_freqs=self.cfg.num_freqs, input_dim=3 ) self.encoder = OneDEncoder( embedder=self.embedder, num_latents=self.cfg.num_encoder_latents, with_cls_token=self.cfg.encoder_with_cls_token, point_feats=self.cfg.point_feats, embed_point_feats=self.cfg.embed_point_feats, width=self.cfg.width, num_heads=self.cfg.num_heads, num_layers=self.cfg.num_encoder_layers, cross_attention_levels=self.cfg.encoder_cross_attention_levels, eps=self.cfg.eps, ) block = SphericalVectorQuantizer( self.cfg.embed_dim, self.cfg.num_codes, self.cfg.width, codebook_regularization="kl", ) self.bottleneck = OneDBottleNeck(block=block) self.decoder = OneDDecoder( num_latents=self.cfg.num_encoder_latents, width=self.cfg.width, num_heads=self.cfg.num_heads, num_layers=self.cfg.num_decoder_layers, eps=self.cfg.eps, ) self.occupancy_decoder = OneDOccupancyDecoder( embedder=self.embedder, out_features=self.cfg.out_dim, width=self.cfg.width, num_heads=self.cfg.num_heads, eps=self.cfg.eps, ) @torch.no_grad() def decode_indices(self, shape_ids: torch.Tensor): """ Decodes the given shape indices into latent representations. Args: shape_ids (torch.Tensor): A tensor containing the shape indices to be decoded. Returns: torch.Tensor: The decoded latent representations corresponding to the input shape indices. """ z_q = self.bottleneck.block.lookup_codebook(shape_ids) latents = self.decode(z_q) return latents @torch.no_grad() def query_embeds(self, shape_ids: torch.Tensor): """ Retrieves the latent embeddings corresponding to the given shape IDs. Args: shape_ids (torch.Tensor): A tensor containing the IDs of the shapes for which the latent embeddings are to be queried. Returns: torch.Tensor: A tensor containing the latent embeddings retrieved from the codebook for the provided shape IDs. """ z_q = self.bottleneck.block.lookup_codebook_latents(shape_ids) return z_q @torch.no_grad() def query_indices(self, shape_embs: torch.Tensor): """ Queries the indices of the quantized embeddings from the bottleneck layer. Args: shape_embs (torch.Tensor): The input tensor containing shape embeddings to be quantized. Returns: torch.Tensor: A tensor containing the quantized indices. """ _, ret_dict = self.bottleneck.block.quantize(shape_embs) return ret_dict["q"] def encode(self, x: torch.Tensor, **kwargs): """ Encodes the input tensor using the encoder and bottleneck layers. Args: x (torch.Tensor): Input tensor with shape (..., N), where the first 3 dimensions represent points (pts) and the remaining dimensions represent features (feats). **kwargs: Additional keyword arguments. Returns: Tuple[torch.Tensor, torch.Tensor, None, dict]: A tuple containing: - z_e (torch.Tensor): Encoded tensor before bottleneck processing. - z (torch.Tensor): Encoded tensor after bottleneck processing. - None: Placeholder for compatibility with other methods. - d (dict): Dictionary containing additional information, including: - "z_cls" (torch.Tensor, optional): Class token if `self.cfg.encoder_with_cls_token` is True. """ pts, feats = x[..., :3], x[..., 3:] z_e = self.encoder(pts, feats) # split class token if self.cfg.encoder_with_cls_token: z_cls = z_e[:, 0, ...] z_e = z_e[:, 1:, ...] # quantize or kl z, d = self.bottleneck(z_e) if self.cfg.encoder_with_cls_token: d["z_cls"] = z_cls return z_e, z, None, d def decode(self, z: torch.Tensor): """ Decodes the latent representation `z` using the decoder network. Args: z (torch.Tensor): The latent representation tensor to be decoded. Returns: torch.Tensor: The decoded output tensor. """ h = self.decoder(z) return h def query(self, queries: torch.Tensor, latents: torch.Tensor): """ Computes the logits by decoding the given queries and latent representations. Args: queries (torch.Tensor): A tensor containing the query points to be decoded. latents (torch.Tensor): A tensor containing the latent representations corresponding to the queries. Returns: torch.Tensor: A tensor containing the decoded logits for the given queries and latents. """ logits = self.occupancy_decoder(queries, latents).squeeze(-1) return logits def forward(self, surface, queries, **kwargs): """ Perform a forward pass through the autoencoder model. Args: surface (torch.Tensor): The input surface tensor to be encoded. queries (torch.Tensor): The query tensor used for generating logits. **kwargs: Additional keyword arguments. Returns: tuple: A tuple containing: - z (torch.Tensor): The latent representation of the input surface. - latents (torch.Tensor): The decoded output from the latent representation. - None: Placeholder for a potential future return value. - logits (torch.Tensor): The logits generated from the queries and latents. - d (torch.Tensor): Additional output from the encoding process. """ _, z, _, d = self.encode(surface) latents = self.decode(z) logits = self.query(queries, latents) return z, latents, None, logits, d @torch.no_grad() def extract_geometry( self, latents: torch.FloatTensor, bounds: list[float] = [ -1.05, -1.05, -1.05, 1.05, 1.05, 1.05, ], resolution_base: float = 9.0, chunk_size: int = 2_000_000, use_warp: bool = False, ): """ Extracts 3D geometry from latent representations using a dense grid sampling and marching cubes algorithm. Args: latents (torch.FloatTensor): A tensor of latent representations with shape (batch_size, latent_dim). bounds (list[float], optional): A list of six floats defining the bounding box for the 3D grid in the format [xmin, ymin, zmin, xmax, ymax, zmax]. Defaults to [-1.05, -1.05, -1.05, 1.05, 1.05, 1.05]. resolution_base (float, optional): The base resolution for the grid. Higher values result in finer grids. Defaults to 9.0. chunk_size (int, optional): The number of grid points to process in a single chunk. Defaults to 2,000,000. use_warp (bool, optional): Whether to use a GPU-accelerated marching cubes implementation. If False, falls back to a CPU implementation. Defaults to False. Returns: tuple: - mesh_v_f (list[tuple]): A list of tuples containing vertices and faces for each batch element. Each tuple is of the form (vertices, faces), where: - vertices (np.ndarray): Array of vertex coordinates with shape (num_vertices, 3). - faces (np.ndarray): Array of face indices with shape (num_faces, 3). If geometry extraction fails for a batch element, the tuple will be (None, None). - has_surface (np.ndarray): A boolean array indicating whether a surface was successfully extracted for each batch element. Raises: Exception: Logs warnings or errors if geometry extraction fails for any batch element or if the marching cubes algorithm encounters issues. """ bbox_min = np.array(bounds[0:3]) bbox_max = np.array(bounds[3:6]) bbox_size = bbox_max - bbox_min xyz_samples, grid_size, length = generate_dense_grid_points( bbox_min=bbox_min, bbox_max=bbox_max, resolution_base=resolution_base, indexing="ij", ) xyz_samples = torch.FloatTensor(xyz_samples) batch_size = latents.shape[0] batch_logits = [] progress_bar = tqdm( range(0, xyz_samples.shape[0], chunk_size), desc=f"extracting geometry", unit="chunk", ) for start in progress_bar: queries = xyz_samples[start : start + chunk_size, :] num_queries = queries.shape[0] if start > 0 and num_queries < chunk_size: queries = F.pad(queries, [0, 0, 0, chunk_size - num_queries]) batch_queries = queries.unsqueeze(0).expand(batch_size, -1, -1).to(latents) logits = self.query(batch_queries, latents)[:, :num_queries] batch_logits.append(logits) grid_logits = ( torch.cat(batch_logits, dim=1) .detach() .view((batch_size, grid_size[0], grid_size[1], grid_size[2])) .float() ) mesh_v_f = [] has_surface = np.zeros((batch_size,), dtype=np.bool_) for i in range(batch_size): try: warp_success = False if use_warp: try: vertices, faces = marching_cubes_with_warp( grid_logits[i], level=0.0, device=grid_logits.device, ) warp_success = True except Exception as e: logging.warning( f"Warning: error in marching cubes with warp: {e}" ) warp_success = False # Fall back to CPU version if not warp_success: logging.warning( "Warning: falling back to CPU version of marching cubes using skimage measure" ) vertices, faces, _, _ = measure.marching_cubes( grid_logits[i].cpu().numpy(), 0, method="lewiner" ) vertices = vertices / grid_size * bbox_size + bbox_min faces = faces[:, [2, 1, 0]] mesh_v_f.append( (vertices.astype(np.float32), np.ascontiguousarray(faces)) ) has_surface[i] = True except Exception as e: logging.error(f"Error: error in extract_geometry: {e}") mesh_v_f.append((None, None)) has_surface[i] = False return mesh_v_f, has_surface