| |
| |
| import math |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from .data import grid_centers |
|
|
|
|
| class Student(nn.Module): |
| """Small image-conditioned coordinate MLP; fixed cells avoid unordered point regression. |
| |
| 15 Gaussian outputs + 2 auxiliary geometry logits. This is a coarse baseline, |
| not an architecture expected to reproduce million-splat teacher detail. |
| """ |
| def __init__(self, grid=16, width=128): |
| super().__init__() |
| if grid not in (8, 16, 32) or width < 16: |
| raise ValueError("Unsupported grid or width") |
| self.config = {"grid": grid, "width": width} |
| self.grid = grid |
| blocks, channels = [], 3 |
| for out in (24, 48, 96, 128): |
| blocks.extend([nn.Conv2d(channels, out, 3, stride=2, padding=1), nn.GroupNorm(8, out), nn.SiLU()]) |
| channels = out |
| self.encoder = nn.Sequential(*blocks, nn.AdaptiveAvgPool2d(4), nn.Flatten(), |
| nn.Linear(128*4*4, width), nn.SiLU()) |
| centers = torch.from_numpy(grid_centers(grid)) |
| freq = 2**torch.arange(4, dtype=torch.float32) * math.pi * 2 |
| angles = centers[..., None] * freq |
| self.register_buffer("centers", centers) |
| self.register_buffer("coords", torch.cat([centers, angles.sin().flatten(1), angles.cos().flatten(1)], dim=1)) |
| self.query = nn.Sequential(nn.Linear(27, width), nn.SiLU()) |
| self.decoder = nn.Sequential(nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width), nn.SiLU(), nn.Linear(width, 17)) |
| with torch.no_grad(): |
| self.decoder[-1].bias.zero_() |
| self.decoder[-1].bias[0] = -2 |
|
|
| def forward(self, images): |
| features = self.encoder(images) |
| raw = self.decoder(features[:, None, :] + self.query(self.coords)[None]) |
| quat = raw[..., 10:14] + raw.new_tensor([1, 0, 0, 0]) |
| return {"occ": raw[..., :1], "xyz": self.centers[None] + raw[..., 1:4].tanh() / (2*self.grid), |
| "rgb": raw[..., 4:7].sigmoid(), |
| "scale": (raw[..., 7:10].clamp(-8, 5) - math.log(self.grid*3)).exp(), |
| "quat": F.normalize(quat, dim=-1, eps=1e-8), "opacity": raw[..., 14:15].sigmoid(), |
| "geometry": raw[..., 15:17]} |
|
|
|
|
| def balanced_bce(logits, target): |
| error = F.binary_cross_entropy_with_logits(logits, target, reduction="none") |
| positive = (error * target).sum() / target.sum().clamp_min(1) |
| negative = (error * (1-target)).sum() / (1-target).sum().clamp_min(1) |
| return (positive + negative) / 2 |
|
|
|
|
| def distillation_loss(pred, target, grid): |
| mask = target["occ"] |
| def masked(value): |
| return (value * mask).sum() / (mask.sum().clamp_min(1) * value.shape[-1]) |
| terms = {"occupancy": balanced_bce(pred["occ"], mask), |
| "position": masked(((pred["xyz"] - target["xyz"]) * grid).square()), |
| "colour": masked((pred["rgb"] - target["rgb"].clamp(0, 1)).square()), |
| "scale": masked((pred["scale"].log() - target["scale"].clamp_min(1e-8).log()).square()), |
| "rotation": masked(1 - (pred["quat"] * target["quat"]).sum(-1, keepdim=True).square().clamp(0, 1)), |
| "opacity": masked((pred["opacity"] - target["opacity"]).square())} |
| geometry = pred["geometry"].sum() * 0 |
| for i in range(2): |
| available = target["geometry_valid"][:, i] > 0 |
| if available.any(): |
| geometry = geometry + balanced_bce(pred["geometry"][available, :, i:i+1], target["geometry"][available, :, i:i+1]) |
| terms["geometry"] = geometry |
| total = terms["occupancy"] + terms["position"] + 2*terms["colour"] + .1*terms["scale"] + .1*terms["rotation"] + terms["opacity"] + .25*geometry |
| return total, terms |
|
|
