"""Learned vertex detector for S23DR 2026. Train a CNN to predict 2D vertex heatmaps from gestalt + depth images. Uses ground-truth 3D wireframe vertices projected to 2D as supervision. """ import torch import torch.nn as nn import torch.nn.functional as F import torchvision.models as models import numpy as np from typing import Tuple, List, Optional class VertexHeatmapNet(nn.Module): def __init__(self, in_channels=7, num_classes=2, pretrained_backbone=True): super().__init__() backbone = models.resnet18(weights='IMAGENET1K_V1' if pretrained_backbone else None) self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False) if pretrained_backbone: with torch.no_grad(): pretrained_weight = backbone.conv1.weight new_weight = torch.zeros(64, in_channels, 7, 7) new_weight[:, :3, :, :] = pretrained_weight avg_weight = pretrained_weight.mean(dim=1, keepdim=True) for i in range(3, in_channels): new_weight[:, i:i+1, :, :] = avg_weight self.conv1.weight = nn.Parameter(new_weight) self.bn1 = backbone.bn1 self.relu = backbone.relu self.maxpool = backbone.maxpool self.layer1 = backbone.layer1 self.layer2 = backbone.layer2 self.layer3 = backbone.layer3 self.layer4 = backbone.layer4 self.up4 = nn.Sequential(nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True)) self.up3 = nn.Sequential(nn.ConvTranspose2d(512, 128, 4, stride=2, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True)) self.up2 = nn.Sequential(nn.ConvTranspose2d(256, 64, 4, stride=2, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True)) self.up1 = nn.Sequential(nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True)) self.head = nn.Sequential(nn.Conv2d(64, 32, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(32, num_classes, 1)) def forward(self, x): x = self.conv1(x); x = self.bn1(x); x0 = self.relu(x); x = self.maxpool(x0) x1 = self.layer1(x); x2 = self.layer2(x1); x3 = self.layer3(x2); x4 = self.layer4(x3) d4 = self.up4(x4); d3 = self.up3(torch.cat([d4, x3], dim=1)) d2 = self.up2(torch.cat([d3, x2], dim=1)); d1 = self.up1(torch.cat([d2, x1], dim=1)) return self.head(d1) def create_vertex_heatmap(vertices_2d, vertex_types, height, width, sigma=3.0): heatmap = np.zeros((2, height, width), dtype=np.float32) type_to_channel = {'apex': 0, 'eave_end_point': 1} for (u, v), vtype in zip(vertices_2d, vertex_types): ch = type_to_channel.get(vtype, 0) u_int, v_int = int(round(u)), int(round(v)) if u_int < 0 or u_int >= width or v_int < 0 or v_int >= height: continue radius = int(3 * sigma) for dy in range(-radius, radius + 1): for dx in range(-radius, radius + 1): py, px = v_int + dy, u_int + dx if 0 <= py < height and 0 <= px < width: val = np.exp(-(dx*dx + dy*dy) / (2 * sigma * sigma)) heatmap[ch, py, px] = max(heatmap[ch, py, px], val) return heatmap def prepare_input_tensor(gestalt_img, depth_img, ade_img, target_size=(192, 256)): H, W = target_size gest = np.array(gestalt_img.resize((W, H))).astype(np.float32) / 255.0 if gest.ndim == 2: gest = np.stack([gest]*3, axis=-1) depth = np.array(depth_img.resize((W, H))).astype(np.float32) / 1000.0 depth = np.clip(depth / 50.0, 0, 1) if depth.ndim == 2: depth = depth[:, :, np.newaxis] ade = np.array(ade_img.resize((W, H))).astype(np.float32) / 255.0 if ade.ndim == 2: ade = np.stack([ade]*3, axis=-1) combined = np.concatenate([gest, depth, ade], axis=-1) return torch.from_numpy(combined).permute(2, 0, 1) def extract_vertices_from_heatmap(heatmap, threshold=0.3, nms_radius=5): from scipy.ndimage import maximum_filter vertices, types = [], [] type_names = ['apex', 'eave_end_point'] for ch in range(heatmap.shape[0]): hm = heatmap[ch] local_max = maximum_filter(hm, size=2*nms_radius+1) peaks = (hm == local_max) & (hm >= threshold) ys, xs = np.where(peaks) for y, x in zip(ys, xs): vertices.append([x, y]); types.append(type_names[ch]) if not vertices: return np.zeros((0, 2)), [] return np.array(vertices), types class VertexDetectorTrainer: def __init__(self, model, lr=1e-4, device='cuda'): self.model = model.to(device); self.device = device self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100, eta_min=1e-6) def focal_loss(self, pred, target, alpha=2.0, beta=4.0): pred = torch.clamp(torch.sigmoid(pred), 1e-6, 1 - 1e-6) pos_mask = (target >= 0.99); neg_mask = ~pos_mask pos_loss = -((1 - pred) ** alpha) * torch.log(pred) * pos_mask.float() neg_loss = -((1 - target) ** beta) * (pred ** alpha) * torch.log(1 - pred) * neg_mask.float() return (pos_loss.sum() + neg_loss.sum()) / pos_mask.float().sum().clamp(min=1) def train_step(self, input_tensor, target_heatmap): self.model.train(); self.optimizer.zero_grad() input_tensor = input_tensor.to(self.device); target_heatmap = target_heatmap.to(self.device) pred = self.model(input_tensor) if pred.shape[-2:] != target_heatmap.shape[-2:]: target_heatmap = F.interpolate(target_heatmap, size=pred.shape[-2:], mode='bilinear', align_corners=False) loss = self.focal_loss(pred, target_heatmap) loss.backward(); torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0); self.optimizer.step() return loss.item() @torch.no_grad() def predict(self, input_tensor): self.model.eval() if input_tensor.ndim == 3: input_tensor = input_tensor.unsqueeze(0) return torch.sigmoid(self.model(input_tensor.to(self.device)))[0].cpu().numpy() def save(self, path): torch.save({ 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'scheduler_state_dict': self.scheduler.state_dict(), }, path) def load(self, path): ckpt = torch.load(path, map_location=self.device, weights_only=True) self.model.load_state_dict(ckpt['model_state_dict']) self.optimizer.load_state_dict(ckpt['optimizer_state_dict']) if 'scheduler_state_dict' in ckpt: self.scheduler.load_state_dict(ckpt['scheduler_state_dict'])