Andrea Pozzetti
2859c0a3e5
Split SECAD-Net nodes and utilities from the CADabra monorepo. Voxel-based sketch-extrude CAD reconstruction. Original source: https://github.com/PozzettiAndrea/ComfyUI-CADabra
199 lines
6.6 KiB
Python
199 lines
6.6 KiB
Python
"""
|
|
SECAD-Net neural network architecture.
|
|
|
|
Paper: SECAD-Net: Self-Supervised CAD Reconstruction by Learning Sketch-Extrude Operations
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
|
|
from .sdfs import sdfExtrusion, transform_points, add_latent
|
|
|
|
|
|
class Encoder(nn.Module):
|
|
"""
|
|
Encoder network that converts 64³ voxel grids to 256D latent codes.
|
|
|
|
Architecture: 5 Conv3D layers with stride=2, LeakyReLU activation.
|
|
"""
|
|
|
|
def __init__(self, ef_dim=32):
|
|
super(Encoder, self).__init__()
|
|
self.ef_dim = ef_dim
|
|
self.conv_1 = nn.Conv3d(1, self.ef_dim, 4, stride=2, padding=1, bias=True)
|
|
self.conv_2 = nn.Conv3d(self.ef_dim, self.ef_dim * 2, 4, stride=2, padding=1, bias=True)
|
|
self.conv_3 = nn.Conv3d(self.ef_dim * 2, self.ef_dim * 4, 4, stride=2, padding=1, bias=True)
|
|
self.conv_4 = nn.Conv3d(self.ef_dim * 4, self.ef_dim * 8, 4, stride=2, padding=1, bias=True)
|
|
self.conv_5 = nn.Conv3d(self.ef_dim * 8, self.ef_dim * 8, 4, stride=1, padding=0, bias=True)
|
|
self._initialize_weights()
|
|
|
|
def _initialize_weights(self):
|
|
for conv in [self.conv_1, self.conv_2, self.conv_3, self.conv_4, self.conv_5]:
|
|
nn.init.xavier_uniform_(conv.weight)
|
|
nn.init.constant_(conv.bias, 0)
|
|
|
|
def forward(self, inputs):
|
|
"""
|
|
Args:
|
|
inputs: Voxel grid tensor [B, 1, 64, 64, 64]
|
|
Returns:
|
|
Latent code tensor [B, 256]
|
|
"""
|
|
d = inputs
|
|
for i in range(1, 6):
|
|
d = getattr(self, f"conv_{i}")(d)
|
|
d = F.leaky_relu(d, negative_slope=0.01)
|
|
d = d.view(-1, self.ef_dim * 8)
|
|
d = F.leaky_relu(d, negative_slope=0.01)
|
|
return d
|
|
|
|
|
|
class Decoder(nn.Module):
|
|
"""
|
|
Decoder network that converts latent codes to primitive parameters.
|
|
|
|
Each primitive has 8 parameters:
|
|
- 4 values: Quaternion (rotation)
|
|
- 3 values: Translation (XYZ position)
|
|
- 1 value: Extrusion height
|
|
"""
|
|
|
|
def __init__(self, ef_dim=32, num_primitives=4):
|
|
super(Decoder, self).__init__()
|
|
self.num_primitives = num_primitives
|
|
self.feature_dim = ef_dim
|
|
|
|
self.num_primitive_parameters_aggregated = 4 + 3 + 1 # quat + trans + height
|
|
self.primitive_linear = nn.Linear(
|
|
self.feature_dim * 8,
|
|
int(self.num_primitives * self.num_primitive_parameters_aggregated),
|
|
bias=True,
|
|
)
|
|
nn.init.xavier_uniform_(self.primitive_linear.weight)
|
|
nn.init.constant_(self.primitive_linear.bias, 0)
|
|
|
|
def forward(self, feature):
|
|
"""
|
|
Args:
|
|
feature: Latent code tensor [B, 256]
|
|
Returns:
|
|
Primitive parameters tensor [B, 8, num_primitives]
|
|
"""
|
|
shapes = self.primitive_linear(feature)
|
|
para_3d = shapes[..., : self.num_primitives * (4 + 3 + 1)].view(
|
|
-1, (4 + 3 + 1), int(self.num_primitives)
|
|
) # B,C,P
|
|
return para_3d
|
|
|
|
|
|
class SketchHead(nn.Module):
|
|
"""
|
|
Neural network head that computes 2D sketch occupancy from point coordinates.
|
|
|
|
Architecture: 3-layer MLP with Softplus activation.
|
|
"""
|
|
|
|
def __init__(self, d_in, dims):
|
|
super().__init__()
|
|
dims = [d_in] + dims + [1]
|
|
self.num_layers = len(dims)
|
|
|
|
for layer in range(0, self.num_layers - 1):
|
|
out_dim = dims[layer + 1]
|
|
lin = nn.Linear(dims[layer], out_dim)
|
|
|
|
if layer == self.num_layers - 2:
|
|
torch.nn.init.normal_(
|
|
lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[layer]), std=0.00001
|
|
)
|
|
torch.nn.init.constant_(lin.bias, -1)
|
|
else:
|
|
torch.nn.init.constant_(lin.bias, 0.0)
|
|
torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
|
|
|
|
setattr(self, "lin" + str(layer), lin)
|
|
|
|
self.activation = nn.Softplus(beta=100)
|
|
|
|
def forward(self, input):
|
|
x = input
|
|
for layer in range(0, self.num_layers - 1):
|
|
lin = getattr(self, "lin" + str(layer))
|
|
x = lin(x)
|
|
if layer < self.num_layers - 2:
|
|
x = self.activation(x)
|
|
else:
|
|
x = x.clamp(-1, 1)
|
|
return x
|
|
|
|
|
|
class Generator(nn.Module):
|
|
"""
|
|
Generator network that computes occupancy fields from primitive parameters.
|
|
|
|
Uses neural sketch heads to compute 2D profiles, then extrudes to 3D.
|
|
"""
|
|
|
|
def __init__(self, num_primitives=4, sharpness=150, test=False):
|
|
super(Generator, self).__init__()
|
|
self.num_primitives = num_primitives
|
|
self.sharpness = sharpness
|
|
self.test = test
|
|
|
|
D_IN = 2
|
|
LATENT_SIZE = 256
|
|
|
|
for i in range(num_primitives):
|
|
setattr(
|
|
self,
|
|
'sketch_head_' + str(i),
|
|
SketchHead(d_in=D_IN + LATENT_SIZE, dims=[512, 512, 512]),
|
|
)
|
|
|
|
def forward(self, sample_point_coordinates, primitive_parameters, code):
|
|
"""
|
|
Args:
|
|
sample_point_coordinates: Query points [B, N, 3]
|
|
primitive_parameters: Parameters from decoder [B, 8, num_primitives]
|
|
code: Latent code [B, 256]
|
|
Returns:
|
|
union_occupancies: Combined occupancy [B, N]
|
|
total_2d_occ: 2D sketch occupancies [B, N, num_primitives]
|
|
transformed_points: Points in primitive-local coordinates
|
|
"""
|
|
B, N = sample_point_coordinates.shape[:2]
|
|
primitive_parameters = primitive_parameters.transpose(2, 1)
|
|
B, K, param_dim = primitive_parameters.shape
|
|
|
|
boxes = primitive_parameters[..., :8]
|
|
|
|
transformed_points = transform_points(
|
|
boxes[..., :4], boxes[..., 4:7], sample_point_coordinates
|
|
) # [B, N, K, 3]
|
|
|
|
latent_points = []
|
|
for i in range(self.num_primitives):
|
|
latent_point = add_latent(transformed_points[..., i, :2], code).float()
|
|
latent_points.append(latent_point)
|
|
|
|
sdfs_2d = [
|
|
getattr(self, f'sketch_head_{i}')(latent_points[i]).reshape(B, N, -1).float()
|
|
for i in range(self.num_primitives)
|
|
]
|
|
sdfs_2d = torch.cat(sdfs_2d, dim=-1)
|
|
|
|
total_2d_occ = torch.sigmoid(-1 * sdfs_2d * self.sharpness)
|
|
|
|
box_ext = sdfExtrusion(sdfs_2d, boxes[..., 7], transformed_points).squeeze(-1)
|
|
primitive_sdf = box_ext
|
|
primitive_occupancies = torch.sigmoid(-1 * primitive_sdf * self.sharpness)
|
|
|
|
with torch.no_grad():
|
|
weights = torch.softmax(primitive_occupancies * 20, dim=-1)
|
|
|
|
union_occupancies = torch.sum(weights * primitive_occupancies, dim=-1)
|
|
|
|
return union_occupancies, total_2d_occ, transformed_points
|