import math import torch import torch.nn as nn import torch.nn.functional as F from SkingToolkit.dense_uv_parser.uv_topology import ( build_outer_uv_graph, build_simple_uv_topology, ) def norm_groups(channels): for groups in (8, 4, 2, 1): if channels % groups == 0: return groups return 1 class ConvBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.GroupNorm(norm_groups(out_channels), out_channels), nn.SiLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.GroupNorm(norm_groups(out_channels), out_channels), nn.SiLU(inplace=True), ) def forward(self, x): return self.block(x) class DownBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.down = nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1) self.block = ConvBlock(out_channels, out_channels) def forward(self, x): return self.block(self.down(x)) class UpBlock(nn.Module): def __init__(self, in_channels, skip_channels, out_channels): super().__init__() self.block = ConvBlock(in_channels + skip_channels, out_channels) def forward(self, x, skip): x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False) return self.block(torch.cat([x, skip], dim=1)) class MultiViewSemanticFusion(nn.Module): """Fuse frozen per-view semantic features into parser bottleneck features.""" def __init__( self, raw_feature_dim, semantic_channels, bottleneck_channels, view_classes, attention_heads=4, layers=1, dropout=0.05, ): super().__init__() if raw_feature_dim < 1 or semantic_channels < 1: raise ValueError("Semantic feature dimensions must be positive.") if view_classes < 1: raise ValueError("Semantic fusion requires view-conditioned parser inputs.") if attention_heads < 1: raise ValueError("semantic_attention_heads must be positive.") if semantic_channels % attention_heads != 0: raise ValueError("semantic_channels must be divisible by attention_heads.") if layers < 1: raise ValueError("semantic_layers must be positive.") if not 0.0 <= dropout < 1.0: raise ValueError("semantic_dropout must be in [0, 1).") self.raw_feature_dim = int(raw_feature_dim) self.semantic_channels = int(semantic_channels) self.view_classes = int(view_classes) self.input_projection = nn.Sequential( nn.LayerNorm(raw_feature_dim), nn.Linear(raw_feature_dim, semantic_channels), nn.GELU(), ) self.view_embedding = nn.Parameter( torch.randn(view_classes, semantic_channels) * 0.02 ) self.encoder = nn.ModuleList( nn.TransformerEncoderLayer( d_model=semantic_channels, nhead=attention_heads, dim_feedforward=semantic_channels * 4, dropout=dropout, activation="gelu", batch_first=True, norm_first=True, ) for _ in range(layers) ) self.modulation = nn.Sequential( nn.LayerNorm(semantic_channels * 2), nn.Linear(semantic_channels * 2, semantic_channels * 2), nn.GELU(), nn.Linear(semantic_channels * 2, bottleneck_channels * 2), ) self.summary = nn.Sequential( nn.LayerNorm(semantic_channels), nn.Linear(semantic_channels, semantic_channels), nn.GELU(), ) # Preserve the geometry-only initialization. The adapter starts by # contributing no FiLM shift/scale and learns semantic corrections. nn.init.zeros_(self.modulation[-1].weight) nn.init.zeros_(self.modulation[-1].bias) def forward(self, raw_features, view_ids, sample_count): if raw_features.dim() == 3: # Training may cache multiple complete view groups per skin, for # example the inference pair followed by a privileged right-side # pair. Each group is still fused independently below. raw_features = raw_features.reshape(-1, raw_features.shape[-1]) if raw_features.dim() != 2 or raw_features.shape != ( sample_count, self.raw_feature_dim, ): raise ValueError( "Semantic features must be shaped NxD or BxVxD; got " f"{tuple(raw_features.shape)}." ) if sample_count % self.view_classes != 0: raise ValueError( f"Semantic sample count {sample_count} is not divisible by " f"view count {self.view_classes}." ) grouped_view_ids = view_ids.reshape(-1, self.view_classes) expected_ids = torch.arange( self.view_classes, device=view_ids.device ).view(1, -1) if not torch.equal(grouped_view_ids, expected_ids.expand_as(grouped_view_ids)): raise ValueError("Semantic fusion requires canonical grouped view order.") batch = sample_count // self.view_classes tokens = self.input_projection(raw_features.float()).reshape( batch, self.view_classes, self.semantic_channels ) tokens = tokens + self.view_embedding.unsqueeze(0) for layer in self.encoder: tokens = layer(tokens) pooled = tokens.mean(dim=1) per_view = torch.cat( [tokens, pooled.unsqueeze(1).expand_as(tokens)], dim=-1 ).reshape(sample_count, self.semantic_channels * 2) return self.modulation(per_view), self.summary(pooled) class SpatialSemanticFusion(nn.Module): """Project frozen 2D semantic features into the parser bottleneck. The final projection starts at zero, preserving the stable geometry parser at initialization while allowing spatial semantic corrections to emerge during training. """ def __init__( self, raw_feature_dim, semantic_channels, bottleneck_channels, ): super().__init__() if raw_feature_dim < 1 or semantic_channels < 1: raise ValueError("Spatial semantic feature dimensions must be positive.") self.raw_feature_dim = int(raw_feature_dim) self.semantic_channels = int(semantic_channels) self.input_norm = nn.LayerNorm(raw_feature_dim) self.input_projection = nn.Conv2d( raw_feature_dim, semantic_channels, kernel_size=1 ) self.activation = nn.GELU() self.output_projection = nn.Conv2d( semantic_channels, bottleneck_channels, kernel_size=1 ) nn.init.zeros_(self.output_projection.weight) nn.init.zeros_(self.output_projection.bias) def forward(self, raw_features, sample_count, output_size): if raw_features.dim() == 5: raw_features = raw_features.reshape( -1, raw_features.shape[-3], raw_features.shape[-2], raw_features.shape[-1], ) if ( raw_features.dim() != 4 or raw_features.shape[0] != sample_count or raw_features.shape[1] != self.raw_feature_dim ): raise ValueError( "Spatial semantic features must be shaped NxCxHxW or " f"BxVxCxHxW; got {tuple(raw_features.shape)}." ) normalized = self.input_norm( raw_features.float().permute(0, 2, 3, 1) ).permute(0, 3, 1, 2) residual = self.output_projection( self.activation(self.input_projection(normalized)) ) return F.interpolate( residual, size=output_size, mode="bilinear", align_corners=False, ) class OuterUVGraphBlock(nn.Module): """Message passing over physical neighbours of Minecraft outer texels.""" def __init__(self, channels, dropout=0.05): super().__init__() self.self_projection = nn.Linear(channels, channels) self.neighbour_projection = nn.Linear(channels, channels) self.norm = nn.LayerNorm(channels) self.dropout = nn.Dropout(dropout) def forward(self, nodes, edge_index, degree): source, target = edge_index neighbour_sum = torch.zeros_like(nodes) neighbour_sum.index_add_(1, target, nodes[:, source]) neighbour_mean = neighbour_sum / degree.view(1, -1, 1).clamp_min(1.0) update = self.self_projection(nodes) update = update + self.neighbour_projection(neighbour_mean) return self.norm(nodes + self.dropout(F.gelu(update))) class ProjectedOuterUVTopologyHead(nn.Module): """Predict outer alpha from projected image features and cube topology.""" def __init__( self, input_channels, global_context_dim, hidden_channels=64, layers=3, dropout=0.05, uv_size=64, ): super().__init__() if layers < 1: raise ValueError("outer_uv_topology_layers must be positive.") flat_indices, edge_index = build_outer_uv_graph() topology = build_simple_uv_topology() parts = topology.part.reshape(-1)[flat_indices] faces = topology.face.reshape(-1)[flat_indices] local_uv = topology.local_uv.reshape(-1, 2)[flat_indices] degree = torch.bincount( edge_index[1], minlength=flat_indices.numel(), ).float() self.uv_size = int(uv_size) self.register_buffer("flat_indices", flat_indices, persistent=False) self.register_buffer("edge_index", edge_index, persistent=False) self.register_buffer("degree", degree, persistent=False) self.register_buffer("parts", parts, persistent=False) self.register_buffer("faces", faces, persistent=False) self.register_buffer("local_uv", local_uv, persistent=False) self.input_projection = nn.Linear(input_channels, hidden_channels) self.part_embedding = nn.Embedding(6, hidden_channels) self.face_embedding = nn.Embedding(6, hidden_channels) self.position_projection = nn.Linear(2, hidden_channels) self.global_projection = nn.Sequential( nn.LayerNorm(global_context_dim), nn.Linear(global_context_dim, hidden_channels), nn.GELU(), ) self.blocks = nn.ModuleList( OuterUVGraphBlock(hidden_channels, dropout=dropout) for _ in range(layers) ) self.output = nn.Sequential( nn.LayerNorm(hidden_channels), nn.Linear(hidden_channels, 1), ) # A very negative initialization made the one-epoch recipe prone to # the all-transparent solution after precision penalties were added. # Keep a sparse prior without starving early positive gradients. nn.init.constant_(self.output[-1].bias, -1.0) def forward(self, atlas_features, global_context): if atlas_features.dim() != 4: raise ValueError("Projected outer UV features must be BCHW.") if atlas_features.shape[-2:] != (self.uv_size, self.uv_size): raise ValueError( "Projected outer UV feature size must match the UV atlas." ) flattened = atlas_features.flatten(2).transpose(1, 2) nodes = flattened.index_select(1, self.flat_indices) nodes = self.input_projection(nodes.float()) nodes = nodes + self.part_embedding(self.parts).unsqueeze(0) nodes = nodes + self.face_embedding(self.faces).unsqueeze(0) nodes = nodes + self.position_projection(self.local_uv).unsqueeze(0) nodes = nodes + self.global_projection(global_context.float()).unsqueeze(1) for block in self.blocks: nodes = block(nodes, self.edge_index, self.degree) node_logits = self.output(nodes).squeeze(-1) atlas_logits = node_logits.new_full( (node_logits.shape[0], self.uv_size * self.uv_size), -12.0, ) atlas_logits[:, self.flat_indices] = node_logits return atlas_logits.reshape( node_logits.shape[0], 1, self.uv_size, self.uv_size, ) class DenseUVParserNet(nn.Module): """Predict dense Minecraft UV routing for each render pixel.""" def __init__( self, input_channels=4, base_channels=32, part_classes=6, face_classes=6, layer_classes=None, layer_face_classes=12, uv_size=64, uv_classification=True, view_classes=0, predict_affine=False, affine_translation_scale=0.0, affine_scale_range=0.0, surface_classes=0, geometry_only=False, feature_dropout=0.0, semantic_feature_dim=0, semantic_channels=128, semantic_attention_heads=4, semantic_layers=1, semantic_dropout=0.05, semantic_spatial_feature_dim=0, semantic_spatial_channels=64, predict_confidence=False, route_role_spatial_prior=False, route_prior_height=32, route_prior_width=16, route_prior_logit_cap=1.5, route_prior_dropout=0.10, predict_head_accessories=False, predict_hat_components=False, predict_head_ownership=False, predict_headwear=False, predict_headwear_presence=False, predict_headphone_presence=False, accessory_route_threshold=0.90, predict_outer_uv_occupancy=False, outer_uv_feature_channels=32, outer_uv_topology_channels=64, outer_uv_topology_layers=3, outer_uv_topology_dropout=0.05, outer_uv_route_evidence_dropout=1.0, ): super().__init__() self.accessory_route_threshold = float(accessory_route_threshold) if not 0.5 <= self.accessory_route_threshold <= 1.: raise ValueError("Accessory seed confidence must lie in [0.5, 1]") self.predict_hat_components = bool(predict_hat_components) self.predict_head_accessories = bool(predict_head_accessories) if self.predict_head_accessories: from SkingToolkit.dense_uv_parser.accessories import HeadAccessoryHead self.accessory_head = HeadAccessoryHead(semantic_spatial_feature_dim, self.predict_hat_components) self.predict_headphone_presence = bool(predict_headphone_presence) self.predict_headwear_presence = bool(predict_headwear_presence) self.predict_headwear = bool(predict_headwear) if self.predict_headwear: from SkingToolkit.dense_uv_parser.headwear import HeadwearHead self.headwear_head = HeadwearHead(semantic_spatial_feature_dim,self.predict_headwear_presence) self.predict_head_ownership = bool(predict_head_ownership) if self.predict_head_ownership: from SkingToolkit.dense_uv_parser.ownership import HeadOwnershipHead self.ownership_head = HeadOwnershipHead(semantic_spatial_feature_dim,self.predict_headphone_presence) self.geometry_only = bool(geometry_only) if layer_classes is None: layer_classes = 3 if self.geometry_only else 2 self.layer_classes = int(layer_classes) self.uv_size = int(uv_size) self.uv_classification = bool(uv_classification) and not self.geometry_only self.view_classes = int(view_classes) self.predict_affine = bool(predict_affine) self.affine_translation_scale = float(affine_translation_scale) self.affine_scale_range = float(affine_scale_range) self.surface_classes = int(surface_classes) self.semantic_feature_dim = int(semantic_feature_dim) self.semantic_channels = int(semantic_channels) self.semantic_attention_heads = int(semantic_attention_heads) self.semantic_layers = int(semantic_layers) self.semantic_dropout = float(semantic_dropout) self.semantic_spatial_feature_dim = int(semantic_spatial_feature_dim) self.semantic_spatial_channels = int(semantic_spatial_channels) self.predict_confidence = bool(predict_confidence) self.route_role_spatial_prior = bool(route_role_spatial_prior) self.route_prior_height = int(route_prior_height) self.route_prior_width = int(route_prior_width) self.route_prior_logit_cap = float(route_prior_logit_cap) self.route_prior_dropout = float(route_prior_dropout) self.predict_outer_uv_occupancy = bool(predict_outer_uv_occupancy) self.outer_uv_feature_channels = int(outer_uv_feature_channels) self.outer_uv_topology_channels = int(outer_uv_topology_channels) self.outer_uv_topology_layers = int(outer_uv_topology_layers) self.outer_uv_topology_dropout = float(outer_uv_topology_dropout) self.outer_uv_route_evidence_dropout = float( outer_uv_route_evidence_dropout ) if self.route_prior_height < 1 or self.route_prior_width < 1: raise ValueError("Route-prior dimensions must be positive.") if self.route_prior_logit_cap <= 0.0: raise ValueError("route_prior_logit_cap must be positive.") if not 0.0 <= self.route_prior_dropout < 1.0: raise ValueError("route_prior_dropout must be in [0, 1).") if self.route_role_spatial_prior and ( not self.geometry_only or self.view_classes < 1 ): raise ValueError( "The fixed-view route-role prior requires geometry_only with view classes." ) self.feature_dropout_probability = float(feature_dropout) if not 0.0 <= self.feature_dropout_probability < 1.0: raise ValueError("feature_dropout must be in [0, 1).") if not 0.0 <= self.outer_uv_route_evidence_dropout <= 1.0: raise ValueError( "outer_uv_route_evidence_dropout must be in [0, 1]." ) c = base_channels self.stem = ConvBlock(input_channels + self.view_classes, c) self.down1 = DownBlock(c, c * 2) self.down2 = DownBlock(c * 2, c * 4) self.down3 = DownBlock(c * 4, c * 8) self.mid = ConvBlock(c * 8, c * 8) self.semantic_fusion = ( MultiViewSemanticFusion( self.semantic_feature_dim, self.semantic_channels, c * 8, self.view_classes, attention_heads=self.semantic_attention_heads, layers=self.semantic_layers, dropout=self.semantic_dropout, ) if self.semantic_feature_dim > 0 else None ) self.semantic_spatial_fusion = ( SpatialSemanticFusion( self.semantic_spatial_feature_dim, self.semantic_spatial_channels, c * 8, ) if self.semantic_spatial_feature_dim > 0 else None ) self.up2 = UpBlock(c * 8, c * 4, c * 4) self.up1 = UpBlock(c * 4, c * 2, c * 2) self.up0 = UpBlock(c * 2, c, c) self.features = nn.Sequential( nn.Conv2d(c, c, kernel_size=3, padding=1), nn.SiLU(inplace=True), ) self.feature_dropout = nn.Dropout2d(self.feature_dropout_probability) self.foreground = nn.Conv2d(c, 1, kernel_size=1) self.layer = nn.Conv2d(c, self.layer_classes, kernel_size=1) self.route_role_prior = ( nn.Parameter( torch.zeros( self.view_classes, self.layer_classes, self.route_prior_height, self.route_prior_width, ) ) if self.route_role_spatial_prior else None ) self.route_confidence = ( nn.Conv2d(c, 1, kernel_size=1) if self.predict_confidence else None ) if self.semantic_fusion is not None: self.outer_presence_head = nn.Linear(self.semantic_channels, 6) self.outer_coverage_head = nn.Linear(self.semantic_channels, 6) if self.predict_outer_uv_occupancy: if self.view_classes < 1: raise ValueError( "Outer UV occupancy prediction requires grouped fixed-view inputs." ) if self.outer_uv_feature_channels < 1: raise ValueError("outer_uv_feature_channels must be positive.") occupancy_context_channels = c * 8 + ( self.semantic_channels if self.semantic_fusion is not None else 0 ) self.outer_uv_feature_projection = nn.Sequential( nn.Conv2d(c, self.outer_uv_feature_channels, kernel_size=1), nn.GELU(), ) # Projected feature mean + mean/max p_outer + foreground coverage # + number of supporting views. occupancy_input_channels = self.outer_uv_feature_channels + 4 self.outer_uv_occupancy_head = ProjectedOuterUVTopologyHead( occupancy_input_channels, occupancy_context_channels, hidden_channels=self.outer_uv_topology_channels, layers=self.outer_uv_topology_layers, dropout=self.outer_uv_topology_dropout, uv_size=uv_size, ) if not self.geometry_only: self.part = nn.Conv2d(c, part_classes, kernel_size=1) self.face = nn.Conv2d(c, face_classes, kernel_size=1) self.layer_face = ( nn.Conv2d(c, layer_face_classes, kernel_size=1) if layer_face_classes > 0 else None ) self.uv = nn.Conv2d(c, 2, kernel_size=1) if self.uv_classification: self.uv_x = nn.Conv2d(c, uv_size, kernel_size=1) self.uv_y = nn.Conv2d(c, uv_size, kernel_size=1) else: self.layer_face = None if self.predict_affine: if not self.geometry_only and self.surface_classes < 2: raise ValueError("Global-affine routing requires at least two static surface classes.") if self.surface_classes == 1: raise ValueError("Surface routing requires at least two static surface classes.") if self.surface_classes > 0: self.surface = nn.Conv2d(c, self.surface_classes, kernel_size=1) hidden = max(c * 4, 32) self.affine_head = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(c * 8, hidden), nn.SiLU(inplace=True), nn.Linear(hidden, 3), ) # Start from the canonical pose; the residual head learns only the # small global transform introduced by the configured augmentation. nn.init.zeros_(self.affine_head[-1].weight) nn.init.zeros_(self.affine_head[-1].bias) self.affine_translation_limit = 2.0 * self.affine_translation_scale lower_log_scale = math.log(max(1.0 - self.affine_scale_range, 1e-6)) upper_log_scale = math.log1p(self.affine_scale_range) self.affine_log_scale_limit = max(abs(lower_log_scale), abs(upper_log_scale)) def _runtime_semantic_features(self, images, foreground=None): backbone = getattr(self, "_runtime_semantic_backbone", None) if backbone is None: raise ValueError( "This parser requires semantic_features or an attached semantic " "runtime backbone." ) rgb = images[:, :3] if foreground is not None: if foreground.dim() == 3: foreground = foreground.unsqueeze(1) if ( foreground.dim() != 4 or foreground.shape[0] != rgb.shape[0] or foreground.shape[1] != 1 ): raise ValueError( "semantic_foreground must be shaped Nx1xHxW or NxHxW; got " f"{tuple(foreground.shape)}." ) if foreground.shape[-2:] != rgb.shape[-2:]: foreground = F.interpolate( foreground.float(), size=rgb.shape[-2:], mode="nearest", ) foreground = foreground.to(device=rgb.device, dtype=rgb.dtype).clamp(0, 1) # A fixed neutral background prevents the frozen spatial tower from # learning accidental correlations with randomized training colors. rgb = rgb * foreground + 0.5 * (1.0 - foreground) with torch.no_grad(): if hasattr(backbone, "encode_dense"): return backbone.encode_dense(rgb) return backbone.encode_global(rgb) def forward( self, x, view_ids=None, semantic_features=None, semantic_foreground=None, ): source_images = x if self.view_classes > 0: if view_ids is None: raise ValueError("view_ids are required for a view-conditioned dense UV parser.") if view_ids.shape != (x.shape[0],): raise ValueError(f"Expected view_ids shape {(x.shape[0],)}, got {tuple(view_ids.shape)}.") if view_ids.min() < 0 or view_ids.max() >= self.view_classes: raise ValueError(f"view_ids must be in [0, {self.view_classes - 1}].") view_one_hot = F.one_hot(view_ids.long(), num_classes=self.view_classes).to(dtype=x.dtype) view_one_hot = view_one_hot.view(x.shape[0], self.view_classes, 1, 1) x = torch.cat([x, view_one_hot.expand(-1, -1, x.shape[2], x.shape[3])], dim=1) s0 = self.stem(x) s1 = self.down1(s0) s2 = self.down2(s1) s3 = self.down3(s2) x = self.mid(s3) if ( self.semantic_fusion is not None or self.semantic_spatial_fusion is not None ) and semantic_features is None: semantic_features = self._runtime_semantic_features( source_images, foreground=semantic_foreground ) if isinstance(semantic_features, dict): semantic_global = semantic_features.get( "raw_global", semantic_features.get("global") ) semantic_spatial = semantic_features.get( "raw_spatial", semantic_features.get("spatial") ) else: semantic_global = semantic_features semantic_spatial = None if self.semantic_spatial_fusion is not None: if semantic_spatial is None: raise ValueError( "This parser checkpoint requires spatial semantic features." ) x = x + self.semantic_spatial_fusion( semantic_spatial, source_images.shape[0], x.shape[-2:], ).to(dtype=x.dtype) visual_summary = x.mean(dim=(2, 3)) if self.predict_outer_uv_occupancy: if visual_summary.shape[0] % self.view_classes != 0: raise ValueError( "Outer UV occupancy prediction requires complete fixed-view groups." ) visual_summary = visual_summary.reshape( -1, self.view_classes, visual_summary.shape[-1] ).mean(dim=1) semantic_summary = None if self.semantic_fusion is not None: if semantic_global is None: raise ValueError( "This parser checkpoint requires global semantic features." ) modulation, semantic_summary = self.semantic_fusion( semantic_global, view_ids, source_images.shape[0], ) scale, shift = modulation.chunk(2, dim=1) x = x * (1.0 + scale.unsqueeze(-1).unsqueeze(-1)) x = x + shift.unsqueeze(-1).unsqueeze(-1) affine = None if self.predict_affine: raw_affine = torch.tanh(self.affine_head(x)) affine = torch.stack( [ raw_affine[:, 0] * self.affine_translation_limit, raw_affine[:, 1] * self.affine_translation_limit, raw_affine[:, 2] * self.affine_log_scale_limit, ], dim=1, ) x = self.up2(x, s2) x = self.up1(x, s1) x = self.up0(x, s0) x = self.features(x) occupancy_feature_source = x x = self.feature_dropout(x) layer_evidence = self.layer(x) outputs = { "foreground": self.foreground(x), "layer": layer_evidence, } if self.route_role_prior is not None: selected_prior_raw = self.route_role_prior.index_select( 0, view_ids.long() ) selected_prior = self.route_prior_logit_cap * torch.tanh( selected_prior_raw / self.route_prior_logit_cap ) if self.training and self.route_prior_dropout > 0.0: keep = ( torch.rand( selected_prior.shape[0], 1, 1, 1, device=selected_prior.device, ) >= self.route_prior_dropout ) selected_prior = selected_prior * keep.to( dtype=selected_prior.dtype ) route_prior = F.interpolate( selected_prior, size=layer_evidence.shape[-2:], mode="bilinear", align_corners=False, ).to(dtype=layer_evidence.dtype) outputs["layer"] = layer_evidence + route_prior outputs["route_role_evidence"] = layer_evidence outputs["route_role_prior"] = route_prior outputs["route_role_prior_raw"] = self.route_role_prior if self.route_confidence is not None: outputs["route_confidence"] = self.route_confidence(x) if semantic_summary is not None: outputs["outer_presence_logits"] = self.outer_presence_head( semantic_summary ) outputs["outer_coverage"] = torch.sigmoid( self.outer_coverage_head(semantic_summary) ) if self.predict_outer_uv_occupancy: occupancy_summary = ( torch.cat([visual_summary, semantic_summary], dim=1) if semantic_summary is not None else visual_summary ) # Occupancy learns from the parser representation without allowing # its sparse atlas loss to destabilize the proven image-space route # trunk. A separate agreement loss later transfers calibrated UV # structure back into the route head after warm-up. outputs["outer_uv_features"] = self.outer_uv_feature_projection( occupancy_feature_source.detach() ) outputs["outer_uv_global_context"] = occupancy_summary.detach() if not self.geometry_only: outputs["part"] = self.part(x) outputs["face"] = self.face(x) outputs["uv"] = torch.sigmoid(self.uv(x)) if self.layer_face is not None: outputs["layer_face"] = self.layer_face(x) if self.uv_classification: outputs["uv_x"] = self.uv_x(x) outputs["uv_y"] = self.uv_y(x) if affine is not None: # [tx, ty, log_scale]. tx/ty are affine_grid normalized coordinates. outputs["affine"] = affine if self.surface_classes > 0: outputs["surface"] = self.surface(x) if self.predict_head_accessories: prediction=self.predict_accessories(source_images,semantic_foreground,return_components=self.predict_hat_components) if self.predict_hat_components: outputs["accessory_logits"],outputs["hat_component_logits"]=prediction else:outputs["accessory_logits"]=prediction outputs["accessory_route_threshold"] = self.accessory_route_threshold if self.predict_headwear: prediction = self.predict_headwear_components(source_images, semantic_foreground,return_presence=self.predict_headwear_presence) if self.predict_headwear_presence:outputs["headwear_logits"],outputs["headwear_presence_logits"]=prediction else:outputs["headwear_logits"]=prediction if self.predict_head_ownership: prediction = self.predict_ownership(source_images, semantic_foreground, return_presence=self.predict_headphone_presence) if self.predict_headphone_presence: outputs["head_ownership_logits"],outputs["headphone_presence_logit"] = prediction else:outputs["head_ownership_logits"] = prediction return outputs def predict_headwear_components(self, images, foreground=None, return_presence=False): from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits crop = head_crop(images, foreground) semantic = self._runtime_semantic_features(crop)["raw_spatial"] prediction=self.headwear_head(crop,semantic,return_presence=return_presence) if return_presence:return restore_logits(prediction[0],images.shape[-2:]),prediction[1] return restore_logits(prediction,images.shape[-2:]) def predict_ownership(self, images, foreground=None, return_presence=False): from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits crop = head_crop(images, foreground) semantic = self._runtime_semantic_features(crop)["raw_spatial"] prediction = self.ownership_head(crop, semantic, return_presence=return_presence) if return_presence:return restore_logits(prediction[0], images.shape[-2:]),prediction[1] return restore_logits(prediction, images.shape[-2:]) def predict_accessories(self, images, foreground=None, return_components=False): from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits crop = head_crop(images, foreground) semantic = self._runtime_semantic_features(crop)["raw_spatial"] logits = self.accessory_head(crop, semantic, return_components=return_components) if isinstance(logits,tuple):return tuple(restore_logits(x,images.shape[-2:]) for x in logits) return restore_logits(logits, images.shape[-2:]) def predict_projected_outer_uv_occupancy( self, atlas_features, global_context, ): if not self.predict_outer_uv_occupancy: raise ValueError("This parser has no outer UV occupancy head.") return self.outer_uv_occupancy_head( atlas_features, global_context, ) def count_parameters(model): return sum(param.numel() for param in model.parameters() if param.requires_grad)