import math import sys from pathlib import Path import torch import torch.nn.functional as F TOOLKIT_ROOT = Path(__file__).resolve().parents[1] WORKSPACE_ROOT = TOOLKIT_ROOT.parent if str(WORKSPACE_ROOT) not in sys.path: sys.path.insert(0, str(WORKSPACE_ROOT)) from SkingToolkit.dense_uv_parser.uv_layout import minecraft_layer_rects # noqa: E402 from SkingToolkit.dense_uv_parser.uv_topology import ( # noqa: E402 build_outer_uv_graph, ) IGNORE_INDEX = 255 PART_CLASSES = 6 FACE_CLASSES = 6 LAYER_CLASSES = 2 LAYER_FACE_CLASSES = LAYER_CLASSES * FACE_CLASSES ROUTE_ROLE_CLASSES = 3 ROUTE_INNER_PRIMARY = 0 ROUTE_OUTER_PRIMARY = 1 ROUTE_SECONDARY = 2 UV_SIZE = 64 SPLAT_COLOR_AGGREGATIONS = ("grid_mode", "texel_center", "exact_mode", "best") SOLID_BACKGROUND_COLOR_TOLERANCE = 16.0 / 255.0 SOLID_BACKGROUND_MIN_CORNER_SUPPORT = 0.5 SECONDARY_TEXEL_MIN_PROBABILITY = 0.5 def parse_views(views): if isinstance(views, str): return [view.strip() for view in views.split(",") if view.strip()] return list(views) def build_part_face_lookups(device=None): inner_part = torch.full((UV_SIZE, UV_SIZE), IGNORE_INDEX, dtype=torch.long, device=device) inner_face = torch.full((UV_SIZE, UV_SIZE), IGNORE_INDEX, dtype=torch.long, device=device) outer_part = torch.full((UV_SIZE, UV_SIZE), IGNORE_INDEX, dtype=torch.long, device=device) outer_face = torch.full((UV_SIZE, UV_SIZE), IGNORE_INDEX, dtype=torch.long, device=device) rects = minecraft_layer_rects(is_slim=False) for rect_index, (inner_x, inner_y, width, height, decor_dx, decor_dy) in enumerate(rects): part = rect_index // FACE_CLASSES face = rect_index % FACE_CLASSES inner_part[inner_y : inner_y + height, inner_x : inner_x + width] = part inner_face[inner_y : inner_y + height, inner_x : inner_x + width] = face outer_x = inner_x + decor_dx outer_y = inner_y + decor_dy outer_part[outer_y : outer_y + height, outer_x : outer_x + width] = part outer_face[outer_y : outer_y + height, outer_x : outer_x + width] = face return { "inner_part": inner_part, "inner_face": inner_face, "outer_part": outer_part, "outer_face": outer_face, } def _grid_to_uv01(grid): return ((grid + 1.0) * 0.5).clamp(0.0, 1.0) def _grid_to_xy(grid): coords = (_grid_to_uv01(grid) * (UV_SIZE - 1)).round().long() x = coords[..., 0].clamp(0, UV_SIZE - 1) y = coords[..., 1].clamp(0, UV_SIZE - 1) return x, y def _renderer_surface_cache(renderer): cache = getattr(renderer, "_dense_uv_surface_cache", None) if cache is None: cache = {} setattr(renderer, "_dense_uv_surface_cache", cache) return cache def _surface_metadata(grids, is_outer, lookups): uv_px = _grid_to_uv01(grids) * (UV_SIZE - 1) rounded_uv_px = uv_px.round() x = rounded_uv_px[..., 0].long().clamp(0, UV_SIZE - 1) y = rounded_uv_px[..., 1].long().clamp(0, UV_SIZE - 1) center_distance_sq = (uv_px - rounded_uv_px).square().sum(dim=-1) # A bilinear renderer is exact at integer UV centers and increasingly # mixes neighboring texels toward a cell boundary. Keep a normalized # quality score for inverse color sampling; geometry/UV ownership still # uses the same rounded integer coordinates as before. texel_center_score = (1.0 - 2.0 * center_distance_sq).clamp(0.0, 1.0) inner_part = lookups["inner_part"][y, x] inner_face = lookups["inner_face"][y, x] outer_part = lookups["outer_part"][y, x] outer_face = lookups["outer_face"][y, x] return { "flat_uv": y * UV_SIZE + x, "layer": is_outer.long(), "part": torch.where(is_outer, outer_part, inner_part), "face": torch.where(is_outer, outer_face, inner_face), "outer_part": outer_part, "texel_center_score": texel_center_score, } def build_static_surface_routing(renderer, view, device): """Build fixed UV candidates matching the renderer's direct/composite geometry paths.""" device = torch.device(device) cache = _renderer_surface_cache(renderer) cache_key = (view, device.type, device.index) if cache_key in cache: return cache[cache_key] inner_grid = getattr(renderer, f"{view}_inner_grid").to(device=device) outer_grid = getattr(renderer, f"{view}_outer_grid").to(device=device) inner_mask = getattr(renderer, f"{view}_inner_mask").to(device=device).bool() outer_mask = getattr(renderer, f"{view}_outer_mask").to(device=device).bool() H, W = inner_mask.shape lookups = build_part_face_lookups(device=device) grids = [inner_grid, outer_grid] masks = [inner_mask, outer_mask] is_outer = [torch.zeros((H, W), dtype=torch.bool, device=device), torch.ones((H, W), dtype=torch.bool, device=device)] composite_count = 0 geometry_count = 0 composite_grid_name = f"{view}_composite_grid_layers" composite_decor_name = f"{view}_composite_is_decor_layers" if hasattr(renderer, composite_grid_name) and hasattr(renderer, composite_decor_name): composite_grids = getattr(renderer, composite_grid_name).to(device=device) composite_masks = getattr(renderer, f"{view}_composite_mask_layers").to(device=device).bool() composite_decor = getattr(renderer, composite_decor_name).to(device=device).bool() composite_count = composite_grids.shape[0] grids.extend(composite_grids.unbind(0)) masks.extend(composite_masks.unbind(0)) is_outer.extend(composite_decor.unbind(0)) geometry_grid_name = f"{view}_geometry_grid_layers" if hasattr(renderer, geometry_grid_name): geometry_grids = getattr(renderer, geometry_grid_name).to(device=device) geometry_masks = getattr(renderer, f"{view}_geometry_mask_layers").to(device=device).bool() geometry_order = getattr(renderer, f"{view}_geometry_sort_indices").to(device=device).long() geometry_count = geometry_grids.shape[0] if geometry_count > 0: geometry_order = geometry_order.clamp(0, geometry_count - 1) sorted_grids = torch.gather( geometry_grids, 0, geometry_order.unsqueeze(-1).expand(-1, -1, -1, 2), ) sorted_masks = torch.gather(geometry_masks, 0, geometry_order) geometry_metadata = _surface_metadata( sorted_grids, torch.zeros_like(sorted_masks), lookups, ) geometry_is_outer = geometry_metadata["outer_part"] != IGNORE_INDEX grids.extend(sorted_grids.unbind(0)) masks.extend(sorted_masks.unbind(0)) is_outer.extend(geometry_is_outer.unbind(0)) surface_grids = torch.stack(grids, dim=0) surface_masks = torch.stack(masks, dim=0) surface_is_outer = torch.stack(is_outer, dim=0) metadata = _surface_metadata(surface_grids, surface_is_outer, lookups) metadata["layer_face"] = combine_layer_face(metadata["layer"], metadata["face"]) static = { "grids": surface_grids, "masks": surface_masks, "flat_uv": metadata["flat_uv"], "layer": metadata["layer"], "part": metadata["part"], "face": metadata["face"], "layer_face": metadata["layer_face"], "texel_center_score": metadata["texel_center_score"], "direct_count": 2, "composite_count": composite_count, "geometry_count": geometry_count, } cache[cache_key] = static return static def surface_class_count(renderer, views): device = renderer.bg_color.device return max(build_static_surface_routing(renderer, view, device)["grids"].shape[0] for view in parse_views(views)) def _select_static_surface(static, surface): B, H, W = surface.shape surface_count = static["grids"].shape[0] in_range = (surface >= 0) & (surface < surface_count) safe_surface = surface.clamp(0, max(surface_count - 1, 0)) selected = {} for name in ( "masks", "flat_uv", "layer", "part", "face", "texel_center_score", ): values = static[name].unsqueeze(0).expand(B, -1, -1, -1) selected[name] = values.gather(1, safe_surface.unsqueeze(1)).squeeze(1) selected["valid"] = in_range & selected["masks"] return selected def _sample_surface_alpha(skins, static): B = skins.shape[0] surface_count, H, W, _ = static["grids"].shape grids = static["grids"].unsqueeze(0).expand(B, -1, -1, -1, -1).reshape(B * surface_count, H, W, 2) alpha = skins[:, 3:4].unsqueeze(1).expand(-1, surface_count, -1, -1, -1) alpha = alpha.reshape(B * surface_count, 1, skins.shape[-2], skins.shape[-1]) sampled = F.grid_sample(alpha, grids, mode="bilinear", padding_mode="zeros", align_corners=True) return sampled.reshape(B, surface_count, H, W) * static["masks"].unsqueeze(0) def classify_route_role(static, selected_layer, selected_flat_uv, valid): """Classify visible pixels as direct inner/outer surfaces or deeper backfaces.""" B, H, W = selected_layer.shape safe_layer = selected_layer.clamp(0, LAYER_CLASSES - 1) direct_masks = static["masks"][:LAYER_CLASSES].unsqueeze(0).expand(B, -1, H, W) direct_uv = static["flat_uv"][:LAYER_CLASSES].unsqueeze(0).expand(B, -1, H, W) direct_valid = direct_masks.gather(1, safe_layer.unsqueeze(1)).squeeze(1) expected_uv = direct_uv.gather(1, safe_layer.unsqueeze(1)).squeeze(1) primary = valid & direct_valid & (selected_flat_uv == expected_uv) route_role = torch.full_like(selected_layer, IGNORE_INDEX) route_role[primary] = selected_layer[primary] route_role[valid & ~primary] = ROUTE_SECONDARY return route_role def build_dense_parser_batch(skins, renderer, view, alpha_threshold=0.5): """Render skins and label the exact fixed surface that produced every pixel.""" device = skins.device dtype = skins.dtype B = skins.shape[0] rendered = renderer.forward_view(skins, view) static = build_static_surface_routing(renderer, view, device) _, _, H, W = rendered.shape sampled_alpha = _sample_surface_alpha(skins, static) surface = torch.full((B, H, W), IGNORE_INDEX, dtype=torch.long, device=device) direct_visible = sampled_alpha[:, :2] > alpha_threshold inner_visible = direct_visible[:, 0] & ~direct_visible[:, 1] outer_visible = direct_visible[:, 1] surface[inner_visible] = 0 surface[outer_visible] = 1 composite_start = static["direct_count"] composite_count = static["composite_count"] composite_visible = None if composite_count > 0: composite_visible = sampled_alpha[:, composite_start : composite_start + composite_count] > alpha_threshold composite_indices = torch.arange(composite_count, device=device).view(1, composite_count, 1, 1) first_composite = torch.where( composite_visible, composite_indices, torch.full_like(composite_indices, composite_count), ).amin(dim=1) has_composite = first_composite < composite_count composite_surface = composite_start + first_composite.clamp(max=composite_count - 1) composite_route = _select_static_surface(static, composite_surface) trust_composite = has_composite & ( (composite_route["layer"] == 1) | (first_composite <= 1) ) surface[trust_composite] = composite_surface[trust_composite] geometry_start = composite_start + composite_count geometry_count = static["geometry_count"] if geometry_count > 0: geometry_visible = sampled_alpha[:, geometry_start : geometry_start + geometry_count] > alpha_threshold geometry_indices = torch.arange(geometry_count, device=device).view(1, geometry_count, 1, 1) last_geometry = torch.where( geometry_visible, geometry_indices, torch.full_like(geometry_indices, -1), ).amax(dim=1) has_geometry = last_geometry >= 0 geometry_surface = geometry_start + last_geometry.clamp_min(0) front_decor = static["masks"][composite_start] & (static["layer"][composite_start] == 1) geometry_fallback = ( front_decor.unsqueeze(0) & ~composite_visible[:, 0] & has_geometry ) surface[geometry_fallback] = geometry_surface[geometry_fallback] selected = _select_static_surface(static, surface) valid = selected["valid"] layer = torch.where(valid, selected["layer"], torch.full_like(surface, IGNORE_INDEX)) route_role = classify_route_role(static, selected["layer"], selected["flat_uv"], valid) part = torch.where(valid, selected["part"], torch.full_like(surface, IGNORE_INDEX)) face = torch.where(valid, selected["face"], torch.full_like(surface, IGNORE_INDEX)) uv = flat_uv_to_uv01(selected["flat_uv"], dtype).masked_fill(~valid.unsqueeze(1), 0.0) targets = { "foreground": valid.unsqueeze(1).to(dtype=dtype), "layer": layer, "route_role": route_role, "part": part, "face": face, "surface": torch.where(valid, surface, torch.full_like(surface, IGNORE_INDEX)), "uv": uv, } return rendered, targets def randomize_render_background(rendered, probability=0.9, bg_color=(128, 128, 128)): """Replace rendered backgrounds while preserving the skin pixels and parser targets. The renderer produces RGBA images composited over a fixed gray background. The parser always receives RGB-style inputs (alpha fixed to one), with a random solid-color background for the requested fraction of samples. """ if rendered.dim() != 4 or rendered.shape[1] != 4: raise ValueError(f"Expected RGBA render tensor as NCHW, got {tuple(rendered.shape)}.") B, _, H, W = rendered.shape device = rendered.device dtype = rendered.dtype alpha = rendered[:, 3:4].clamp(0.0, 1.0) source_bg = rendered.new_tensor(bg_color).view(1, 3, 1, 1) / 255.0 probability = max(0.0, min(float(probability), 1.0)) if probability > 0: active = torch.rand(B, device=device) < probability random_background = torch.rand(B, 3, 1, 1, device=device, dtype=dtype).expand(-1, -1, H, W) background = torch.where(active.view(B, 1, 1, 1), random_background, source_bg.expand(B, -1, H, W)) else: background = source_bg.expand(B, -1, H, W) foreground_rgb = (rendered[:, :3] - (1.0 - alpha) * source_bg) / alpha.clamp_min(1e-4) composited_rgb = (alpha * foreground_rgb + (1.0 - alpha) * background).clamp(0.0, 1.0) return torch.cat([composited_rgb, torch.ones_like(alpha)], dim=1) def estimate_solid_background_color( rendered, color_tolerance=SOLID_BACKGROUND_COLOR_TOLERANCE, min_corner_support=SOLID_BACKGROUND_MIN_CORNER_SUPPORT, ): if rendered.dim() != 4 or rendered.shape[1] < 3: raise ValueError(f"Expected NCHW RGB(A), got {tuple(rendered.shape)}.") if not 0.0 <= color_tolerance <= 1.0: raise ValueError("color_tolerance must be in [0, 1].") if not 0.0 <= min_corner_support <= 1.0: raise ValueError("min_corner_support must be in [0, 1].") rgb = rendered[:, :3].float() batch, _, height, width = rgb.shape patch = max(2, min(32, min(height, width) // 16)) corners = torch.cat( [ rgb[:, :, :patch, :patch].flatten(2), rgb[:, :, :patch, -patch:].flatten(2), rgb[:, :, -patch:, :patch].flatten(2), rgb[:, :, -patch:, -patch:].flatten(2), ], dim=2, ) background_rgb = corners.median(dim=2).values.view(batch, 3, 1, 1) corner_distance = (corners - background_rgb.flatten(2)).abs().amax(dim=1) solid_background = (corner_distance <= color_tolerance).float().mean(dim=1) >= min_corner_support return background_rgb, solid_background def estimate_solid_background_foreground( rendered, color_tolerance=SOLID_BACKGROUND_COLOR_TOLERANCE, min_corner_support=SOLID_BACKGROUND_MIN_CORNER_SUPPORT, ): """Find non-background pixels without deleting matching colors enclosed by the subject.""" if min(rendered.shape[-2:]) < 8: return torch.ones( rendered.shape[0], *rendered.shape[-2:], dtype=torch.bool, device=rendered.device, ) rgb = rendered[:, :3].float() batch, _, height, width = rgb.shape background_rgb, solid_background = estimate_solid_background_color( rendered, color_tolerance=color_tolerance, min_corner_support=min_corner_support, ) background_candidate = (rgb - background_rgb).abs().amax(dim=1) <= color_tolerance candidate_u8 = background_candidate.to(torch.uint8) connected_to_border = ( candidate_u8.cumprod(dim=2).bool() | candidate_u8.flip(2).cumprod(dim=2).flip(2).bool() | candidate_u8.cumprod(dim=1).bool() | candidate_u8.flip(1).cumprod(dim=1).flip(1).bool() ) connected_to_border &= solid_background.view(batch, 1, 1) max_steps = max(height, width) for _ in range(max_steps): expanded = F.max_pool2d( connected_to_border.unsqueeze(1).float(), 3, stride=1, padding=1 )[:, 0] > 0.0 updated = connected_to_border | (background_candidate & expanded) if torch.equal(updated, connected_to_border): break connected_to_border = updated return ~connected_to_border def estimate_top_left_flood_foreground(rendered, color_tolerance=8.0 / 255.0): """Remove only seed-colored pixels connected to the top-left image pixel. Each image in the batch owns its own RGB seed. Pixels with the same color elsewhere inside the character remain foreground unless a path of seed-colored pixels connects them to the top-left corner. """ if rendered.dim() != 4 or rendered.shape[1] < 3: raise ValueError(f"Expected NCHW RGB(A), got {tuple(rendered.shape)}.") if not 0.0 <= color_tolerance <= 1.0: raise ValueError("color_tolerance must be in [0, 1].") rgb = rendered[:, :3].float() seed_rgb = rgb[:, :, :1, :1] background_candidate = ( (rgb - seed_rgb).abs().amax(dim=1) <= color_tolerance ) connected = torch.zeros_like(background_candidate) connected[:, 0, 0] = background_candidate[:, 0, 0] height, width = rendered.shape[-2:] for _ in range(height + width): # Four-connected expansion cannot leak diagonally through a one-pixel # corner where two foreground blocks merely touch. expanded = connected.clone() expanded[:, 1:, :] |= connected[:, :-1, :] expanded[:, :-1, :] |= connected[:, 1:, :] expanded[:, :, 1:] |= connected[:, :, :-1] expanded[:, :, :-1] |= connected[:, :, 1:] updated = connected | (background_candidate & expanded) if torch.equal(updated, connected): break connected = updated return ~connected def build_color_sampling_support( rendered, observed_foreground, background_rgb, background_tolerance=None, foreground_inset=1, ): """Separate route occupancy from source pixels that are safe for RGB pickup. Flood fill intentionally keeps colors that merely resemble the background when they are not connected to the image border. That is correct for occupancy, but an isolated background pocket or antialiased silhouette fringe must not beat a real character sample during inverse UV splatting. Only background-like *boundary* candidates are rejected. Interior pixels remain valid even when a skin intentionally uses the background color. """ if rendered.dim() != 4 or rendered.shape[1] < 3: raise ValueError(f"Expected NCHW RGB(A), got {tuple(rendered.shape)}.") expected = rendered.shape[:1] + rendered.shape[-2:] if observed_foreground.shape != expected: raise ValueError( f"Expected observed_foreground shape {expected}, got " f"{tuple(observed_foreground.shape)}." ) if background_rgb.shape not in ((rendered.shape[0], 3), (rendered.shape[0], 3, 1, 1)): raise ValueError( "background_rgb must have shape (N, 3) or (N, 3, 1, 1)." ) if background_tolerance is not None and not 0.0 <= background_tolerance <= 1.0: raise ValueError("background_tolerance must be in [0, 1].") if foreground_inset < 0: raise ValueError("foreground_inset must be non-negative.") observed = observed_foreground.to(device=rendered.device, dtype=torch.bool) interior = observed for _ in range(int(foreground_inset)): adjacent_background = F.max_pool2d( (~interior).unsqueeze(1).float(), kernel_size=3, stride=1, padding=1, )[:, 0] > 0.0 interior = interior & ~adjacent_background background = background_rgb.to( device=rendered.device, dtype=rendered.dtype ).reshape(rendered.shape[0], 3, 1, 1) background_like = ( (rendered[:, :3] - background).abs().amax(dim=1) <= float(background_tolerance) if background_tolerance is not None else torch.zeros_like(observed) ) boundary = observed & ~interior rejected = boundary & background_like return { "valid": observed & ~rejected, "interior": interior, "background_like": background_like, "rejected": rejected, } def affine_to_canonical_grid(affine, output_shape): """Build a grid that samples an augmented render at canonical coordinates.""" if affine.dim() != 2 or affine.shape[1] != 3: raise ValueError(f"Expected affine shape (N, 3), got {tuple(affine.shape)}.") if affine.shape[0] != output_shape[0]: raise ValueError( f"Affine batch size {affine.shape[0]} does not match output batch size {output_shape[0]}." ) scale = affine[:, 2].exp() theta = affine.new_zeros(affine.shape[0], 2, 3) theta[:, 0, 0] = scale theta[:, 1, 1] = scale theta[:, 0, 2] = affine[:, 0] theta[:, 1, 2] = affine[:, 1] return F.affine_grid(theta, output_shape, align_corners=False) def canonicalize_tensor(tensor, affine, mode="bilinear"): """Undo the parser's predicted global translation/scale for an NCHW tensor.""" affine = affine.to(device=tensor.device, dtype=tensor.dtype) grid = affine_to_canonical_grid(affine, tensor.shape) return F.grid_sample( tensor, grid, mode=mode, padding_mode="zeros", align_corners=False, ) def canonicalize_foreground_mask(mask, affine, min_coverage=1e-3): """Warp a foreground mask without dropping thin silhouette coverage. Nearest-neighbour sampling can erase a one-pixel hat brim or outer-layer protrusion after even a small predicted affine scale. Bilinear coverage with a low positive threshold keeps every output pixel touched by the transformed source mask. Later color/background and minimum-source filters remain responsible for rejecting unsupported edge colors. """ if mask.dim() == 3: mask = mask.unsqueeze(1) if mask.dim() != 4 or mask.shape[1] != 1: raise ValueError( "Foreground mask must have shape NHW or N1HW, got " f"{tuple(mask.shape)}." ) if not 0.0 < float(min_coverage) <= 1.0: raise ValueError("min_coverage must be in (0, 1].") coverage = canonicalize_tensor( mask.to(device=affine.device, dtype=torch.float32), affine, mode="bilinear", ) return coverage[:, 0] >= float(min_coverage) def canonicalize_index_tensor(index, affine): """Undo a global transform while preserving IGNORE_INDEX outside valid labels.""" valid = (index != IGNORE_INDEX).unsqueeze(1).float() values = index.masked_fill(index == IGNORE_INDEX, 0).float().unsqueeze(1) affine = affine.to(device=values.device, dtype=values.dtype) grid = affine_to_canonical_grid(affine, values.shape) sampled_valid = F.grid_sample( valid, grid, mode="nearest", padding_mode="zeros", align_corners=False, ).squeeze(1) > 0.5 sampled_values = F.grid_sample( values, grid, mode="nearest", padding_mode="zeros", align_corners=False, ).squeeze(1).round().long() sampled_values[~sampled_valid] = IGNORE_INDEX return sampled_values def canonicalize_dense_targets(targets): """Return parser targets in the fixed renderer coordinate system.""" if "affine" not in targets: raise ValueError("Global-affine targets require an 'affine' tensor.") affine = targets["affine"] canonical = { "foreground": canonicalize_tensor(targets["foreground"], affine, mode="nearest"), "layer": canonicalize_index_tensor(targets["layer"], affine), "route_role": canonicalize_index_tensor(targets["route_role"], affine), "part": canonicalize_index_tensor(targets["part"], affine), "face": canonicalize_index_tensor(targets["face"], affine), "surface": canonicalize_index_tensor(targets["surface"], affine), "uv": canonicalize_tensor(targets["uv"], affine, mode="nearest"), "affine": affine, } valid = (canonical["layer"] != IGNORE_INDEX).unsqueeze(1) canonical["foreground"] = canonical["foreground"] * valid.to(canonical["foreground"].dtype) canonical["uv"] = torch.where(valid, canonical["uv"], torch.zeros_like(canonical["uv"])) return canonical def canonicalize_parser_outputs(outputs): """Warp dense parser logits into the fixed renderer coordinate system.""" if "affine" not in outputs: return outputs affine = outputs["affine"] canonical = {} for name, value in outputs.items(): if ( name in ( "affine", "route_role_prior_raw", "outer_uv_occupancy_logits", "outer_uv_occupancy_support", ) or name == "outer_uv_global_context" or not torch.is_tensor(value) or value.dim() != 4 ): canonical[name] = value else: canonical[name] = canonicalize_tensor(value, affine, mode="bilinear") return canonical def canonicalize_parser_render(rendered, outputs, mode="nearest", fill_color=None): """Undo the global transform for colors without blending Minecraft texels.""" if "affine" not in outputs: return rendered if fill_color is None: return canonicalize_tensor(rendered, outputs["affine"], mode=mode) if fill_color.dim() == 2: fill_color = fill_color.unsqueeze(-1).unsqueeze(-1) if fill_color.shape != (rendered.shape[0], rendered.shape[1], 1, 1): raise ValueError( "fill_color must have shape " f"{(rendered.shape[0], rendered.shape[1], 1, 1)}, got {tuple(fill_color.shape)}." ) fill_color = fill_color.to(device=rendered.device, dtype=rendered.dtype) return canonicalize_tensor( rendered - fill_color, outputs["affine"], mode=mode ) + fill_color def _mask_moments(weights): weights = weights.float() N, H, W = weights.shape y = torch.arange(H, device=weights.device, dtype=weights.dtype).view(1, H, 1) x = torch.arange(W, device=weights.device, dtype=weights.dtype).view(1, 1, W) total = weights.sum(dim=(1, 2)).clamp_min(1e-6) mean_x = (weights * x).sum(dim=(1, 2)) / total mean_y = (weights * y).sum(dim=(1, 2)) / total var_x = (weights * (x - mean_x.view(N, 1, 1)).square()).sum(dim=(1, 2)) / total var_y = (weights * (y - mean_y.view(N, 1, 1)).square()).sum(dim=(1, 2)) / total return total, mean_x, mean_y, var_x.clamp_min(1e-6).sqrt(), var_y.clamp_min(1e-6).sqrt() def _soft_mask_dice(probability, target, support): probability = probability.float() * support.float() target = target.float() intersection = (probability * target).sum(dim=(1, 2, 3)) denominator = probability.sum(dim=(1, 2, 3)) + target.sum(dim=(1, 2, 3)) return (2.0 * intersection + 1e-6) / (denominator + 1e-6) def refine_parser_affine( outputs, renderer, views, translation_radius_px=8.0, scale_radius=0.0, translation_step_px=0.5, observed_foreground=None, ): """Refine parser affine predictions against the fixed inner-layer silhouette.""" if "affine" not in outputs or "foreground" not in outputs: return outputs.get("affine"), None if translation_radius_px < 0 or scale_radius < 0: raise ValueError("Affine refinement radii must be non-negative.") if translation_step_px <= 0: raise ValueError("translation_step_px must be positive.") views = parse_views(views) affine = outputs["affine"] if observed_foreground is None: foreground_prob = torch.sigmoid(outputs["foreground"]) else: if observed_foreground.shape != outputs["foreground"].shape[:1] + outputs["foreground"].shape[-2:]: raise ValueError( "observed_foreground must match parser spatial output shape, got " f"{tuple(observed_foreground.shape)}." ) foreground_prob = observed_foreground.to( device=outputs["foreground"].device, dtype=outputs["foreground"].dtype, ).unsqueeze(1) N, _, H, W = foreground_prob.shape if not views or N % len(views) != 0: raise ValueError(f"N={N} must be divisible by the number of views ({len(views)}).") target = foreground_prob.new_zeros(N, 1, H, W) for view_index, view in enumerate(views): static = build_static_surface_routing(renderer, view, foreground_prob.device) inner_mask = static["masks"][0] if inner_mask.shape != (H, W): raise ValueError( f"View {view!r} mapping shape {tuple(inner_mask.shape)} does not match parser shape {(H, W)}." ) target[view_index::len(views), 0] = inner_mask.to(dtype=target.dtype) support_radius = max(math.ceil(translation_radius_px) + 2, 2) support = F.max_pool2d( target, kernel_size=support_radius * 2 + 1, stride=1, padding=support_radius, ).clamp(0.0, 1.0) base_canonical = canonicalize_tensor(foreground_prob, affine, mode="bilinear") base_score = _soft_mask_dice(base_canonical, target, support) target_stats = _mask_moments(target[:, 0]) base_stats = _mask_moments(base_canonical[:, 0] * support[:, 0]) target_total, _, _, target_std_x, target_std_y = target_stats base_total, _, _, base_std_x, base_std_y = base_stats enough_foreground = base_total > target_total * 0.25 if scale_radius > 0: scale_delta = 0.5 * ( (base_std_x / target_std_x).clamp_min(1e-6).log() + (base_std_y / target_std_y).clamp_min(1e-6).log() ) scale_delta = scale_delta.clamp(-scale_radius, scale_radius) scale_affine = affine.clone() scale_affine[:, 2] = scale_affine[:, 2] + scale_delta.to(dtype=scale_affine.dtype) scale_canonical = canonicalize_tensor(foreground_prob, scale_affine, mode="bilinear") else: scale_delta = torch.zeros_like(base_score) scale_affine = affine scale_canonical = base_canonical _, pred_mean_x, pred_mean_y, _, _ = _mask_moments(scale_canonical[:, 0] * support[:, 0]) _, target_mean_x, target_mean_y, _, _ = target_stats delta_x_px = (pred_mean_x - target_mean_x).clamp(-translation_radius_px, translation_radius_px) delta_y_px = (pred_mean_y - target_mean_y).clamp(-translation_radius_px, translation_radius_px) delta_x_px = (delta_x_px / translation_step_px).round() * translation_step_px delta_y_px = (delta_y_px / translation_step_px).round() * translation_step_px candidate_affine = scale_affine.clone() candidate_affine[:, 0] = candidate_affine[:, 0] + (2.0 * delta_x_px / W).to(candidate_affine.dtype) candidate_affine[:, 1] = candidate_affine[:, 1] + (2.0 * delta_y_px / H).to(candidate_affine.dtype) candidate_canonical = canonicalize_tensor(foreground_prob, candidate_affine, mode="bilinear") candidate_score = _soft_mask_dice(candidate_canonical, target, support) accepted = enough_foreground & (candidate_score > base_score + 1e-5) refined_affine = torch.where(accepted.unsqueeze(1), candidate_affine, affine) accepted_float = accepted.to(dtype=delta_x_px.dtype) details = { "accepted": accepted, "translation_px": torch.stack( [delta_x_px * accepted_float, delta_y_px * accepted_float], dim=1, ), "scale_percent": (scale_delta.exp() - 1.0) * 100.0 * accepted_float, "score_before": base_score, "score_after": torch.where(accepted, candidate_score, base_score), } return refined_affine, details def _routing_from_affine_outputs(renderer, views, outputs, fg_threshold=0.5, semantic_gate=True): views = parse_views(views) if not views: raise ValueError("At least one renderer view is required for deterministic UV routing.") if "surface" not in outputs: raise ValueError("Affine parser outputs must include the static surface classifier.") foreground_prob = torch.sigmoid(outputs["foreground"])[:, 0] surface_prob = torch.softmax(outputs["surface"], dim=1) layer_prob = torch.softmax(outputs["layer"], dim=1) N, _, H, W = surface_prob.shape if N % len(views) != 0: raise ValueError(f"N={N} must be divisible by {len(views)} renderer views.") fg = torch.zeros_like(foreground_prob, dtype=torch.bool) layer = torch.zeros((N, H, W), dtype=torch.long, device=surface_prob.device) flat_uv = torch.zeros_like(layer) surface = torch.full_like(layer, IGNORE_INDEX) confidence = torch.zeros_like(foreground_prob) confidence_margin = torch.zeros_like(foreground_prob) confidence_margin_ratio = torch.zeros_like(foreground_prob) texel_center_score = torch.zeros_like(foreground_prob) semantic_fallback = torch.zeros_like(foreground_prob, dtype=torch.bool) expected_part = torch.full_like(layer, IGNORE_INDEX) expected_face = torch.full_like(layer, IGNORE_INDEX) for view_index, view in enumerate(views): selection = slice(view_index, N, len(views)) static = build_static_surface_routing(renderer, view, surface_prob.device) if static["masks"].shape[-2:] != (H, W): raise ValueError( f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} does not match parser shape {(H, W)}." ) view_surface_prob = surface_prob[selection] view_batch = view_surface_prob.shape[0] surface_count = static["masks"].shape[0] if view_surface_prob.shape[1] < surface_count: raise ValueError( f"Parser has {view_surface_prob.shape[1]} surface classes, but view {view!r} requires {surface_count}." ) candidate_score = view_surface_prob[:, :surface_count] physical_valid = static["masks"].unsqueeze(0).expand(view_batch, -1, -1, -1) candidate_valid = physical_valid.clone() semantic_score = torch.ones_like(candidate_score) # Re-rank only surfaces that physically exist at this screen pixel. This # converts the semantic heads from a rejection gate into useful routing # evidence and avoids holes caused by a globally invalid surface argmax. semantic_names = ( ("part", "layer_face") if "layer_face" in outputs else ("layer", "part", "face") ) for name in semantic_names: if name not in outputs: continue probabilities = torch.softmax(outputs[name][selection], dim=1) expected = static[name] known_semantic = expected != IGNORE_INDEX expected_safe = expected.clamp(0, probabilities.shape[1] - 1) expected_probability = probabilities.gather( 1, expected_safe.unsqueeze(0).expand(view_batch, -1, -1, -1), ) semantic_score = semantic_score * torch.where( known_semantic.unsqueeze(0), expected_probability, torch.ones_like(expected_probability), ) if semantic_gate: candidate_valid = candidate_valid & ( ~known_semantic.unsqueeze(0) | (probabilities.argmax(dim=1).unsqueeze(1) == expected_safe.unsqueeze(0)) ) candidate_score = candidate_score * semantic_score.clamp_min(1e-12).sqrt() if "uv_x" in outputs and "uv_y" in outputs: uv_x_prob = torch.softmax(outputs["uv_x"][selection], dim=1) uv_y_prob = torch.softmax(outputs["uv_y"][selection], dim=1) candidate_x = static["flat_uv"].remainder(UV_SIZE) candidate_y = static["flat_uv"].div(UV_SIZE, rounding_mode="floor") candidate_x_prob = uv_x_prob.gather( 1, candidate_x.unsqueeze(0).expand(view_batch, -1, -1, -1), ) candidate_y_prob = uv_y_prob.gather( 1, candidate_y.unsqueeze(0).expand(view_batch, -1, -1, -1), ) candidate_score = candidate_score * ( candidate_x_prob * candidate_y_prob ).clamp_min(1e-12).sqrt() gated_score = candidate_score.masked_fill(~candidate_valid, -1.0) has_gated_candidate = candidate_valid.any(dim=1, keepdim=True) if semantic_gate: fallback_score = candidate_score.masked_fill(~physical_valid, -1.0) candidate_score = torch.where(has_gated_candidate, gated_score, fallback_score) semantic_fallback[selection] = ( ~has_gated_candidate.squeeze(1) & physical_valid.any(dim=1) ) else: candidate_score = gated_score best_score, selected_surface = candidate_score.max(dim=1) if surface_count > 1: top_scores = candidate_score.topk(k=2, dim=1).values route_margin = (top_scores[:, 0] - top_scores[:, 1]).clamp_min(0.0) route_margin_ratio = ( route_margin / top_scores[:, 0].abs().clamp_min(1e-8) ).clamp(0.0, 1.0) else: route_margin = best_score.clamp_min(0.0) route_margin_ratio = (best_score > 0.0).to(best_score.dtype) routed = _select_static_surface(static, selected_surface) selected_layer = routed["layer"] selected_part = routed["part"] selected_face = routed["face"] has_candidate = best_score >= 0.0 base_silhouette = static["masks"][0].unsqueeze(0).expand(view_batch, -1, -1) routed_fg = ( ((foreground_prob[selection] > fg_threshold) | base_silhouette) & has_candidate & routed["valid"] ) fg[selection] = routed_fg layer[selection] = selected_layer flat_uv[selection] = routed["flat_uv"] surface[selection] = torch.where( has_candidate, selected_surface, torch.full_like(selected_surface, IGNORE_INDEX), ) confidence[selection] = foreground_prob[selection] * best_score.clamp_min(0.0) confidence_margin[selection] = route_margin confidence_margin_ratio[selection] = route_margin_ratio expected_part[selection] = selected_part expected_face[selection] = selected_face texel_center_score[selection] = routed["texel_center_score"] return { "foreground": fg, "layer": layer, "flat_uv": flat_uv, "surface": surface, "confidence": confidence, "confidence_margin": confidence_margin, "confidence_margin_ratio": confidence_margin_ratio, "semantic_fallback": semantic_fallback, "part": expected_part, "face": expected_face, "texel_center_score": texel_center_score, } def _aggregate_role_probabilities_by_direct_texel( role_prob, static, foreground_prob=None, center_power=2.0, ): """Center-weight route probabilities within each projected Minecraft texel.""" role_prob = role_prob.float() batch, role_count, height, width = role_prob.shape if foreground_prob is not None and foreground_prob.shape != ( batch, height, width, ): raise ValueError( "foreground_prob must match the route map, got " f"{tuple(foreground_prob.shape)} and {(batch, height, width)}." ) if center_power < 0: raise ValueError("center_power must be non-negative.") aggregated = [] for layer_index in range(LAYER_CLASSES): valid = static["masks"][layer_index].reshape(1, 1, -1) flat_uv = static["flat_uv"][layer_index].reshape(1, 1, -1) center_weight = static["texel_center_score"][layer_index].float() center_weight = center_weight.clamp(0.0, 1.0).pow(center_power) pixel_weight = center_weight.reshape(1, 1, -1) * valid if foreground_prob is not None: pixel_weight = ( pixel_weight * foreground_prob.float().clamp(0.0, 1.0).reshape(batch, 1, -1) ) pixel_weight = pixel_weight.expand(batch, 1, -1) sums = role_prob.new_zeros(batch, role_count, UV_SIZE * UV_SIZE) sums.scatter_add_( 2, flat_uv.expand(batch, role_count, -1), role_prob.flatten(2) * pixel_weight, ) counts = role_prob.new_zeros(batch, 1, UV_SIZE * UV_SIZE) counts.scatter_add_( 2, flat_uv.expand(batch, 1, -1), pixel_weight, ) means = sums / counts.clamp_min(1e-6) aggregated.append( means.gather(2, flat_uv.expand(batch, role_count, -1)).reshape( batch, role_count, height, width ) ) return torch.stack(aggregated, dim=1) def _project_outer_uv_occupancy_to_view( occupancy_probability, static, view_batch, ): """Gather a grouped 64x64 outer-alpha prior into one projected view.""" if occupancy_probability is None: return None if occupancy_probability.shape != (view_batch, 1, UV_SIZE, UV_SIZE): raise ValueError( "outer_uv_occupancy probability must be grouped as " f"{(view_batch, 1, UV_SIZE, UV_SIZE)}, got " f"{tuple(occupancy_probability.shape)}." ) probability = occupancy_probability.float().flatten(2) flat_uv = static["flat_uv"][ROUTE_OUTER_PRIMARY].reshape( 1, 1, -1 ).expand(view_batch, 1, -1) projected = probability.gather(2, flat_uv).reshape( view_batch, *static["flat_uv"].shape[-2:] ) outer_valid = static["masks"][ROUTE_OUTER_PRIMARY].unsqueeze(0) return torch.where(outer_valid, projected, torch.zeros_like(projected)) def _outer_graph_component_labels(candidate, edge_index, max_steps=64): """Label candidate outer nodes by bounded min-label propagation.""" batch, node_count = candidate.shape labels = torch.arange( node_count, device=candidate.device, dtype=torch.long, ).view(1, -1).expand(batch, -1).clone() labels[~candidate] = node_count source, target = edge_index expanded_target = target.view(1, -1).expand(batch, -1) connected_edge = candidate[:, source] & candidate[:, target] for _ in range(int(max_steps)): neighbour = torch.where( connected_edge, labels[:, source], torch.full_like(labels[:, source], node_count), ) updated = labels.clone() updated.scatter_reduce_( 1, expanded_target, neighbour, reduce="amin", include_self=True, ) labels = updated return labels def outer_uv_topology_hysteresis( probability, seed_threshold=0.80, grow_threshold=0.50, min_component_size=2, max_steps=64, ): """Grow high-confidence outer seeds through topology-supported candidates.""" if probability.dim() == 3: probability = probability.unsqueeze(1) if probability.dim() != 4 or probability.shape[1:] != ( 1, UV_SIZE, UV_SIZE, ): raise ValueError("Outer occupancy probability must be Bx1x64x64.") if not 0.0 <= float(grow_threshold) <= float(seed_threshold) <= 1.0: raise ValueError( "Require 0 <= grow_threshold <= seed_threshold <= 1." ) if min_component_size < 1: raise ValueError("min_component_size must be positive.") flat_indices, edge_index = build_outer_uv_graph() flat_indices = flat_indices.to(probability.device) edge_index = edge_index.to(probability.device) nodes = probability.float().flatten(2)[:, 0].index_select( 1, flat_indices ) candidate = nodes >= float(grow_threshold) seeds = nodes >= float(seed_threshold) source, target = edge_index expanded_target = target.view(1, -1).expand(nodes.shape[0], -1) active = seeds.clone() connected_edge = candidate[:, source] & candidate[:, target] for _ in range(int(max_steps)): neighbour_active = torch.where( connected_edge, active[:, source], torch.zeros_like(active[:, source]), ) propagated_active = torch.zeros_like(active, dtype=torch.uint8) propagated_active.scatter_reduce_( 1, expanded_target, neighbour_active.to(torch.uint8), reduce="amax", include_self=False, ) active = active | (candidate & propagated_active.bool()) labels = _outer_graph_component_labels( candidate, edge_index, max_steps=max_steps, ) node_count = nodes.shape[1] batch_offsets = ( torch.arange(nodes.shape[0], device=nodes.device).view(-1, 1) * (node_count + 1) ) component_index = labels + batch_offsets component_sizes = nodes.new_zeros(nodes.shape[0] * (node_count + 1)) component_sizes.scatter_add_( 0, component_index[candidate], torch.ones_like(nodes[candidate]), ) node_component_size = component_sizes[component_index] active = active & (node_component_size >= int(min_component_size)) grown = active & ~seeds # Connectivity determines eligibility, not confidence. Propagating a # single seed's score across the whole component made a 0.50 neighbour # behave like a 0.90 texel and amplified one false seed into a large outer # patch. Accepted nodes therefore retain their own calibrated score. structured_nodes = torch.where(active, nodes, torch.zeros_like(nodes)) def scatter_nodes(node_values, dtype=None): atlas = node_values.new_zeros( node_values.shape[0], UV_SIZE * UV_SIZE, dtype=dtype, ) atlas[:, flat_indices] = node_values.to(dtype=atlas.dtype) return atlas.reshape(-1, 1, UV_SIZE, UV_SIZE) return { "probability": scatter_nodes(structured_nodes), "accepted": scatter_nodes(active, dtype=torch.bool), "seed": scatter_nodes(seeds & active, dtype=torch.bool), "grown": scatter_nodes(grown, dtype=torch.bool), "rejected_candidate": scatter_nodes( candidate & ~active, dtype=torch.bool, ), } def aggregate_direct_outer_values_by_view( renderer, views, values, center_power=2.0, ): """Pool image-space values into direct outer UV texels, keeping views separate.""" views = parse_views(views) if not views: raise ValueError("At least one renderer view is required.") if values.dim() != 4: raise ValueError("values must be shaped NCHW.") if values.shape[0] % len(views) != 0: raise ValueError("The value batch must be divisible by the view count.") if center_power < 0.0: raise ValueError("center_power must be non-negative.") group_count = values.shape[0] // len(views) channel_count = values.shape[1] uv_count = UV_SIZE * UV_SIZE pooled_views = [] supported_views = [] for view_index, view in enumerate(views): selection = slice(view_index, values.shape[0], len(views)) selected = values[selection].float() static = build_static_surface_routing(renderer, view, values.device) if static["masks"].shape[-2:] != values.shape[-2:]: raise ValueError( f"View {view!r} mapping shape " f"{tuple(static['masks'].shape[-2:])} does not match " f"value shape {tuple(values.shape[-2:])}." ) valid = static["masks"][ROUTE_OUTER_PRIMARY].reshape(1, 1, -1) flat_uv = static["flat_uv"][ROUTE_OUTER_PRIMARY].reshape(1, 1, -1) center_weight = static["texel_center_score"][ ROUTE_OUTER_PRIMARY ].float().clamp(0.0, 1.0).pow(float(center_power)) weight = center_weight.reshape(1, 1, -1) * valid weight = weight.expand(group_count, 1, -1) indices = flat_uv.expand(group_count, channel_count, -1) sums = selected.new_zeros(group_count, channel_count, uv_count) sums.scatter_add_( 2, indices, selected.flatten(2) * weight, ) counts = selected.new_zeros(group_count, 1, uv_count) counts.scatter_add_( 2, flat_uv.expand(group_count, 1, -1), weight, ) pooled_views.append(sums / counts.clamp_min(1e-6)) supported_views.append(counts[:, 0] > 1e-6) return ( torch.stack(pooled_views, dim=1), torch.stack(supported_views, dim=1), ) def attach_projected_outer_uv_occupancy( model, outputs, renderer, views, observed_foreground=None, center_power=2.0, ): """Project dense parser features into UV space and run the topology head.""" if not getattr(model, "predict_outer_uv_occupancy", False): return outputs if "outer_uv_features" not in outputs: raise ValueError( "Projected occupancy parser did not return outer_uv_features." ) views = parse_views(views) feature_map = outputs["outer_uv_features"] layer_logits = outputs["layer"] foreground = torch.sigmoid(outputs["foreground"].float()) if "affine" in outputs: affine = outputs["affine"] feature_map = canonicalize_tensor(feature_map, affine, mode="bilinear") layer_logits = canonicalize_tensor( layer_logits, affine, mode="bilinear" ) foreground = canonicalize_tensor( foreground, affine, mode="bilinear" ) if observed_foreground is not None: if observed_foreground.dim() == 3: observed_foreground = observed_foreground.unsqueeze(1) observed_foreground = canonicalize_tensor( observed_foreground.float(), affine, mode="bilinear", ) elif observed_foreground is not None and observed_foreground.dim() == 3: observed_foreground = observed_foreground.unsqueeze(1) if observed_foreground is not None: foreground = observed_foreground.to( device=foreground.device, dtype=foreground.dtype, ).clamp(0.0, 1.0) outer_probability = torch.softmax( layer_logits.float(), dim=1 )[:, ROUTE_OUTER_PRIMARY : ROUTE_OUTER_PRIMARY + 1].detach() values = torch.cat( [ feature_map.float(), outer_probability, foreground.float().detach(), ], dim=1, ) pooled, support = aggregate_direct_outer_values_by_view( renderer, views, values, center_power=center_power, ) support_weight = support.unsqueeze(2).to(dtype=pooled.dtype) support_count = support_weight.sum(dim=1).clamp_min(1.0) feature_channels = feature_map.shape[1] pooled_mean = ( (pooled * support_weight).sum(dim=1) / support_count ) mean_features = pooled_mean[:, :feature_channels] mean_outer = pooled_mean[:, feature_channels : feature_channels + 1] mean_foreground = pooled_mean[:, feature_channels + 1 : feature_channels + 2] per_view_outer = pooled[:, :, feature_channels : feature_channels + 1] max_outer = torch.where( support_weight.bool(), per_view_outer, torch.zeros_like(per_view_outer), ).amax(dim=1) route_evidence = torch.cat([mean_outer, max_outer], dim=1) route_evidence_dropout = float( getattr(model, "outer_uv_route_evidence_dropout", 0.0) ) if route_evidence_dropout >= 1.0: # The projected ownership classifier must remain independent from the # already-unreliable per-pixel route prediction. Keep the two channels # for checkpoint compatibility, but disable the shortcut in both # training and inference. route_evidence = torch.zeros_like(route_evidence) elif model.training and route_evidence_dropout > 0.0: # Drop the entire route-evidence pair per grouped skin. This prevents # the occupancy head from learning the shortcut "copy p_outer" while # leaving half of the batches available for useful route context. keep_route_evidence = ( torch.rand( (route_evidence.shape[0],) + (1,) * (route_evidence.ndim - 1), device=route_evidence.device, ) >= route_evidence_dropout ) route_evidence = route_evidence * keep_route_evidence.to( dtype=route_evidence.dtype ) support_fraction = ( support_count / max(float(len(views)), 1.0) ) atlas_features = torch.cat( [ mean_features, route_evidence, mean_foreground, support_fraction, ], dim=1, ).reshape( mean_features.shape[0], feature_channels + 4, UV_SIZE, UV_SIZE, ) logits = model.predict_projected_outer_uv_occupancy( atlas_features, outputs["outer_uv_global_context"], ) outputs["outer_uv_occupancy_logits"] = logits outputs["outer_uv_occupancy_support"] = ( support.any(dim=1).reshape(-1, 1, UV_SIZE, UV_SIZE) ) return outputs def _project_grouped_outer_texels_to_views( renderer, views, texel_values, ): """Project grouped outer-UV values back into every fixed input view.""" views = parse_views(views) if texel_values.dim() != 2 or texel_values.shape[1] != UV_SIZE * UV_SIZE: raise ValueError( "texel_values must be shaped " f"(B, {UV_SIZE * UV_SIZE})." ) projected = [] for view in views: static = build_static_surface_routing( renderer, view, texel_values.device ) flat_uv = static["flat_uv"][ROUTE_OUTER_PRIMARY].reshape(1, -1) gathered = texel_values.gather( 1, flat_uv.expand(texel_values.shape[0], -1) ).reshape(texel_values.shape[0], *static["masks"].shape[-2:]) valid = static["masks"][ROUTE_OUTER_PRIMARY].unsqueeze(0) projected.append(torch.where(valid, gathered, torch.zeros_like(gathered))) return torch.stack(projected, dim=1).reshape( texel_values.shape[0] * len(views), *projected[0].shape[-2:], ) def cross_view_outer_visibility_evidence( renderer, views, role_prob, observed_foreground, center_power=2.0, positive_confidence=0.70, positive_margin=0.20, negative_confidence=0.70, negative_margin=0.20, background_max_coverage=0.25, min_views=2, ): """Find shared outer texels with high-confidence cross-view contradictions. A contradiction requires at least one view that strongly supports outer and another view that either sees background at the projected outer texel or strongly supports inner. Weak predictions, missing mappings, and texels visible from only one view never form a veto. """ for name, value in ( ("positive_confidence", positive_confidence), ("positive_margin", positive_margin), ("negative_confidence", negative_confidence), ("negative_margin", negative_margin), ("background_max_coverage", background_max_coverage), ): if not 0.0 <= float(value) <= 1.0: raise ValueError(f"{name} must be in [0, 1].") if min_views < 2: raise ValueError("min_views must be at least 2.") if observed_foreground.shape != ( role_prob.shape[0], *role_prob.shape[-2:], ): raise ValueError("observed_foreground must match the role probability map.") values = torch.cat( [ role_prob.float(), observed_foreground.unsqueeze(1).float(), ], dim=1, ) pooled, support = aggregate_direct_outer_values_by_view( renderer, views, values, center_power=center_power, ) role_count = role_prob.shape[1] pooled_role = pooled[:, :, :role_count] foreground_coverage = pooled[:, :, role_count].clamp(0.0, 1.0) outer_score = pooled_role[:, :, ROUTE_OUTER_PRIMARY] outer_competitor = torch.maximum( pooled_role[:, :, ROUTE_INNER_PRIMARY], pooled_role[:, :, ROUTE_SECONDARY], ) outer_margin_ratio = ( (outer_score - outer_competitor).clamp_min(0.0) / outer_score.clamp_min(1e-8) ).clamp(0.0, 1.0) inner_score = pooled_role[:, :, ROUTE_INNER_PRIMARY] inner_competitor = torch.maximum( pooled_role[:, :, ROUTE_OUTER_PRIMARY], pooled_role[:, :, ROUTE_SECONDARY], ) inner_margin_ratio = ( (inner_score - inner_competitor).clamp_min(0.0) / inner_score.clamp_min(1e-8) ).clamp(0.0, 1.0) positive = ( support & (foreground_coverage > float(background_max_coverage)) & (outer_score >= float(positive_confidence)) & (outer_margin_ratio >= float(positive_margin)) ) background_negative = support & ( foreground_coverage <= float(background_max_coverage) ) inner_negative = ( support & ~background_negative & (inner_score >= float(negative_confidence)) & (inner_margin_ratio >= float(negative_margin)) ) negative = background_negative | inner_negative shared = support.sum(dim=1) >= int(min_views) conflict = shared & positive.any(dim=1) & negative.any(dim=1) evidence_weight = support.to(dtype=outer_score.dtype) prior = ( (outer_score * foreground_coverage * evidence_weight).sum(dim=1) / evidence_weight.sum(dim=1).clamp_min(1.0) ).clamp(0.0, 1.0) return { "prior": _project_grouped_outer_texels_to_views( renderer, views, prior ), "shared": _project_grouped_outer_texels_to_views( renderer, views, shared.float() ) > 0.5, "conflict": _project_grouped_outer_texels_to_views( renderer, views, conflict.float() ) > 0.5, "positive": _project_grouped_outer_texels_to_views( renderer, views, positive.any(dim=1).float() ) > 0.5, "negative": _project_grouped_outer_texels_to_views( renderer, views, negative.any(dim=1).float() ) > 0.5, "shared_texels": shared, "conflict_texels": conflict, } def _soft_route_role_consensus( role_prob, static, foreground_prob, consensus_weight=0.60, preserve_outer_confidence=0.80, preserve_outer_margin=0.35, outer_confidence=0.80, outer_margin=0.35, outer_uv_occupancy=None, occupancy_blend_weight=0.0, occupancy_gate_threshold=0.15, occupancy_rescue_threshold=0.70, occupancy_rescue_route_threshold=0.30, cross_view_outer_prior=None, cross_view_outer_valid=None, cross_view_outer_veto=None, cross_view_outer_weight=0.50, ): """Denoise route roles without erasing strong local outer-layer evidence.""" for name, value in ( ("consensus_weight", consensus_weight), ("preserve_outer_confidence", preserve_outer_confidence), ("preserve_outer_margin", preserve_outer_margin), ("outer_confidence", outer_confidence), ("outer_margin", outer_margin), ("occupancy_blend_weight", occupancy_blend_weight), ("occupancy_gate_threshold", occupancy_gate_threshold), ("occupancy_rescue_threshold", occupancy_rescue_threshold), ( "occupancy_rescue_route_threshold", occupancy_rescue_route_threshold, ), ("cross_view_outer_weight", cross_view_outer_weight), ): if not 0.0 <= float(value) <= 1.0: raise ValueError(f"{name} must be in [0, 1].") role_prob = role_prob.float() raw_role = role_prob.argmax(dim=1) raw_top = role_prob.topk(2, dim=1).values raw_margin = raw_top[:, 0] - raw_top[:, 1] raw_selected_score = role_prob.gather(1, raw_role.unsqueeze(1)).squeeze(1) raw_margin_ratio = ( raw_margin / raw_selected_score.clamp_min(1e-8) ).clamp(0.0, 1.0) direct_scores = _aggregate_role_probabilities_by_direct_texel( role_prob, static, foreground_prob=foreground_prob, center_power=2.0, ) batch = role_prob.shape[0] inner_valid = static["masks"][0].unsqueeze(0).expand(batch, -1, -1) outer_valid = static["masks"][1].unsqueeze(0).expand_as(inner_valid) has_direct_candidate = inner_valid | outer_valid texel_inner = direct_scores[:, 0, ROUTE_INNER_PRIMARY] texel_outer = direct_scores[:, 1, ROUTE_OUTER_PRIMARY] unavailable = role_prob.new_full(texel_inner.shape, -1.0) texel_secondary = torch.maximum( torch.where( inner_valid, direct_scores[:, 0, ROUTE_SECONDARY], unavailable, ), torch.where( outer_valid, direct_scores[:, 1, ROUTE_SECONDARY], unavailable, ), ) local_weight = 1.0 - float(consensus_weight) fused_inner = torch.where( inner_valid, local_weight * role_prob[:, ROUTE_INNER_PRIMARY] + float(consensus_weight) * texel_inner, unavailable, ) fused_outer = torch.where( outer_valid, local_weight * role_prob[:, ROUTE_OUTER_PRIMARY] + float(consensus_weight) * texel_outer, unavailable, ) if cross_view_outer_prior is not None: if ( cross_view_outer_valid is None or cross_view_outer_prior.shape != texel_outer.shape or cross_view_outer_valid.shape != texel_outer.shape ): raise ValueError( "Cross-view outer prior and validity must match the route map." ) cross_view_outer_prior = cross_view_outer_prior.float().clamp(0.0, 1.0) cross_view_outer_valid = cross_view_outer_valid.bool() fused_outer = torch.where( outer_valid & cross_view_outer_valid, (1.0 - float(cross_view_outer_weight)) * fused_outer + float(cross_view_outer_weight) * cross_view_outer_prior, fused_outer, ) if cross_view_outer_veto is None: cross_view_outer_veto = torch.zeros_like(outer_valid) elif cross_view_outer_veto.shape != texel_outer.shape: raise ValueError("Cross-view outer veto must match the route map.") else: cross_view_outer_veto = cross_view_outer_veto.bool() occupancy_available = outer_uv_occupancy is not None if occupancy_available: if outer_uv_occupancy.shape != texel_outer.shape: raise ValueError( "Projected outer occupancy must match the route map, got " f"{tuple(outer_uv_occupancy.shape)} and " f"{tuple(texel_outer.shape)}." ) occupancy_prior = outer_uv_occupancy.float().clamp(0.0, 1.0) fused_outer = torch.where( outer_valid, (1.0 - float(occupancy_blend_weight)) * fused_outer + float(occupancy_blend_weight) * occupancy_prior, unavailable, ) occupancy_allows_outer = ( occupancy_prior >= float(occupancy_gate_threshold) ) else: occupancy_prior = torch.ones_like(texel_outer) occupancy_allows_outer = torch.ones_like(outer_valid) # Secondary/backface observations remain strict cell-level decisions. The # soft blend is needed for the ambiguous inner/outer pair; retaining an # isolated secondary pixel would reintroduce exactly the speckle that # projected-texel consensus is meant to remove. fused_secondary = torch.where( has_direct_candidate, texel_secondary.clamp_min(0.0), unavailable, ) raw_outer_score = role_prob[:, ROUTE_OUTER_PRIMARY] raw_outer_competitor = torch.maximum( role_prob[:, ROUTE_INNER_PRIMARY], role_prob[:, ROUTE_SECONDARY], ) raw_outer_margin = (raw_outer_score - raw_outer_competitor).clamp_min(0.0) raw_outer_margin_ratio = ( raw_outer_margin / raw_outer_score.clamp_min(1e-8) ).clamp(0.0, 1.0) raw_preserve_outer = ( outer_valid & ~cross_view_outer_veto & occupancy_allows_outer & (raw_role == ROUTE_OUTER_PRIMARY) & (raw_outer_score >= float(preserve_outer_confidence)) & (raw_outer_margin_ratio >= float(preserve_outer_margin)) ) fused_outer_competitor = torch.maximum(fused_inner, fused_secondary) fused_outer_margin = (fused_outer - fused_outer_competitor).clamp_min(0.0) fused_outer_margin_ratio = ( fused_outer_margin / fused_outer.clamp_min(1e-8) ).clamp(0.0, 1.0) fused_outer_admitted = ( outer_valid & (fused_outer >= float(outer_confidence)) & (fused_outer_margin_ratio >= float(outer_margin)) ) # Raw outer confidence is deliberately not an admission shortcut. A # coherent inner texture (eyes, mouth, shirt graphics) can make the local # route head extremely confident while still being semantically wrong. # Every ordinary outer decision must therefore pass the fused texel gate. # Occupancy evidence remains the only explicit bypass because it predicts # atlas-level outer existence rather than repeating the local route score. outer_admitted = fused_outer_admitted & ( (raw_role != ROUTE_OUTER_PRIMARY) | raw_preserve_outer ) preserve_outer = raw_preserve_outer & fused_outer_admitted outer_gate_rejected = ( outer_valid & ~cross_view_outer_veto & occupancy_allows_outer & (raw_role == ROUTE_OUTER_PRIMARY) & ~outer_admitted ) occupancy_rescue = ( outer_valid & occupancy_allows_outer & (occupancy_prior >= float(occupancy_rescue_threshold)) & ( raw_outer_score >= float(occupancy_rescue_route_threshold) ) ) if occupancy_available else torch.zeros_like(outer_valid) allow_outer = ( outer_valid & ~cross_view_outer_veto & occupancy_allows_outer & (outer_admitted | occupancy_rescue) ) outer_candidate_score = torch.where( occupancy_rescue, torch.maximum( fused_outer, occupancy_prior, ), fused_outer, ) eligible_outer = torch.where(allow_outer, outer_candidate_score, unavailable) eligible_secondary = torch.where( (fused_secondary >= SECONDARY_TEXEL_MIN_PROBABILITY) & has_direct_candidate, fused_secondary, unavailable, ) eligible_scores = torch.stack( [fused_inner, eligible_outer, eligible_secondary], dim=1, ) selected_score, selected_role = eligible_scores.max(dim=1) selected_role = torch.where(has_direct_candidate, selected_role, raw_role) selected_score = torch.where( has_direct_candidate, selected_score, raw_selected_score, ) top_scores = eligible_scores.topk(2, dim=1).values selected_margin = (top_scores[:, 0] - top_scores[:, 1]).clamp_min(0.0) selected_margin_ratio = ( selected_margin / selected_score.clamp_min(1e-8) ).clamp(0.0, 1.0) # If fused outer still wins but does not clear the asymmetric admission # gate, keep its outer UV identity while marking it untrusted. Downstream # filtering can then leave the observation unknown (and preserve it as # rejected context) instead of writing an uncertain outer color into the # inner atlas. When fused inner actually wins, the ordinary argmax above # performs the intended outer-to-inner correction. deferred_outer = outer_gate_rejected & ( fused_outer >= fused_outer_competitor ) selected_role = torch.where( deferred_outer, torch.full_like(selected_role, ROUTE_OUTER_PRIMARY), selected_role, ) selected_score = torch.where(deferred_outer, fused_outer, selected_score) selected_margin = torch.where( deferred_outer, fused_outer_margin, selected_margin, ) selected_margin_ratio = torch.where( deferred_outer, fused_outer_margin_ratio, selected_margin_ratio, ) return { "role": selected_role, "score": selected_score, "margin": selected_margin, "margin_ratio": selected_margin_ratio, "scores": torch.stack( [fused_inner, fused_outer, fused_secondary], dim=1, ), "has_direct_candidate": has_direct_candidate, "preserved_outer": preserve_outer, "outer_gate_rejected": outer_gate_rejected, "occupancy_rescued_outer": ( occupancy_rescue & (raw_role != ROUTE_OUTER_PRIMARY) & (selected_role == ROUTE_OUTER_PRIMARY) ), "occupancy_rejected_outer": ( outer_valid & (raw_role == ROUTE_OUTER_PRIMARY) & ~occupancy_allows_outer ), "cross_view_outer_veto": cross_view_outer_veto, "outer_uv_occupancy": ( occupancy_prior if occupancy_available else None ), "changed": has_direct_candidate & (selected_role != raw_role), "inner_to_outer": ( has_direct_candidate & (raw_role == ROUTE_INNER_PRIMARY) & (selected_role == ROUTE_OUTER_PRIMARY) ), "outer_to_inner": ( has_direct_candidate & (raw_role == ROUTE_OUTER_PRIMARY) & (selected_role == ROUTE_INNER_PRIMARY) ), } def _routing_from_geometry_outputs( renderer, views, outputs, fg_threshold=0.5, outer_threshold=0.5, texel_consensus=True, texel_consensus_weight=0.60, preserve_outer_confidence=0.80, preserve_outer_margin=0.35, consensus_outer_confidence=0.70, consensus_outer_margin=0.20, outer_uv_occupancy=False, outer_uv_occupancy_blend_weight=0.0, outer_uv_occupancy_gate_threshold=0.15, outer_uv_occupancy_rescue_threshold=0.70, outer_uv_occupancy_rescue_route_threshold=0.30, outer_uv_component_routing=False, outer_uv_component_seed_threshold=0.80, outer_uv_component_grow_threshold=0.50, outer_uv_component_min_size=2, observed_foreground=None, cross_view_outer_consistency=True, cross_view_outer_weight=0.50, cross_view_outer_positive_confidence=0.70, cross_view_outer_positive_margin=0.20, cross_view_outer_negative_confidence=0.70, cross_view_outer_negative_margin=0.20, cross_view_outer_background_max_coverage=0.25, cross_view_outer_min_views=2, ): """Route a fitted Steve render through fixed inner/outer cuboid UV maps.""" views = parse_views(views) if not views: raise ValueError("At least one renderer view is required for geometry routing.") foreground_prob = torch.sigmoid(outputs["foreground"])[:, 0] if outputs["layer"].shape[1] != ROUTE_ROLE_CLASSES: raise ValueError( "geometry_fit requires a three-class route-role head " "(inner_primary, outer_primary, secondary_backface). Retrain the parser." ) role_prob = torch.softmax(outputs["layer"].float(), dim=1) raw_route_role = role_prob.argmax(dim=1) route_role = raw_route_role.clone() outer_prob = role_prob[:, ROUTE_OUTER_PRIMARY].clone() inner_prob = role_prob[:, ROUTE_INNER_PRIMARY].clone() top_prob = role_prob.topk(2, dim=1).values role_margin = top_prob[:, 0] - top_prob[:, 1] N, H, W = outer_prob.shape if N % len(views) != 0: raise ValueError(f"N={N} must be divisible by {len(views)} renderer views.") fg = torch.zeros_like(outer_prob, dtype=torch.bool) layer = torch.zeros((N, H, W), dtype=torch.long, device=outer_prob.device) flat_uv = torch.zeros_like(layer) surface = torch.full_like(layer, IGNORE_INDEX) confidence = torch.zeros_like(outer_prob) confidence_margin = torch.zeros_like(outer_prob) confidence_margin_ratio = torch.zeros_like(outer_prob) texel_center_score = torch.zeros_like(outer_prob) part = torch.full_like(layer, IGNORE_INDEX) face = torch.full_like(layer, IGNORE_INDEX) consensus_changed = torch.zeros_like(fg) consensus_inner_to_outer = torch.zeros_like(fg) consensus_outer_to_inner = torch.zeros_like(fg) consensus_preserved_outer = torch.zeros_like(fg) consensus_outer_gate_rejected = torch.zeros_like(fg) occupancy_rescued_outer = torch.zeros_like(fg) occupancy_rejected_outer = torch.zeros_like(fg) projected_outer_occupancy = torch.zeros_like(outer_prob) projected_outer_component_seed = torch.zeros_like(fg) projected_outer_component_grown = torch.zeros_like(fg) projected_outer_component_rejected = torch.zeros_like(fg) occupancy_logits = ( outputs.get("outer_uv_occupancy_logits") if outer_uv_occupancy else None ) occupancy_details = None occupancy_probability = None if occupancy_logits is not None: raw_occupancy_probability = torch.sigmoid(occupancy_logits.float()) if outer_uv_component_routing: occupancy_details = outer_uv_topology_hysteresis( raw_occupancy_probability, seed_threshold=outer_uv_component_seed_threshold, grow_threshold=outer_uv_component_grow_threshold, min_component_size=outer_uv_component_min_size, ) occupancy_probability = occupancy_details["probability"] else: occupancy_probability = raw_occupancy_probability cross_view = None if cross_view_outer_consistency and len(views) >= cross_view_outer_min_views: if observed_foreground is None: observed_foreground = foreground_prob > float(fg_threshold) cross_view = cross_view_outer_visibility_evidence( renderer, views, role_prob, observed_foreground, positive_confidence=cross_view_outer_positive_confidence, positive_margin=cross_view_outer_positive_margin, negative_confidence=cross_view_outer_negative_confidence, negative_margin=cross_view_outer_negative_margin, background_max_coverage=( cross_view_outer_background_max_coverage ), min_views=cross_view_outer_min_views, ) for view_index, view in enumerate(views): selection = slice(view_index, N, len(views)) static = build_static_surface_routing(renderer, view, outer_prob.device) if static["masks"].shape[-2:] != (H, W): raise ValueError( f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} " f"does not match parser shape {(H, W)}." ) inner_valid = static["masks"][0].unsqueeze(0).expand(outer_prob[selection].shape[0], -1, -1) outer_valid = static["masks"][1].unsqueeze(0).expand_as(inner_valid) if ( texel_consensus or cross_view is not None or occupancy_probability is not None ): view_occupancy = _project_outer_uv_occupancy_to_view( occupancy_probability, static, outer_prob[selection].shape[0], ) if occupancy_details is not None: projected_outer_component_seed[selection] = ( _project_outer_uv_occupancy_to_view( occupancy_details["seed"].float(), static, outer_prob[selection].shape[0], ) > 0.5 ) projected_outer_component_grown[selection] = ( _project_outer_uv_occupancy_to_view( occupancy_details["grown"].float(), static, outer_prob[selection].shape[0], ) > 0.5 ) projected_outer_component_rejected[selection] = ( _project_outer_uv_occupancy_to_view( occupancy_details["rejected_candidate"].float(), static, outer_prob[selection].shape[0], ) > 0.5 ) consensus = _soft_route_role_consensus( role_prob[selection], static, foreground_prob[selection], consensus_weight=( texel_consensus_weight if texel_consensus else 0.0 ), preserve_outer_confidence=preserve_outer_confidence, preserve_outer_margin=preserve_outer_margin, outer_confidence=consensus_outer_confidence, outer_margin=consensus_outer_margin, outer_uv_occupancy=view_occupancy, occupancy_blend_weight=( outer_uv_occupancy_blend_weight ), occupancy_gate_threshold=( outer_uv_occupancy_gate_threshold ), occupancy_rescue_threshold=( outer_uv_occupancy_rescue_threshold ), occupancy_rescue_route_threshold=( outer_uv_occupancy_rescue_route_threshold ), cross_view_outer_prior=( cross_view["prior"][selection] if cross_view is not None else None ), cross_view_outer_valid=( cross_view["shared"][selection] if cross_view is not None else None ), cross_view_outer_veto=( cross_view["conflict"][selection] if cross_view is not None else None ), cross_view_outer_weight=cross_view_outer_weight, ) has_direct_candidate = consensus["has_direct_candidate"] route_role[selection] = consensus["role"] inner_prob[selection] = torch.where( has_direct_candidate, torch.where( consensus["role"] == ROUTE_INNER_PRIMARY, consensus["score"], consensus["scores"][:, ROUTE_INNER_PRIMARY], ).clamp_min(0.0), inner_prob[selection], ) outer_prob[selection] = torch.where( has_direct_candidate, torch.where( consensus["role"] == ROUTE_OUTER_PRIMARY, consensus["score"], consensus["scores"][:, ROUTE_OUTER_PRIMARY], ).clamp_min(0.0), outer_prob[selection], ) role_margin[selection] = consensus["margin"] consensus_changed[selection] = consensus["changed"] consensus_inner_to_outer[selection] = consensus["inner_to_outer"] consensus_outer_to_inner[selection] = consensus["outer_to_inner"] consensus_preserved_outer[selection] = consensus["preserved_outer"] consensus_outer_gate_rejected[selection] = consensus[ "outer_gate_rejected" ] occupancy_rescued_outer[selection] = consensus[ "occupancy_rescued_outer" ] occupancy_rejected_outer[selection] = consensus[ "occupancy_rejected_outer" ] if view_occupancy is not None: projected_outer_occupancy[selection] = view_occupancy visible_outer = ( (route_role[selection] == ROUTE_OUTER_PRIMARY) & (outer_prob[selection] >= outer_threshold) ) # The Steve base layer is opaque and therefore supplies a complete # silhouette. Outer-only protrusions still need image foreground evidence. choose_outer = outer_valid & visible_outer & ( inner_valid | (foreground_prob[selection] > fg_threshold) ) choose_inner = inner_valid & (route_role[selection] == ROUTE_INNER_PRIMARY) selected_fg = choose_inner | choose_outer inner_flat_uv = static["flat_uv"][0].unsqueeze(0).expand_as(layer[selection]) outer_flat_uv = static["flat_uv"][1].unsqueeze(0).expand_as(layer[selection]) inner_part = static["part"][0].unsqueeze(0).expand_as(layer[selection]) outer_part = static["part"][1].unsqueeze(0).expand_as(layer[selection]) inner_face = static["face"][0].unsqueeze(0).expand_as(layer[selection]) outer_face = static["face"][1].unsqueeze(0).expand_as(layer[selection]) inner_center_score = static["texel_center_score"][0].unsqueeze(0).expand_as( confidence[selection] ) outer_center_score = static["texel_center_score"][1].unsqueeze(0).expand_as( confidence[selection] ) fg[selection] = selected_fg layer[selection] = choose_outer.long() flat_uv[selection] = torch.where(choose_outer, outer_flat_uv, inner_flat_uv) surface[selection] = torch.where( selected_fg, choose_outer.long(), torch.full_like(layer[selection], IGNORE_INDEX), ) part[selection] = torch.where( selected_fg, torch.where(choose_outer, outer_part, inner_part), torch.full_like(layer[selection], IGNORE_INDEX), ) face[selection] = torch.where( selected_fg, torch.where(choose_outer, outer_face, inner_face), torch.full_like(layer[selection], IGNORE_INDEX), ) selected_confidence = torch.where(choose_outer, outer_prob[selection], inner_prob[selection]) margin = role_margin[selection] confidence[selection] = selected_confidence confidence_margin[selection] = margin confidence_margin_ratio[selection] = ( margin / selected_confidence.clamp_min(1e-8) ).clamp(0.0, 1.0) texel_center_score[selection] = torch.where( choose_outer, outer_center_score, inner_center_score ) secondary = (route_role == ROUTE_SECONDARY) & (foreground_prob > fg_threshold) return { "foreground": fg, "layer": layer, "flat_uv": flat_uv, "surface": surface, "confidence": confidence, "confidence_margin": confidence_margin, "confidence_margin_ratio": confidence_margin_ratio, "semantic_fallback": torch.zeros_like(fg), "part": part, "face": face, "raw_route_role": raw_route_role, "route_role": route_role, "secondary": secondary, "texel_center_score": texel_center_score, "consensus_changed": consensus_changed, "consensus_inner_to_outer": consensus_inner_to_outer, "consensus_outer_to_inner": consensus_outer_to_inner, "consensus_preserved_outer": consensus_preserved_outer, "consensus_outer_gate_rejected": consensus_outer_gate_rejected, "occupancy_rescued_outer": occupancy_rescued_outer, "occupancy_rejected_outer": occupancy_rejected_outer, "projected_outer_occupancy": projected_outer_occupancy, "projected_outer_component_seed": projected_outer_component_seed, "projected_outer_component_grown": projected_outer_component_grown, "projected_outer_component_rejected": ( projected_outer_component_rejected ), "outer_uv_occupancy_available": torch.full_like( fg, occupancy_logits is not None ), "cross_view_outer_shared": ( cross_view["shared"] if cross_view is not None else torch.zeros_like(fg) ), "cross_view_outer_conflict": ( cross_view["conflict"] if cross_view is not None else torch.zeros_like(fg) ), "cross_view_outer_vetoed": ( cross_view["conflict"] & (raw_route_role == ROUTE_OUTER_PRIMARY) & (route_role != ROUTE_OUTER_PRIMARY) if cross_view is not None else torch.zeros_like(fg) ), "cross_view_outer_positive": ( cross_view["positive"] if cross_view is not None else torch.zeros_like(fg) ), "cross_view_outer_negative": ( cross_view["negative"] if cross_view is not None else torch.zeros_like(fg) ), } def _aggregate_surface_scores_by_candidate_texel(candidate_score, static): """Average each surface candidate inside its projected Minecraft UV texel.""" candidate_score = candidate_score.float() batch, surface_count, height, width = candidate_score.shape flat_uv = static["flat_uv"][:surface_count].reshape(1, surface_count, -1) valid = static["masks"][:surface_count].reshape(1, surface_count, -1) indices = flat_uv.expand(batch, -1, -1) values = candidate_score.flatten(2) * valid.to(dtype=candidate_score.dtype) sums = candidate_score.new_zeros(batch, surface_count, UV_SIZE * UV_SIZE) sums.scatter_add_(2, indices, values) counts = candidate_score.new_zeros(1, surface_count, UV_SIZE * UV_SIZE) counts.scatter_add_( 2, flat_uv, valid.to(dtype=candidate_score.dtype), ) means = sums / counts.clamp_min(1.0) return means.gather(2, indices).reshape(batch, surface_count, height, width) def _static_candidate_roles(static): """Return the Minecraft layer and route role implemented by every surface pixel.""" candidate_layer = static["layer"].clamp(0, LAYER_CLASSES - 1) direct_masks = static["masks"][:LAYER_CLASSES].gather(0, candidate_layer) direct_uv = static["flat_uv"][:LAYER_CLASSES].gather(0, candidate_layer) primary = static["masks"] & direct_masks & (static["flat_uv"] == direct_uv) candidate_role = torch.where( primary, candidate_layer, torch.full_like(candidate_layer, ROUTE_SECONDARY), ) return candidate_layer, candidate_role def _routing_from_geometry_surface_outputs( renderer, views, outputs, fg_threshold=0.5, texel_consensus=True, texel_consensus_weight=0.60, preserve_outer_confidence=0.80, preserve_outer_margin=0.35, consensus_outer_confidence=0.70, consensus_outer_margin=0.20, outer_uv_occupancy=False, outer_uv_occupancy_blend_weight=0.0, outer_uv_occupancy_gate_threshold=0.15, outer_uv_occupancy_rescue_threshold=0.70, outer_uv_occupancy_rescue_route_threshold=0.30, outer_uv_component_routing=False, outer_uv_component_seed_threshold=0.80, outer_uv_component_grow_threshold=0.50, outer_uv_component_min_size=2, observed_foreground=None, cross_view_outer_consistency=True, cross_view_outer_weight=0.50, cross_view_outer_positive_confidence=0.70, cross_view_outer_positive_margin=0.20, cross_view_outer_negative_confidence=0.70, cross_view_outer_negative_margin=0.20, cross_view_outer_background_max_coverage=0.25, cross_view_outer_min_views=2, ): """Route primary layers first, then use the surface head for secondary UVs. The route-role head owns the inner/outer/secondary decision. The surface head only disambiguates renderer slots that implement the selected role. Keeping those decisions hierarchical is important: many composite and geometry slots overlap a direct cuboid face, so a strong wrong surface logit must not turn a primary pixel into a secondary/back-facing pixel. """ views = parse_views(views) if not views: raise ValueError("At least one renderer view is required for geometry-surface routing.") if outputs["layer"].shape[1] != ROUTE_ROLE_CLASSES: raise ValueError("Geometry-surface routing requires the three-class route-role head.") if "surface" not in outputs: raise ValueError("Geometry-surface routing requires a static surface classifier.") foreground_prob = torch.sigmoid(outputs["foreground"])[:, 0] role_prob = torch.softmax(outputs["layer"].float(), dim=1) surface_prob = torch.softmax(outputs["surface"].float(), dim=1) raw_route_role = role_prob.argmax(dim=1) N, _, H, W = surface_prob.shape if N % len(views) != 0: raise ValueError(f"N={N} must be divisible by {len(views)} renderer views.") fg = torch.zeros((N, H, W), dtype=torch.bool, device=surface_prob.device) layer = torch.zeros((N, H, W), dtype=torch.long, device=surface_prob.device) flat_uv = torch.zeros_like(layer) surface = torch.full_like(layer, IGNORE_INDEX) route_role = torch.full_like(layer, IGNORE_INDEX) confidence = torch.zeros_like(foreground_prob) confidence_margin = torch.zeros_like(foreground_prob) confidence_margin_ratio = torch.zeros_like(foreground_prob) texel_center_score = torch.zeros_like(foreground_prob) part = torch.full_like(layer, IGNORE_INDEX) face = torch.full_like(layer, IGNORE_INDEX) consensus_changed = torch.zeros_like(fg) consensus_inner_to_outer = torch.zeros_like(fg) consensus_outer_to_inner = torch.zeros_like(fg) consensus_preserved_outer = torch.zeros_like(fg) consensus_outer_gate_rejected = torch.zeros_like(fg) occupancy_rescued_outer = torch.zeros_like(fg) occupancy_rejected_outer = torch.zeros_like(fg) projected_outer_occupancy = torch.zeros_like(foreground_prob) projected_outer_component_seed = torch.zeros_like(fg) projected_outer_component_grown = torch.zeros_like(fg) projected_outer_component_rejected = torch.zeros_like(fg) occupancy_logits = ( outputs.get("outer_uv_occupancy_logits") if outer_uv_occupancy else None ) occupancy_details = None occupancy_probability = None if occupancy_logits is not None: raw_occupancy_probability = torch.sigmoid(occupancy_logits.float()) if outer_uv_component_routing: occupancy_details = outer_uv_topology_hysteresis( raw_occupancy_probability, seed_threshold=outer_uv_component_seed_threshold, grow_threshold=outer_uv_component_grow_threshold, min_component_size=outer_uv_component_min_size, ) occupancy_probability = occupancy_details["probability"] else: occupancy_probability = raw_occupancy_probability cross_view = None if cross_view_outer_consistency and len(views) >= cross_view_outer_min_views: if observed_foreground is None: observed_foreground = foreground_prob > float(fg_threshold) cross_view = cross_view_outer_visibility_evidence( renderer, views, role_prob, observed_foreground, positive_confidence=cross_view_outer_positive_confidence, positive_margin=cross_view_outer_positive_margin, negative_confidence=cross_view_outer_negative_confidence, negative_margin=cross_view_outer_negative_margin, background_max_coverage=( cross_view_outer_background_max_coverage ), min_views=cross_view_outer_min_views, ) for view_index, view in enumerate(views): selection = slice(view_index, N, len(views)) static = build_static_surface_routing(renderer, view, surface_prob.device) if static["masks"].shape[-2:] != (H, W): raise ValueError( f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} " f"does not match parser shape {(H, W)}." ) surface_count = static["masks"].shape[0] if surface_prob.shape[1] < surface_count: raise ValueError( f"Parser has {surface_prob.shape[1]} surface classes, but view {view!r} " f"requires {surface_count}." ) candidate_layer, candidate_role = _static_candidate_roles(static) view_batch = surface_prob[selection].shape[0] view_role_prob = role_prob[selection] selected_role = raw_route_role[selection] selected_role_score = view_role_prob.gather( 1, selected_role.unsqueeze(1) ).squeeze(1) role_top = view_role_prob.topk(2, dim=1).values role_margin = role_top[:, 0] - role_top[:, 1] if ( texel_consensus or cross_view is not None or occupancy_probability is not None ): view_occupancy = _project_outer_uv_occupancy_to_view( occupancy_probability, static, view_batch, ) if occupancy_details is not None: projected_outer_component_seed[selection] = ( _project_outer_uv_occupancy_to_view( occupancy_details["seed"].float(), static, view_batch, ) > 0.5 ) projected_outer_component_grown[selection] = ( _project_outer_uv_occupancy_to_view( occupancy_details["grown"].float(), static, view_batch, ) > 0.5 ) projected_outer_component_rejected[selection] = ( _project_outer_uv_occupancy_to_view( occupancy_details["rejected_candidate"].float(), static, view_batch, ) > 0.5 ) consensus = _soft_route_role_consensus( view_role_prob, static, foreground_prob[selection], consensus_weight=( texel_consensus_weight if texel_consensus else 0.0 ), preserve_outer_confidence=preserve_outer_confidence, preserve_outer_margin=preserve_outer_margin, outer_confidence=consensus_outer_confidence, outer_margin=consensus_outer_margin, outer_uv_occupancy=view_occupancy, occupancy_blend_weight=( outer_uv_occupancy_blend_weight ), occupancy_gate_threshold=( outer_uv_occupancy_gate_threshold ), occupancy_rescue_threshold=( outer_uv_occupancy_rescue_threshold ), occupancy_rescue_route_threshold=( outer_uv_occupancy_rescue_route_threshold ), cross_view_outer_prior=( cross_view["prior"][selection] if cross_view is not None else None ), cross_view_outer_valid=( cross_view["shared"][selection] if cross_view is not None else None ), cross_view_outer_veto=( cross_view["conflict"][selection] if cross_view is not None else None ), cross_view_outer_weight=cross_view_outer_weight, ) selected_role = consensus["role"] selected_role_score = consensus["score"] role_margin = consensus["margin"] consensus_changed[selection] = consensus["changed"] consensus_inner_to_outer[selection] = consensus["inner_to_outer"] consensus_outer_to_inner[selection] = consensus["outer_to_inner"] consensus_preserved_outer[selection] = consensus["preserved_outer"] consensus_outer_gate_rejected[selection] = consensus[ "outer_gate_rejected" ] occupancy_rescued_outer[selection] = consensus[ "occupancy_rescued_outer" ] occupancy_rejected_outer[selection] = consensus[ "occupancy_rejected_outer" ] if view_occupancy is not None: projected_outer_occupancy[selection] = view_occupancy candidate_score = surface_prob[selection, :surface_count] candidate_valid = static["masks"].unsqueeze(0).expand( view_batch, -1, -1, -1 ) candidate_valid = candidate_valid & ( candidate_role.unsqueeze(0) == selected_role.unsqueeze(1) ) if texel_consensus: consensus_score = _aggregate_surface_scores_by_candidate_texel( candidate_score, static, ) local_scores = candidate_score.masked_fill(~candidate_valid, -1.0) local_top = local_scores.topk(k=min(2, surface_count), dim=1).values if surface_count > 1: local_margin_ratio = ( (local_top[:, 0] - local_top[:, 1]).clamp_min(0.0) / local_top[:, 0].abs().clamp_min(1e-8) ).clamp(0.0, 1.0) consensus_weight = (1.0 - local_margin_ratio).unsqueeze(1) candidate_score = ( candidate_score * (1.0 - consensus_weight) + consensus_score * consensus_weight ) candidate_score = candidate_score.masked_fill(~candidate_valid, -1.0) best_surface_score, selected_surface = candidate_score.max(dim=1) selected_secondary = selected_role == ROUTE_SECONDARY # Primary routes already have an exact direct cuboid surface and UV. # Do not let an equivalent composite slot bypass direct-layer coverage # filtering; the learned surface head is needed only for secondary UVs. selected_surface = torch.where( selected_secondary, selected_surface, selected_role, ) if surface_count > 1: top_scores = candidate_score.topk(k=2, dim=1).values surface_margin = (top_scores[:, 0] - top_scores[:, 1]).clamp_min(0.0) surface_margin_ratio = ( surface_margin / top_scores[:, 0].abs().clamp_min(1e-8) ).clamp(0.0, 1.0) else: surface_margin = best_surface_score.clamp_min(0.0) surface_margin_ratio = (best_surface_score > 0.0).to( best_surface_score.dtype ) routed = _select_static_surface(static, selected_surface) has_candidate = best_surface_score >= 0.0 base_silhouette = static["masks"][0].unsqueeze(0).expand(view_batch, -1, -1) routed_fg = ( ((foreground_prob[selection] > fg_threshold) | base_silhouette) & has_candidate & routed["valid"] ) role_margin_ratio = ( role_margin / selected_role_score.clamp_min(1e-8) ).clamp(0.0, 1.0) selected_surface_score = best_surface_score.clamp_min(0.0) selected_confidence = torch.where( selected_secondary, (selected_role_score * selected_surface_score).sqrt(), selected_role_score, ) if "route_confidence" in outputs: # The learned trust target already says whether the argmax route is # correct (and includes surface correctness for secondary routes). # Multiplying it by the class probability counts the same # uncertainty twice and makes a 0.80 outer gate nearly # unreachable. A geometric mean keeps both signals necessary # without systematically collapsing otherwise well-calibrated # confidence. learned_trust = torch.sigmoid( outputs["route_confidence"][selection, 0].float() ) primary_confidence = ( selected_role_score * learned_trust ).clamp_min(0.0).sqrt() secondary_confidence = ( selected_role_score * selected_surface_score * learned_trust ).clamp_min(0.0).pow(1.0 / 3.0) selected_confidence = torch.where( selected_secondary, secondary_confidence, primary_confidence, ) selected_margin = torch.where( selected_secondary, torch.minimum(role_margin, surface_margin), role_margin, ) selected_margin_ratio = torch.where( selected_secondary, torch.minimum(role_margin_ratio, surface_margin_ratio), role_margin_ratio, ) fg[selection] = routed_fg layer[selection] = routed["layer"] flat_uv[selection] = routed["flat_uv"] surface[selection] = torch.where( routed_fg, selected_surface, torch.full_like(selected_surface, IGNORE_INDEX), ) route_role[selection] = torch.where( routed_fg, selected_role, torch.full_like(selected_role, IGNORE_INDEX), ) confidence[selection] = selected_confidence confidence_margin[selection] = selected_margin confidence_margin_ratio[selection] = selected_margin_ratio part[selection] = routed["part"] face[selection] = routed["face"] texel_center_score[selection] = routed["texel_center_score"] secondary = fg & (route_role == ROUTE_SECONDARY) return { "foreground": fg, "layer": layer, "flat_uv": flat_uv, "surface": surface, "confidence": confidence, "confidence_margin": confidence_margin, "confidence_margin_ratio": confidence_margin_ratio, "semantic_fallback": torch.zeros_like(fg), "part": part, "face": face, "raw_route_role": raw_route_role, "route_role": route_role, "secondary": secondary, "secondary_routed": secondary, "texel_center_score": texel_center_score, "consensus_changed": consensus_changed, "consensus_inner_to_outer": consensus_inner_to_outer, "consensus_outer_to_inner": consensus_outer_to_inner, "consensus_preserved_outer": consensus_preserved_outer, "consensus_outer_gate_rejected": consensus_outer_gate_rejected, "occupancy_rescued_outer": occupancy_rescued_outer, "occupancy_rejected_outer": occupancy_rejected_outer, "projected_outer_occupancy": projected_outer_occupancy, "projected_outer_component_seed": projected_outer_component_seed, "projected_outer_component_grown": projected_outer_component_grown, "projected_outer_component_rejected": ( projected_outer_component_rejected ), "outer_uv_occupancy_available": torch.full_like( fg, occupancy_logits is not None ), "cross_view_outer_shared": ( cross_view["shared"] if cross_view is not None else torch.zeros_like(fg) ), "cross_view_outer_conflict": ( cross_view["conflict"] if cross_view is not None else torch.zeros_like(fg) ), "cross_view_outer_vetoed": ( cross_view["conflict"] & (raw_route_role == ROUTE_OUTER_PRIMARY) & (route_role != ROUTE_OUTER_PRIMARY) if cross_view is not None else torch.zeros_like(fg) ), "cross_view_outer_positive": ( cross_view["positive"] if cross_view is not None else torch.zeros_like(fg) ), "cross_view_outer_negative": ( cross_view["negative"] if cross_view is not None else torch.zeros_like(fg) ), "learned_trust": ( torch.sigmoid(outputs["route_confidence"][:, 0].float()) if "route_confidence" in outputs else torch.ones_like(foreground_prob) ), } def _surface_aware_outer_coverage(routing, trusted, selected_outer, renderer, views): """Measure outer support within the selected view/surface/UV texel.""" views = parse_views(views) views_per_group = len(views) group_count = trusted.shape[0] // views_per_group static_views = [ build_static_surface_routing(renderer, view, trusted.device) for view in views ] surface_count = max(static["masks"].shape[0] for static in static_views) uv_count = UV_SIZE * UV_SIZE view_stride = surface_count * uv_count group_stride = views_per_group * view_stride expected = routing["confidence"].new_zeros(group_stride) for view_index, static in enumerate(static_views): static_surface_count = static["masks"].shape[0] surface_offsets = ( torch.arange(static_surface_count, device=trusted.device).view(-1, 1, 1) * uv_count ) expected_indices = ( view_index * view_stride + surface_offsets + static["flat_uv"][:static_surface_count] ) expected_mask = ( static["masks"][:static_surface_count] & (static["layer"][:static_surface_count] == 1) ) selected_indices = expected_indices[expected_mask] expected.scatter_add_( 0, selected_indices, torch.ones( selected_indices.shape[0], device=trusted.device, dtype=expected.dtype, ), ) item_indices = torch.arange(trusted.shape[0], device=trusted.device) item_groups = item_indices // views_per_group item_views = item_indices % views_per_group selected_surface = routing["surface"].clamp(0, surface_count - 1) route_index = ( item_views.view(-1, 1, 1) * view_stride + selected_surface * uv_count + routing["flat_uv"] ) grouped_route_index = route_index + item_groups.view(-1, 1, 1) * group_stride observed_mask = trusted & selected_outer observed = routing["confidence"].new_zeros(group_count * group_stride) observed_indices = grouped_route_index[observed_mask] observed.scatter_add_( 0, observed_indices, torch.ones( observed_indices.shape[0], device=trusted.device, dtype=observed.dtype, ), ) coverage = observed.reshape(group_count, group_stride) / expected.clamp_min(1.0) pixel_coverage = coverage[item_groups.view(-1, 1, 1), route_index] # Composite and geometry slots can be partially occluded by nearer surfaces, # so their static masks are not valid coverage denominators. Their exact slot # classifier and texel consensus already reject isolated fragments. return torch.where( selected_surface >= LAYER_CLASSES, torch.ones_like(pixel_coverage), pixel_coverage, ) def _outer_silhouette_coverage( observed_foreground, flat_uv, renderer, views, dilation=0, min_pixels=4, ): """Measure whether an outer texel's protruding pixels exist in the cutout. Only pixels covered by the outer cuboid but not by the inner cuboid carry useful silhouette evidence. Interior projections are deliberately left unassessed because an absent outer texel and the inner surface behind it have the same foreground mask there. Coverage is calculated independently for every view and the least supported assessed view wins. Averaging views would let a texel that is visible in one view but absent in another pass at exactly 50% coverage, which creates the characteristic head-top and clothing-edge spikes. """ views = parse_views(views) if not views: raise ValueError("At least one renderer view is required.") if observed_foreground.shape != flat_uv.shape: raise ValueError( "observed_foreground and flat_uv must have identical shapes." ) if observed_foreground.shape[0] % len(views) != 0: raise ValueError("The routed batch must contain complete view groups.") if dilation < 0: raise ValueError("dilation must be non-negative.") if min_pixels < 1: raise ValueError("min_pixels must be positive.") observed = observed_foreground.to(dtype=torch.bool) if dilation > 0: kernel_size = 2 * int(dilation) + 1 observed = F.max_pool2d( observed.unsqueeze(1).float(), kernel_size=kernel_size, stride=1, padding=int(dilation), )[:, 0] > 0.5 views_per_group = len(views) group_count = observed.shape[0] // views_per_group uv_count = UV_SIZE * UV_SIZE expected = observed.new_zeros( (views_per_group, uv_count), dtype=torch.float32 ) supported = observed.new_zeros( (group_count, views_per_group, uv_count), dtype=torch.float32 ) for view_index, view in enumerate(views): static = build_static_surface_routing( renderer, view, observed.device ) inner_silhouette = static["masks"][ROUTE_INNER_PRIMARY] outer_silhouette = static["masks"][ROUTE_OUTER_PRIMARY] protruding = outer_silhouette & ~inner_silhouette if not protruding.any(): continue protruding_uv = static["flat_uv"][ROUTE_OUTER_PRIMARY][ protruding ] expected[view_index].scatter_add_( 0, protruding_uv, torch.ones_like(protruding_uv, dtype=expected.dtype), ) observed_view = observed[view_index::views_per_group, protruding] supported[:, view_index].scatter_add_( 1, protruding_uv.view(1, -1).expand(group_count, -1), observed_view.to(dtype=supported.dtype), ) assessed_by_view = expected >= float(min_pixels) coverage_by_view = supported / expected.clamp_min(1.0).unsqueeze(0) coverage_by_view = torch.where( assessed_by_view.unsqueeze(0), coverage_by_view, torch.ones_like(coverage_by_view), ) coverage_uv = coverage_by_view.amin(dim=1) assessed_uv = assessed_by_view.any(dim=0) item_groups = ( torch.arange(observed.shape[0], device=observed.device) // views_per_group ) pixel_coverage = coverage_uv[ item_groups.view(-1, 1, 1), flat_uv ] pixel_assessed = assessed_uv[flat_uv] return pixel_coverage, pixel_assessed def _grouped_uv_support(mask, flat_uv, group_size): """Broadcast per-pixel evidence to matching UV texels across all input views.""" if mask.shape != flat_uv.shape: raise ValueError("mask and flat_uv must have identical shapes.") if group_size <= 0 or mask.shape[0] % group_size != 0: raise ValueError("group_size must evenly divide the routed batch.") group_count = mask.shape[0] // group_size item_groups = torch.arange(mask.shape[0], device=mask.device) // group_size grouped_uv = flat_uv + item_groups.view(-1, 1, 1) * (UV_SIZE * UV_SIZE) support_count = torch.zeros( group_count * UV_SIZE * UV_SIZE, device=mask.device, dtype=torch.float32, ) selected_uv = grouped_uv[mask] support_count.scatter_add_( 0, selected_uv, torch.ones(selected_uv.shape[0], device=mask.device, dtype=support_count.dtype), ) return support_count[grouped_uv] > 0 def _geometry_supported_outer_texels(routing, renderer, views): """Find outer texels proven by outer-only silhouette or an exact secondary slot. A direct outer pixel outside the base cuboid silhouette cannot be explained by the inner layer. Likewise a routed secondary surface is tied to an exact renderer slot. Either observation is safe evidence for relaxing confidence gates on the same outer UV texel in another pixel or view. """ views = parse_views(views) if not views: raise ValueError("At least one renderer view is required for outer rescue.") selected_outer = routing["layer"] == 1 if selected_outer.shape[0] % len(views) != 0: raise ValueError("Routed batch must be divisible by the number of views.") outer_only = torch.zeros_like(selected_outer) selected_surface = routing.get("surface") for view_index, view in enumerate(views): selection = slice(view_index, selected_outer.shape[0], len(views)) static = build_static_surface_routing(renderer, view, selected_outer.device) direct_outer = static["masks"][1] & ~static["masks"][0] if selected_surface is not None: direct_outer = direct_outer.unsqueeze(0) & ( selected_surface[selection] == ROUTE_OUTER_PRIMARY ) else: direct_outer = direct_outer.unsqueeze(0).expand( selected_outer[selection].shape[0], -1, -1 ) outer_only[selection] = direct_outer secondary = routing.get("secondary", torch.zeros_like(selected_outer)) geometry_seed = selected_outer & (outer_only | secondary) supported_texel = _grouped_uv_support( geometry_seed, routing["flat_uv"], len(views), ) return selected_outer & supported_texel def _semantic_supported_outer_pixels( outputs, routing, group_size, presence_threshold=0.80, coverage_threshold=0.20, ): """Use image-level part semantics to identify plausible outer-layer pixels.""" selected_outer = routing["layer"] == 1 presence_logits = outputs.get("outer_presence_logits") predicted_coverage = outputs.get("outer_coverage") routed_part = routing.get("part") if ( presence_logits is None or predicted_coverage is None or routed_part is None ): return torch.zeros_like(selected_outer) group_count = selected_outer.shape[0] // group_size if presence_logits.shape[:2] != (group_count, PART_CLASSES): return torch.zeros_like(selected_outer) if predicted_coverage.shape[:2] != (group_count, PART_CLASSES): return torch.zeros_like(selected_outer) item_groups = torch.arange( selected_outer.shape[0], device=selected_outer.device ) // group_size safe_part = routed_part.clamp(0, PART_CLASSES - 1) valid_part = routed_part != IGNORE_INDEX expanded_groups = item_groups.view(-1, 1, 1).expand_as(safe_part) presence = torch.sigmoid(presence_logits.float())[ expanded_groups, safe_part ] coverage = predicted_coverage.float()[expanded_groups, safe_part] return ( selected_outer & valid_part & (presence >= float(presence_threshold)) & (coverage >= float(coverage_threshold)) ) def _scatter_soft_uv(rgb_sum, weight_sum, source_rgb, weight, flat_uv, valid): """Differentiably accumulate weighted render pixels into fixed UV indices.""" if not valid.any(): return rgb_sum, weight_sum batch = source_rgb.shape[0] indices = flat_uv[valid].reshape(1, -1).expand(batch, -1) selected_weight = weight[:, valid] selected_rgb = source_rgb[:, :, valid] * selected_weight.unsqueeze(1) weight_contribution = weight_sum.new_zeros(batch, UV_SIZE * UV_SIZE).scatter_add( 1, indices, selected_weight, ) rgb_contribution = rgb_sum.new_zeros(batch, 3, UV_SIZE * UV_SIZE).scatter_add( 2, indices.unsqueeze(1).expand(-1, 3, -1), selected_rgb, ) return rgb_sum + rgb_contribution, weight_sum + weight_contribution def soft_splat_geometry_predictions_to_uv( rendered, outputs, renderer, views, group_size=None, temperature=1.0, canonicalize=True, eps=1e-6, return_details=False, ): """Build a differentiable provisional UV atlas from geometry-parser probabilities. Primary inner/outer probabilities are splatted through their direct cuboid maps. Secondary probability is distributed through the learned exact surface probabilities. UV scatter indices stay fixed while all weights and the predicted affine remain differentiable. """ if rendered.dim() != 4: raise ValueError(f"Expected rendered NCHW tensor, got {tuple(rendered.shape)}.") views = parse_views(views) if not views: raise ValueError("At least one renderer view is required for soft geometry splatting.") if group_size is None: group_size = len(views) if group_size != len(views): raise ValueError(f"group_size={group_size} must equal the number of views ({len(views)}).") if rendered.shape[0] % group_size != 0: raise ValueError(f"N={rendered.shape[0]} must be divisible by group_size={group_size}.") if temperature <= 0: raise ValueError("temperature must be positive.") if outputs["layer"].shape[1] != ROUTE_ROLE_CLASSES or "surface" not in outputs: raise ValueError( "Soft geometry splatting requires three route-role logits and exact surface logits." ) if canonicalize: canonical_outputs = canonicalize_parser_outputs(outputs) canonical_rendered = canonicalize_parser_render( rendered, outputs, mode="bilinear", ).float() else: canonical_outputs = outputs canonical_rendered = rendered.float() role_prob = torch.softmax(canonical_outputs["layer"].float() / temperature, dim=1) surface_logits = canonical_outputs["surface"] surface_log_normalizer = torch.logsumexp(surface_logits / temperature, dim=1) foreground_prob = torch.sigmoid(canonical_outputs["foreground"][:, 0].float()) groups = rendered.shape[0] // group_size reference = canonical_rendered rgb_sums = [ reference.new_zeros(groups, 3, UV_SIZE * UV_SIZE) for _ in range(LAYER_CLASSES) ] weight_sums = [ reference.new_zeros(groups, UV_SIZE * UV_SIZE) for _ in range(LAYER_CLASSES) ] expected = reference.new_zeros(LAYER_CLASSES, UV_SIZE * UV_SIZE) for view_index, view in enumerate(views): selection = slice(view_index, rendered.shape[0], group_size) static = build_static_surface_routing(renderer, view, rendered.device) if static["masks"].shape[-2:] != canonical_rendered.shape[-2:]: raise ValueError( f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} " f"does not match parser shape {tuple(canonical_rendered.shape[-2:])}." ) surface_count = static["masks"].shape[0] if surface_logits.shape[1] < surface_count: raise ValueError( f"Parser has {surface_logits.shape[1]} surface classes, but view {view!r} " f"requires {surface_count}." ) source_rgb = canonical_rendered[selection, :3] view_foreground = foreground_prob[selection] view_role_prob = role_prob[selection] view_surface_logits = surface_logits[selection] view_surface_log_normalizer = surface_log_normalizer[selection] for layer_index in range(LAYER_CLASSES): valid = static["masks"][layer_index] flat_uv = static["flat_uv"][layer_index] weight = view_foreground * view_role_prob[:, layer_index] rgb_sums[layer_index], weight_sums[layer_index] = _scatter_soft_uv( rgb_sums[layer_index], weight_sums[layer_index], source_rgb, weight, flat_uv, valid, ) expected[layer_index].scatter_add_( 0, flat_uv[valid], torch.ones_like(flat_uv[valid], dtype=expected.dtype), ) candidate_layer, candidate_role = _static_candidate_roles(static) for surface_index in range(LAYER_CLASSES, surface_count): secondary = ( static["masks"][surface_index] & (candidate_role[surface_index] == ROUTE_SECONDARY) ) if not secondary.any(): continue surface_weight = ( view_foreground * view_role_prob[:, ROUTE_SECONDARY] * torch.exp( view_surface_logits[:, surface_index].float() / temperature - view_surface_log_normalizer.float() ) ) for layer_index in range(LAYER_CLASSES): valid = secondary & (candidate_layer[surface_index] == layer_index) if not valid.any(): continue rgb_sums[layer_index], weight_sums[layer_index] = _scatter_soft_uv( rgb_sums[layer_index], weight_sums[layer_index], source_rgb, surface_weight, static["flat_uv"][surface_index], valid, ) stacked_rgb_sum = torch.stack(rgb_sums, dim=1) stacked_weight_sum = torch.stack(weight_sums, dim=1).unsqueeze(2) expected = expected.view(1, LAYER_CLASSES, 1, UV_SIZE * UV_SIZE).clamp_min(1.0) layer_alpha = (stacked_weight_sum / expected).clamp(0.0, 1.0) total_rgb_sum = stacked_rgb_sum.sum(dim=1) total_weight = stacked_weight_sum.sum(dim=1) rgb = total_rgb_sum / total_weight.clamp_min(eps) alpha = layer_alpha.amax(dim=1) pred_uv = torch.cat([rgb, alpha], dim=1).reshape(groups, 4, UV_SIZE, UV_SIZE) support = total_weight.reshape(groups, 1, UV_SIZE, UV_SIZE) if return_details: return pred_uv, { "support": support, "layer_weight": stacked_weight_sum.reshape( groups, LAYER_CLASSES, UV_SIZE, UV_SIZE ), "layer_alpha": layer_alpha.reshape( groups, LAYER_CLASSES, UV_SIZE, UV_SIZE ), "canonical_rendered": canonical_rendered, "canonical_outputs": canonical_outputs, } return pred_uv def render_direct_uv(skins, renderer, view): """Differentiably render the direct inner/outer cuboids without deep composites.""" if skins.dim() != 4 or skins.shape[1:] != (4, UV_SIZE, UV_SIZE): raise ValueError(f"Expected skins shaped (N, 4, 64, 64), got {tuple(skins.shape)}.") batch = skins.shape[0] dtype = skins.dtype inner_grid = getattr(renderer, f"{view}_inner_grid").to(dtype=dtype) outer_grid = getattr(renderer, f"{view}_outer_grid").to(dtype=dtype) inner_mask = getattr(renderer, f"{view}_inner_mask").to(dtype=dtype) outer_mask = getattr(renderer, f"{view}_outer_mask").to(dtype=dtype) inner = F.grid_sample( skins, inner_grid.unsqueeze(0).expand(batch, -1, -1, -1), mode="bilinear", padding_mode="zeros", align_corners=True, ) * inner_mask.view(1, 1, *inner_mask.shape) outer = F.grid_sample( skins, outer_grid.unsqueeze(0).expand(batch, -1, -1, -1), mode="bilinear", padding_mode="zeros", align_corners=True, ) * outer_mask.view(1, 1, *outer_mask.shape) bg = renderer.bg_color.to(device=skins.device, dtype=dtype).view(1, 3, 1, 1) inner_rgb = inner[:, 3:4] * inner[:, :3] + (1.0 - inner[:, 3:4]) * bg rgb = outer[:, 3:4] * outer[:, :3] + (1.0 - outer[:, 3:4]) * inner_rgb alpha = outer[:, 3:4] + (1.0 - outer[:, 3:4]) * inner[:, 3:4] return torch.cat([rgb, alpha], dim=1) def render_outer_alpha_uv(skins, renderer, view): """Differentiably project only the outer-layer alpha into one view.""" if skins.dim() != 4 or skins.shape[1:] != (4, UV_SIZE, UV_SIZE): raise ValueError( f"Expected skins shaped (N, 4, 64, 64), got {tuple(skins.shape)}." ) batch = skins.shape[0] dtype = skins.dtype outer_grid = getattr(renderer, f"{view}_outer_grid").to(dtype=dtype) outer_mask = getattr(renderer, f"{view}_outer_mask").to(dtype=dtype) return F.grid_sample( skins[:, 3:4], outer_grid.unsqueeze(0).expand(batch, -1, -1, -1), mode="bilinear", padding_mode="zeros", align_corners=True, ) * outer_mask.view(1, 1, *outer_mask.shape) def flat_uv_to_uv01(flat_uv, dtype): return torch.stack( [ (flat_uv.remainder(UV_SIZE) / float(UV_SIZE - 1)).to(dtype), (flat_uv.div(UV_SIZE, rounding_mode="floor") / float(UV_SIZE - 1)).to(dtype), ], dim=1, ) def splat_parser_predictions_to_uv_conditioning( rendered, outputs, renderer=None, views=None, group_size=1, fg_threshold=0.5, bg_color=(128, 128, 128), semantic_gate=True, affine_refine=True, affine_refine_translation_px=8.0, affine_refine_scale=0.0, route_confidence_threshold=0.0, route_margin_threshold=0.0, outer_route_confidence_threshold=None, outer_route_margin_threshold=None, outer_uv_min_coverage=0.5, outer_uv_min_source_pixels=1, outer_silhouette_consistency=True, outer_silhouette_min_coverage=0.50, outer_silhouette_dilation=0, outer_silhouette_min_pixels=4, outer_geometry_rescue=False, outer_semantic_rescue=True, outer_semantic_presence_threshold=0.80, outer_semantic_coverage_threshold=0.20, outer_rescue_confidence_threshold=0.60, outer_rescue_margin_threshold=0.25, outer_rescue_min_coverage=0.10, color_aggregation="grid_mode", geometry_outer_threshold=0.5, geometry_route_texel_consensus=True, geometry_route_texel_consensus_weight=0.60, geometry_route_preserve_outer_confidence=0.80, geometry_route_preserve_outer_margin=0.35, geometry_route_consensus_outer_confidence=0.70, geometry_route_consensus_outer_margin=0.20, geometry_cross_view_outer_consistency=False, geometry_cross_view_outer_weight=0.50, geometry_cross_view_outer_positive_confidence=0.70, geometry_cross_view_outer_positive_margin=0.20, geometry_cross_view_outer_negative_confidence=0.70, geometry_cross_view_outer_negative_margin=0.20, geometry_cross_view_outer_background_max_coverage=0.25, geometry_cross_view_outer_min_views=2, outer_uv_occupancy=False, outer_uv_occupancy_blend_weight=0.0, outer_uv_occupancy_gate_threshold=0.15, outer_uv_occupancy_rescue_threshold=0.70, outer_uv_occupancy_rescue_route_threshold=0.30, outer_uv_component_routing=False, outer_uv_component_seed_threshold=0.80, outer_uv_component_grow_threshold=0.50, outer_uv_component_min_size=2, observed_foreground=None, background_color_tolerance=SOLID_BACKGROUND_COLOR_TOLERANCE, color_background_tolerance=None, color_foreground_inset=1, reject_semantic_fallback=False, reject_inner_semantic_fallback=False, include_rejected_context=True, rejected_context_confidence_threshold=0.35, rejected_context_margin_threshold=0.10, rejected_context_alpha_confidence_threshold=0.50, rejected_context_alpha_margin_threshold=0.10, include_confidence=False, return_details=False, observed_color_support=None, head_color_separation_inset=0, ): """Route parser outputs to UV, using static mappings for affine-parser checkpoints.""" if not 0.0 <= route_confidence_threshold <= 1.0: raise ValueError("route_confidence_threshold must be in [0, 1].") if not 0.0 <= route_margin_threshold <= 1.0: raise ValueError("route_margin_threshold must be in [0, 1].") if outer_route_confidence_threshold is None: outer_route_confidence_threshold = route_confidence_threshold if outer_route_margin_threshold is None: outer_route_margin_threshold = route_margin_threshold if not 0.0 <= outer_route_confidence_threshold <= 1.0: raise ValueError("outer_route_confidence_threshold must be in [0, 1].") if not 0.0 <= outer_route_margin_threshold <= 1.0: raise ValueError("outer_route_margin_threshold must be in [0, 1].") if not 0.0 <= outer_uv_min_coverage <= 1.0: raise ValueError("outer_uv_min_coverage must be in [0, 1].") if outer_uv_min_source_pixels < 1: raise ValueError("outer_uv_min_source_pixels must be positive.") if not 0.0 <= outer_silhouette_min_coverage <= 1.0: raise ValueError("outer_silhouette_min_coverage must be in [0, 1].") if outer_silhouette_dilation < 0: raise ValueError("outer_silhouette_dilation must be non-negative.") if outer_silhouette_min_pixels < 1: raise ValueError("outer_silhouette_min_pixels must be positive.") if not ( 0.0 <= outer_uv_component_grow_threshold <= outer_uv_component_seed_threshold <= 1.0 ): raise ValueError( "Require outer component grow <= seed threshold in [0, 1]." ) if outer_uv_component_min_size < 1: raise ValueError("outer_uv_component_min_size must be positive.") if not 0.0 <= outer_rescue_confidence_threshold <= 1.0: raise ValueError("outer_rescue_confidence_threshold must be in [0, 1].") if not 0.0 <= outer_rescue_margin_threshold <= 1.0: raise ValueError("outer_rescue_margin_threshold must be in [0, 1].") if not 0.0 <= outer_rescue_min_coverage <= 1.0: raise ValueError("outer_rescue_min_coverage must be in [0, 1].") if not 0.0 <= outer_semantic_presence_threshold <= 1.0: raise ValueError("outer_semantic_presence_threshold must be in [0, 1].") if not 0.0 <= outer_semantic_coverage_threshold <= 1.0: raise ValueError("outer_semantic_coverage_threshold must be in [0, 1].") for name, value in ( ( "geometry_route_texel_consensus_weight", geometry_route_texel_consensus_weight, ), ( "geometry_route_preserve_outer_confidence", geometry_route_preserve_outer_confidence, ), ( "geometry_route_preserve_outer_margin", geometry_route_preserve_outer_margin, ), ( "geometry_route_consensus_outer_confidence", geometry_route_consensus_outer_confidence, ), ( "geometry_route_consensus_outer_margin", geometry_route_consensus_outer_margin, ), ( "geometry_cross_view_outer_weight", geometry_cross_view_outer_weight, ), ( "geometry_cross_view_outer_positive_confidence", geometry_cross_view_outer_positive_confidence, ), ( "geometry_cross_view_outer_positive_margin", geometry_cross_view_outer_positive_margin, ), ( "geometry_cross_view_outer_negative_confidence", geometry_cross_view_outer_negative_confidence, ), ( "geometry_cross_view_outer_negative_margin", geometry_cross_view_outer_negative_margin, ), ( "geometry_cross_view_outer_background_max_coverage", geometry_cross_view_outer_background_max_coverage, ), ( "outer_uv_occupancy_blend_weight", outer_uv_occupancy_blend_weight, ), ( "outer_uv_occupancy_gate_threshold", outer_uv_occupancy_gate_threshold, ), ( "outer_uv_occupancy_rescue_threshold", outer_uv_occupancy_rescue_threshold, ), ( "outer_uv_occupancy_rescue_route_threshold", outer_uv_occupancy_rescue_route_threshold, ), ): if not 0.0 <= float(value) <= 1.0: raise ValueError(f"{name} must be in [0, 1].") if not 0.0 <= rejected_context_confidence_threshold <= 1.0: raise ValueError("rejected_context_confidence_threshold must be in [0, 1].") if not 0.0 <= rejected_context_margin_threshold <= 1.0: raise ValueError("rejected_context_margin_threshold must be in [0, 1].") if not 0.0 <= rejected_context_alpha_confidence_threshold <= 1.0: raise ValueError( "rejected_context_alpha_confidence_threshold must be in [0, 1]." ) if not 0.0 <= rejected_context_alpha_margin_threshold <= 1.0: raise ValueError( "rejected_context_alpha_margin_threshold must be in [0, 1]." ) if not 0.0 <= background_color_tolerance <= 1.0: raise ValueError("background_color_tolerance must be in [0, 1].") if ( color_background_tolerance is not None and not 0.0 <= color_background_tolerance <= 1.0 ): raise ValueError("color_background_tolerance must be in [0, 1].") if color_foreground_inset < 0: raise ValueError("color_foreground_inset must be non-negative.") if geometry_cross_view_outer_min_views < 2: raise ValueError("geometry_cross_view_outer_min_views must be at least 2.") if "affine" not in outputs: conditioning = splat_predictions_to_uv_conditioning( rendered, outputs, group_size=group_size, fg_threshold=fg_threshold, bg_color=bg_color, ) if return_details: return conditioning, {"rendered": rendered, "outputs": outputs, "routing": None, "alignment": None} return conditioning if renderer is None or views is None: raise ValueError("Affine parser routing requires both renderer and views.") views = parse_views(views) if group_size != len(views): raise ValueError(f"group_size={group_size} must equal the number of views ({len(views)}).") if observed_foreground is None: observed_foreground = estimate_solid_background_foreground( rendered, color_tolerance=background_color_tolerance, ) else: if observed_foreground.shape != rendered.shape[:1] + rendered.shape[-2:]: raise ValueError( "observed_foreground must have shape " f"{rendered.shape[:1] + rendered.shape[-2:]}, got {tuple(observed_foreground.shape)}." ) observed_foreground = observed_foreground.to( device=rendered.device, dtype=torch.bool ) source_background_rgb, _ = estimate_solid_background_color(rendered) source_fill = rendered.new_ones(rendered.shape[0], rendered.shape[1], 1, 1) source_fill[:, :3] = source_background_rgb.to(dtype=rendered.dtype) alignment = None routing_outputs = outputs if affine_refine: refined_affine, alignment = refine_parser_affine( outputs, renderer, views, translation_radius_px=affine_refine_translation_px, scale_radius=affine_refine_scale, observed_foreground=observed_foreground, ) routing_outputs = dict(outputs) routing_outputs["affine"] = refined_affine canonical_rendered = canonicalize_parser_render( rendered, routing_outputs, mode="nearest", fill_color=source_fill, ) canonical_observed_foreground = canonicalize_foreground_mask( observed_foreground, routing_outputs["affine"], ) canonical_observed_foreground_nearest = canonicalize_tensor( observed_foreground.unsqueeze(1).to(dtype=rendered.dtype), routing_outputs["affine"], mode="nearest", )[:, 0] > 0.5 canonical_foreground_coverage_rescued = ( canonical_observed_foreground & ~canonical_observed_foreground_nearest ) color_support = build_color_sampling_support( canonical_rendered, canonical_observed_foreground, source_background_rgb, background_tolerance=color_background_tolerance, foreground_inset=color_foreground_inset, ) if observed_color_support is not None: if observed_color_support.shape != observed_foreground.shape: raise ValueError('Observed color support must match observed foreground') safe_source=canonicalize_tensor(observed_color_support[:,None].to(rendered.device,rendered.dtype),routing_outputs['affine'],mode='nearest')[:,0]>.5 color_support['rejected'] |= color_support['valid'] & ~safe_source color_support['valid'] &= safe_source canonical_outputs = canonicalize_parser_outputs(routing_outputs) if "surface" in canonical_outputs and "part" not in canonical_outputs: routing = _routing_from_geometry_surface_outputs( renderer, views, canonical_outputs, fg_threshold=fg_threshold, texel_consensus=geometry_route_texel_consensus, texel_consensus_weight=geometry_route_texel_consensus_weight, preserve_outer_confidence=geometry_route_preserve_outer_confidence, preserve_outer_margin=geometry_route_preserve_outer_margin, consensus_outer_confidence=geometry_route_consensus_outer_confidence, consensus_outer_margin=geometry_route_consensus_outer_margin, outer_uv_occupancy=outer_uv_occupancy, outer_uv_occupancy_blend_weight=( outer_uv_occupancy_blend_weight ), outer_uv_occupancy_gate_threshold=( outer_uv_occupancy_gate_threshold ), outer_uv_occupancy_rescue_threshold=( outer_uv_occupancy_rescue_threshold ), outer_uv_occupancy_rescue_route_threshold=( outer_uv_occupancy_rescue_route_threshold ), outer_uv_component_routing=outer_uv_component_routing, outer_uv_component_seed_threshold=( outer_uv_component_seed_threshold ), outer_uv_component_grow_threshold=( outer_uv_component_grow_threshold ), outer_uv_component_min_size=outer_uv_component_min_size, observed_foreground=canonical_observed_foreground, cross_view_outer_consistency=( geometry_cross_view_outer_consistency ), cross_view_outer_weight=geometry_cross_view_outer_weight, cross_view_outer_positive_confidence=( geometry_cross_view_outer_positive_confidence ), cross_view_outer_positive_margin=( geometry_cross_view_outer_positive_margin ), cross_view_outer_negative_confidence=( geometry_cross_view_outer_negative_confidence ), cross_view_outer_negative_margin=( geometry_cross_view_outer_negative_margin ), cross_view_outer_background_max_coverage=( geometry_cross_view_outer_background_max_coverage ), cross_view_outer_min_views=geometry_cross_view_outer_min_views, ) elif "surface" not in canonical_outputs and "part" not in canonical_outputs: routing = _routing_from_geometry_outputs( renderer, views, canonical_outputs, fg_threshold=fg_threshold, outer_threshold=geometry_outer_threshold, texel_consensus=geometry_route_texel_consensus, texel_consensus_weight=geometry_route_texel_consensus_weight, preserve_outer_confidence=geometry_route_preserve_outer_confidence, preserve_outer_margin=geometry_route_preserve_outer_margin, consensus_outer_confidence=geometry_route_consensus_outer_confidence, consensus_outer_margin=geometry_route_consensus_outer_margin, outer_uv_occupancy=outer_uv_occupancy, outer_uv_occupancy_blend_weight=( outer_uv_occupancy_blend_weight ), outer_uv_occupancy_gate_threshold=( outer_uv_occupancy_gate_threshold ), outer_uv_occupancy_rescue_threshold=( outer_uv_occupancy_rescue_threshold ), outer_uv_occupancy_rescue_route_threshold=( outer_uv_occupancy_rescue_route_threshold ), outer_uv_component_routing=outer_uv_component_routing, outer_uv_component_seed_threshold=( outer_uv_component_seed_threshold ), outer_uv_component_grow_threshold=( outer_uv_component_grow_threshold ), outer_uv_component_min_size=outer_uv_component_min_size, observed_foreground=canonical_observed_foreground, cross_view_outer_consistency=( geometry_cross_view_outer_consistency ), cross_view_outer_weight=geometry_cross_view_outer_weight, cross_view_outer_positive_confidence=( geometry_cross_view_outer_positive_confidence ), cross_view_outer_positive_margin=( geometry_cross_view_outer_positive_margin ), cross_view_outer_negative_confidence=( geometry_cross_view_outer_negative_confidence ), cross_view_outer_negative_margin=( geometry_cross_view_outer_negative_margin ), cross_view_outer_background_max_coverage=( geometry_cross_view_outer_background_max_coverage ), cross_view_outer_min_views=geometry_cross_view_outer_min_views, ) else: routing = _routing_from_affine_outputs( renderer, views, canonical_outputs, fg_threshold=fg_threshold, semantic_gate=semantic_gate, ) from SkingToolkit.dense_uv_parser.accessories import apply_accessory_routing accessory_supported = apply_accessory_routing( routing, canonical_outputs, canonical_observed_foreground, renderer, views) raw_foreground = routing["foreground"] routing["secondary"] = routing.get( "secondary", torch.zeros_like(raw_foreground) ) & canonical_observed_foreground selected_outer = routing["layer"] == 1 confidence_threshold = torch.where( selected_outer, routing["confidence"].new_tensor(outer_route_confidence_threshold), routing["confidence"].new_tensor(route_confidence_threshold), ) margin_threshold = torch.where( selected_outer, routing["confidence_margin_ratio"].new_tensor(outer_route_margin_threshold), routing["confidence_margin_ratio"].new_tensor(route_margin_threshold), ) strict_trusted = ( raw_foreground & canonical_observed_foreground & (routing["confidence"] >= confidence_threshold) & (routing["confidence_margin_ratio"] >= margin_threshold) ) outer_consensus_gate_rejected = routing.get( "consensus_outer_gate_rejected", torch.zeros_like(selected_outer), ) # A deferred outer route is intentionally unknown. It may be restored by # exact geometry or atlas occupancy evidence below, but broad semantic part # presence must not turn the rejected local observation back into outer. deferred_outer_gate_rejected = ( selected_outer & outer_consensus_gate_rejected ) strict_trusted = strict_trusted & ~deferred_outer_gate_rejected if reject_semantic_fallback: rejected_fallback = routing["semantic_fallback"] & ( selected_outer | reject_inner_semantic_fallback ) strict_trusted = strict_trusted & ~rejected_fallback geometry_supported_outer = torch.zeros_like(selected_outer) semantic_supported_outer = torch.zeros_like(selected_outer) occupancy_supported_outer = torch.zeros_like(selected_outer) outer_geometry_rescued = torch.zeros_like(selected_outer) outer_semantic_rescued = torch.zeros_like(selected_outer) outer_occupancy_rescued = torch.zeros_like(selected_outer) if outer_geometry_rescue: geometry_supported_outer = _geometry_supported_outer_texels( routing, renderer, views, ) rescue_trusted = ( raw_foreground & canonical_observed_foreground & geometry_supported_outer & (routing["confidence"] >= outer_rescue_confidence_threshold) & ( routing["confidence_margin_ratio"] >= outer_rescue_margin_threshold ) ) if reject_semantic_fallback: rescue_trusted = rescue_trusted & ~rejected_fallback outer_geometry_rescued = rescue_trusted & ~strict_trusted trusted = strict_trusted | rescue_trusted else: trusted = strict_trusted if outer_semantic_rescue: semantic_supported_outer = _semantic_supported_outer_pixels( canonical_outputs, routing, group_size, presence_threshold=outer_semantic_presence_threshold, coverage_threshold=outer_semantic_coverage_threshold, ) semantic_rescue_trusted = ( raw_foreground & canonical_observed_foreground & semantic_supported_outer & ~deferred_outer_gate_rejected & (routing["confidence"] >= outer_rescue_confidence_threshold) & ( routing["confidence_margin_ratio"] >= outer_rescue_margin_threshold ) ) if reject_semantic_fallback: semantic_rescue_trusted = semantic_rescue_trusted & ~rejected_fallback outer_semantic_rescued = semantic_rescue_trusted & ~trusted trusted = trusted | semantic_rescue_trusted occupancy_available = routing.get( "outer_uv_occupancy_available", torch.zeros_like(selected_outer), ) occupancy_supported_outer = ( selected_outer & occupancy_available & ( routing.get( "projected_outer_occupancy", torch.zeros_like(routing["confidence"]), ) >= float(outer_uv_occupancy_rescue_threshold) ) ) if outer_uv_occupancy and occupancy_supported_outer.any(): occupancy_rescue_trusted = ( raw_foreground & canonical_observed_foreground & occupancy_supported_outer & (routing["confidence"] >= outer_rescue_confidence_threshold) & ( routing["confidence_margin_ratio"] >= outer_rescue_margin_threshold ) ) if reject_semantic_fallback: occupancy_rescue_trusted = ( occupancy_rescue_trusted & ~rejected_fallback ) outer_occupancy_rescued = occupancy_rescue_trusted & ~trusted trusted = trusted | occupancy_rescue_trusted trusted = trusted | accessory_supported outer_silhouette_coverage = torch.ones_like(routing["confidence"]) outer_silhouette_assessed = torch.zeros_like(selected_outer) outer_silhouette_rejected = torch.zeros_like(selected_outer) if outer_silhouette_consistency: ( outer_silhouette_coverage, outer_silhouette_assessed, ) = _outer_silhouette_coverage( # A one-pixel antialias/background fringe often survives flood # fill. It is unsafe silhouette evidence even when it is already # excluded from RGB pickup. Use the inset foreground here so that # such a fringe cannot validate a false outer shell. color_support["interior"], routing["flat_uv"], renderer, views, dilation=outer_silhouette_dilation, min_pixels=outer_silhouette_min_pixels, ) outer_silhouette_rejected = ( trusted & selected_outer & outer_silhouette_assessed & ( outer_silhouette_coverage < float(outer_silhouette_min_coverage) ) ) # Thin crown teeth need not fill most of a projected square. A learned # whole-object decision plus confident pixel semantics may preserve them; # foreground and valid outer geometry are still mandatory. All other # objects retain the established silhouette veto. crown_rescue = torch.zeros_like(trusted) if 'headwear_presence_logits' in canonical_outputs and 'headwear_supported' in routing: crown_rescue = (routing['headwear_supported'] & (routing['headwear_family']==4) & (routing['headwear_probability']>=.8) & canonical_observed_foreground & selected_outer) routing['headwear_silhouette_rescued'] = outer_silhouette_rejected & crown_rescue outer_silhouette_rejected &= ~crown_rescue trusted = trusted & ~outer_silhouette_rejected outer_uv_coverage = torch.ones_like(routing["confidence"]) outer_required_coverage = torch.zeros_like(routing["confidence"]) if outer_uv_min_coverage > 0.0: views_per_group = len(views) group_count = trusted.shape[0] // views_per_group if "surface" in canonical_outputs: pixel_coverage = _surface_aware_outer_coverage( routing, trusted, selected_outer, renderer, views, ) else: expected = routing["confidence"].new_zeros(UV_SIZE * UV_SIZE) for view in views: static = build_static_surface_routing(renderer, view, trusted.device) expected_uv = static["flat_uv"][1][static["masks"][1]] expected.scatter_add_( 0, expected_uv, torch.ones(expected_uv.shape[0], device=trusted.device, dtype=expected.dtype), ) item_groups = torch.arange(trusted.shape[0], device=trusted.device) // views_per_group grouped_uv = routing["flat_uv"] + item_groups.view(-1, 1, 1) * (UV_SIZE * UV_SIZE) observed_mask = trusted & selected_outer observed = routing["confidence"].new_zeros(group_count * UV_SIZE * UV_SIZE) observed_uv = grouped_uv[observed_mask] observed.scatter_add_( 0, observed_uv, torch.ones(observed_uv.shape[0], device=trusted.device, dtype=observed.dtype), ) coverage = observed.reshape(group_count, UV_SIZE * UV_SIZE) / expected.clamp_min(1.0) pixel_coverage = coverage[item_groups.view(-1, 1, 1), routing["flat_uv"]] outer_uv_coverage = torch.where(selected_outer, pixel_coverage, outer_uv_coverage) required_coverage = torch.full_like(pixel_coverage, outer_uv_min_coverage) if ( outer_geometry_rescue or outer_semantic_rescue or outer_uv_occupancy ): required_coverage = torch.where( geometry_supported_outer | semantic_supported_outer | occupancy_supported_outer, required_coverage.new_tensor( min(outer_uv_min_coverage, outer_rescue_min_coverage) ), required_coverage, ) outer_required_coverage = torch.where( selected_outer, required_coverage, outer_required_coverage, ) trusted = trusted & (~selected_outer | (pixel_coverage >= required_coverage) | accessory_supported) outer_source_support = torch.zeros_like(routing["flat_uv"]) outer_source_rejected = torch.zeros_like(selected_outer) if outer_uv_min_source_pixels > 1: views_per_group = len(views) group_count = trusted.shape[0] // views_per_group item_groups = ( torch.arange(trusted.shape[0], device=trusted.device) // views_per_group ) grouped_uv = routing["flat_uv"] + item_groups.view(-1, 1, 1) * ( UV_SIZE * UV_SIZE ) outer_candidates = trusted & selected_outer source_counts = torch.zeros( group_count * UV_SIZE * UV_SIZE, dtype=torch.long, device=trusted.device, ) source_counts.scatter_add_( 0, grouped_uv[outer_candidates], torch.ones_like(grouped_uv[outer_candidates]), ) outer_source_support = source_counts[grouped_uv] outer_source_rejected = ( outer_candidates & (outer_source_support < int(outer_uv_min_source_pixels)) & ~accessory_supported ) trusted = trusted & ~outer_source_rejected routing["raw_foreground"] = raw_foreground routing["observed_foreground"] = canonical_observed_foreground routing["canonical_foreground_coverage_rescued"] = ( canonical_foreground_coverage_rescued ) routing["background_rejected"] = raw_foreground & ~canonical_observed_foreground routing["rejected"] = raw_foreground & ~trusted routing["outer_geometry_supported"] = geometry_supported_outer routing["outer_geometry_rescued"] = outer_geometry_rescued & trusted routing["outer_semantic_supported"] = semantic_supported_outer routing["outer_semantic_rescued"] = outer_semantic_rescued & trusted routing["outer_occupancy_supported"] = occupancy_supported_outer routing["outer_occupancy_rescued"] = outer_occupancy_rescued & trusted routing["outer_required_coverage"] = outer_required_coverage routing["outer_silhouette_coverage"] = outer_silhouette_coverage routing["outer_silhouette_assessed"] = outer_silhouette_assessed routing["outer_silhouette_rejected"] = outer_silhouette_rejected routing["outer_source_support"] = outer_source_support routing["outer_source_rejected"] = outer_source_rejected routing["foreground"] = trusted color_foreground = trusted & color_support["valid"] boundary_excluded = torch.zeros_like(trusted) if head_color_separation_inset > 0: # A layer boundary is also an RGB mixing boundary. Keep occupancy, but # do not let a thin contaminated inner fringe become an inpainting seed. radius=int(head_color_separation_inset) head_outer=trusted&(routing['layer']==1)&(routing['part']==0) near_outer=F.max_pool2d(head_outer[:,None].float(),2*radius+1,1,radius)[:,0]>.5 boundary_excluded=near_outer&(routing['layer']==0)&(routing['part']==0) phone=routing.get('headphone_colour_conflict') if phone is not None: # Compare a semantic accessory observation with nearby accepted outer # geometry; exclude contradictory RGB without inventing outer texels. nearby=F.max_pool2d(head_outer[:,None].float(),17,1,8)[:,0]>.5 conflict=phone&nearby&(routing['layer']==0)&(routing['part']==0) boundary_excluded|=conflict routing['headphone_colour_rejected']=conflict color_foreground &= ~boundary_excluded headwear_layer = routing.get('headwear_layer') if headwear_layer is not None: # A component outside the chosen layer's projected support is unknown; # sampling it into the other layer would duplicate its material. conflict = (headwear_layer >= 0) & (routing['layer'] != headwear_layer) & (routing['part'] == 0) color_foreground &= ~conflict family = routing['headwear_family'] boundary = torch.zeros_like(trusted) for component in (1,2,3,4): region = (family == component) & routing['headwear_supported'] interior = -F.max_pool2d(-region[:,None].float(),3,1,1)[:,0] > .5 boundary |= region & ~interior color_foreground &= ~boundary routing['headwear_color_rejected'] = conflict | boundary routing['head_color_boundary_excluded']=boundary_excluded routing["color_foreground"] = color_foreground routing["color_interior"] = color_support["interior"] routing["color_background_like"] = color_support["background_like"] routing["color_rejected"] = trusted & color_support["rejected"] if "secondary_routed" in routing: routing["secondary_rejected"] = routing["secondary"] & ~trusted routing["secondary_routed"] = routing["secondary"] & trusted routing["outer_uv_coverage"] = outer_uv_coverage conditioning = splat_to_uv_conditioning( canonical_rendered, color_foreground, routing["layer"], routing["flat_uv"], group_size=group_size, bg_color=bg_color, confidence=routing["confidence"], sampling_quality=( routing.get("texel_center_score") + 0.05 * color_support["interior"].to( dtype=routing["confidence"].dtype ) ), color_aggregation=color_aggregation, include_confidence=include_confidence, ) rejected_context = ( raw_foreground & canonical_observed_foreground & ~trusted & color_support["valid"] & (routing["confidence"] >= rejected_context_confidence_threshold) & ( routing["confidence_margin_ratio"] >= rejected_context_margin_threshold ) ) if reject_semantic_fallback: rejected_context = rejected_context & ~rejected_fallback rejected_context_alpha_supported = ( rejected_context & selected_outer & (geometry_supported_outer | semantic_supported_outer) & ( routing["confidence"] >= rejected_context_alpha_confidence_threshold ) & ( routing["confidence_margin_ratio"] >= rejected_context_alpha_margin_threshold ) ) context_alpha_rescue_uv = conditioning.new_zeros( conditioning.shape[0], 1, UV_SIZE, UV_SIZE ) if include_rejected_context and include_confidence and rejected_context.any(): context = splat_to_uv_conditioning( canonical_rendered, rejected_context, routing["layer"], routing["flat_uv"], group_size=group_size, bg_color=bg_color, confidence=routing["confidence"], sampling_quality=( routing.get("texel_center_score") + 0.05 * color_support["interior"].to( dtype=routing["confidence"].dtype ) ), color_aggregation=color_aggregation, include_confidence=True, ) for offset in (0, 6): trusted_evidence = conditioning[:, offset + 4 : offset + 5] > 0.5 context_evidence = context[:, offset + 4 : offset + 5] > 0.5 fill = ~trusted_evidence & context_evidence conditioning[:, offset : offset + 4] = torch.where( fill.expand(-1, 4, -1, -1), context[:, offset : offset + 4], conditioning[:, offset : offset + 4], ) conditioning[:, offset + 5 : offset + 6] = torch.where( fill, context[:, offset + 5 : offset + 6], conditioning[:, offset + 5 : offset + 6], ) if rejected_context_alpha_supported.any(): alpha_context = splat_to_uv_conditioning( canonical_rendered, rejected_context_alpha_supported, routing["layer"], routing["flat_uv"], group_size=group_size, bg_color=bg_color, confidence=routing["confidence"], sampling_quality=( routing.get("texel_center_score") + 0.05 * color_support["interior"].to( dtype=routing["confidence"].dtype ) ), color_aggregation=color_aggregation, include_confidence=True, ) context_alpha_rescue_uv = alpha_context[:, 10:11] > 0.5 routing["rejected_context"] = rejected_context routing["rejected_context_alpha_supported"] = ( rejected_context_alpha_supported ) if return_details: return conditioning, { "rendered": canonical_rendered, "color_source_support": color_support["valid"], "outputs": canonical_outputs, "routing": routing, "alignment": alignment, "context_alpha_rescue_uv": context_alpha_rescue_uv, } return conditioning def splat_deterministic_targets_to_uv_conditioning( rendered, targets, renderer, views, group_size=1, bg_color=(128, 128, 128), color_aggregation="grid_mode", ): """Splat ground-truth labels through the same fixed mapping used by affine mode.""" views = parse_views(views) if group_size != len(views): raise ValueError(f"group_size={group_size} must equal the number of views ({len(views)}).") canonical_rendered = canonicalize_tensor(rendered, targets["affine"], mode="nearest") canonical_targets = canonicalize_dense_targets(targets) requested_surface = canonical_targets["surface"] N, H, W = requested_surface.shape fg = torch.zeros_like(requested_surface, dtype=torch.bool) routed_layer = torch.zeros_like(requested_surface) flat_uv = torch.zeros_like(requested_surface) texel_center_score = canonical_rendered.new_zeros(requested_surface.shape) for view_index, view in enumerate(views): selection = slice(view_index, N, len(views)) static = build_static_surface_routing(renderer, view, requested_surface.device) if static["masks"].shape[-2:] != (H, W): raise ValueError( f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} does not match target shape {(H, W)}." ) routed = _select_static_surface(static, requested_surface[selection]) fg[selection] = ( (canonical_targets["foreground"][selection, 0] > 0.5) & (requested_surface[selection] != IGNORE_INDEX) & routed["valid"] ) routed_layer[selection] = routed["layer"] flat_uv[selection] = routed["flat_uv"] texel_center_score[selection] = routed["texel_center_score"] return splat_to_uv_conditioning( canonical_rendered, fg, routed_layer, flat_uv, group_size=group_size, bg_color=bg_color, sampling_quality=texel_center_score, color_aggregation=color_aggregation, ) def splat_predictions_to_uv_conditioning( rendered, outputs, group_size=1, fg_threshold=0.5, bg_color=(128, 128, 128), ): """Splat parser predictions into the parser's layered UV evidence layout.""" foreground_prob = torch.sigmoid(outputs["foreground"])[:, 0] fg = foreground_prob > fg_threshold fg = fg & (rendered[:, 3] > 1e-4) layer_prob = torch.softmax(outputs["layer"], dim=1) layer_confidence, layer = layer_prob.max(dim=1) flat_uv = prediction_flat_uv(outputs) confidence = foreground_prob * layer_confidence if "uv_x" in outputs and "uv_y" in outputs: uv_x_confidence = torch.softmax(outputs["uv_x"], dim=1).amax(dim=1) uv_y_confidence = torch.softmax(outputs["uv_y"], dim=1).amax(dim=1) confidence = confidence * torch.sqrt((uv_x_confidence * uv_y_confidence).clamp_min(1e-8)) return splat_to_uv_conditioning( rendered, fg, layer, flat_uv, group_size=group_size, bg_color=bg_color, confidence=confidence, ) def prediction_flat_uv(outputs): if "uv_x" in outputs and "uv_y" in outputs: x = outputs["uv_x"].argmax(dim=1).clamp(0, UV_SIZE - 1) y = outputs["uv_y"].argmax(dim=1).clamp(0, UV_SIZE - 1) else: uv = outputs["uv"].clamp(0.0, 1.0) x = (uv[:, 0] * (UV_SIZE - 1)).round().long().clamp(0, UV_SIZE - 1) y = (uv[:, 1] * (UV_SIZE - 1)).round().long().clamp(0, UV_SIZE - 1) return y * UV_SIZE + x def prediction_uv01(outputs): if "uv_x" in outputs and "uv_y" in outputs: x = outputs["uv_x"].argmax(dim=1).to(dtype=outputs["uv"].dtype) / (UV_SIZE - 1) y = outputs["uv_y"].argmax(dim=1).to(dtype=outputs["uv"].dtype) / (UV_SIZE - 1) return torch.stack([x, y], dim=1) return outputs["uv"].clamp(0.0, 1.0) def splat_targets_to_uv_conditioning(rendered, targets, group_size=1, bg_color=(128, 128, 128)): if rendered.dim() != 4: raise ValueError(f"Expected rendered tensor as NCHW, got {tuple(rendered.shape)}.") layer = targets["layer"] fg = (targets["foreground"][:, 0] > 0.5) & (layer != IGNORE_INDEX) & (rendered[:, 3] > 1e-4) uv = targets["uv"].clamp(0.0, 1.0) x = (uv[:, 0] * (UV_SIZE - 1)).round().long().clamp(0, UV_SIZE - 1) y = (uv[:, 1] * (UV_SIZE - 1)).round().long().clamp(0, UV_SIZE - 1) flat_uv = y * UV_SIZE + x safe_layer = torch.where(layer == IGNORE_INDEX, torch.zeros_like(layer), layer) return splat_to_uv_conditioning(rendered, fg, safe_layer, flat_uv, group_size=group_size, bg_color=bg_color) def _select_exact_mode_candidates( values, target_uv, scores, color_tie_break=None, ): """Select one real source pixel from the most frequent 8-bit RGB value per UV texel.""" if target_uv.numel() == 0: return values, target_uv, scores rgb8 = (values[:3].clamp(0.0, 1.0) * 255.0).round().long() color_code = rgb8[0] | (rgb8[1] << 8) | (rgb8[2] << 16) color_count = 1 << 24 combined_key = target_uv.long() * color_count + color_code unique_keys, inverse = torch.unique(combined_key, sorted=False, return_inverse=True) votes = scores.new_zeros(unique_keys.shape[0]) votes.index_add_(0, inverse, torch.ones_like(scores)) tie_break_votes = None if color_tie_break is not None: tie_break_votes = scores.new_zeros(unique_keys.shape[0]) tie_break_votes.index_add_(0, inverse, color_tie_break.float()) positions = torch.arange(target_uv.shape[0], device=target_uv.device) first_position = torch.full( (unique_keys.shape[0],), target_uv.shape[0], dtype=torch.long, device=target_uv.device, ) first_position.scatter_reduce_(0, inverse, positions, reduce="amin", include_self=True) unique_uv = unique_keys.div(color_count, rounding_mode="floor") best_votes = votes.new_full((UV_SIZE * UV_SIZE,), -torch.inf) best_votes.scatter_reduce_(0, unique_uv, votes, reduce="amax", include_self=True) tied_winner = votes >= (best_votes[unique_uv] - 1e-7) if tie_break_votes is not None: best_tie_break = tie_break_votes.new_full( (UV_SIZE * UV_SIZE,), -torch.inf ) best_tie_break.scatter_reduce_( 0, unique_uv[tied_winner], tie_break_votes[tied_winner], reduce="amax", include_self=True, ) tied_winner = tied_winner & ( tie_break_votes >= (best_tie_break[unique_uv] - 1e-7) ) earliest_winner = torch.full( (UV_SIZE * UV_SIZE,), target_uv.shape[0], dtype=torch.long, device=target_uv.device, ) earliest_winner.scatter_reduce_( 0, unique_uv[tied_winner], first_position[tied_winner], reduce="amin", include_self=True, ) winning_unique = tied_winner & (first_position == earliest_winner[unique_uv]) winning_key = torch.full( (UV_SIZE * UV_SIZE,), -1, dtype=torch.long, device=target_uv.device, ) winning_key[unique_uv[winning_unique]] = unique_keys[winning_unique] in_winning_color = combined_key == winning_key[target_uv] values = values[:, in_winning_color] target_uv = target_uv[in_winning_color] scores = scores[in_winning_color] best_scores = scores.new_full((UV_SIZE * UV_SIZE,), -torch.inf) best_scores.scatter_reduce_(0, target_uv, scores, reduce="amax", include_self=True) tied_representative = scores >= (best_scores[target_uv] - 1e-7) positions = torch.arange(target_uv.shape[0], device=target_uv.device) first_representative = torch.full( (UV_SIZE * UV_SIZE,), target_uv.shape[0], dtype=torch.long, device=target_uv.device, ) first_representative.scatter_reduce_( 0, target_uv[tied_representative], positions[tied_representative], reduce="amin", include_self=True, ) selected = tied_representative & (positions == first_representative[target_uv]) return values[:, selected], target_uv[selected], scores[selected] def splat_to_uv_conditioning( rendered, fg, layer, flat_uv, group_size=1, bg_color=(128, 128, 128), confidence=None, sampling_quality=None, color_aggregation="grid_mode", include_confidence=False, ): if rendered.dim() != 4: raise ValueError(f"Expected rendered tensor as NCHW, got {tuple(rendered.shape)}.") N, _, _, _ = rendered.shape if N % group_size != 0: raise ValueError(f"N={N} must be divisible by group_size={group_size}.") if color_aggregation not in SPLAT_COLOR_AGGREGATIONS: raise ValueError( f"Unknown color_aggregation={color_aggregation!r}; " f"expected one of {SPLAT_COLOR_AGGREGATIONS}." ) groups = N // group_size device = rendered.device dtype = rendered.dtype accum = rendered.new_zeros(groups, LAYER_CLASSES, 4, UV_SIZE * UV_SIZE) counts = rendered.new_zeros(groups, LAYER_CLASSES, 1, UV_SIZE * UV_SIZE) confidence_map = rendered.new_zeros( groups, LAYER_CLASSES, 1, UV_SIZE * UV_SIZE ) if confidence is None: confidence = rendered.new_ones(fg.shape) select_highest_confidence = False else: select_highest_confidence = True if sampling_quality is None: sampling_quality = rendered.new_ones(fg.shape) elif sampling_quality.shape != fg.shape: raise ValueError("sampling_quality must have the same shape as fg.") for group in range(groups): group_start = group * group_size group_end = group_start + group_size for layer_index in range(LAYER_CLASSES): candidate_values = [] candidate_uv = [] candidate_confidence = [] candidate_sampling_quality = [] for item in range(group_start, group_end): item_mask = fg[item] & (layer[item] == layer_index) if not item_mask.any(): continue candidate_values.append(rendered[item, :, item_mask]) candidate_uv.append(flat_uv[item, item_mask]) candidate_confidence.append(confidence[item, item_mask]) candidate_sampling_quality.append( sampling_quality[item, item_mask] ) if not candidate_uv: continue values = torch.cat(candidate_values, dim=1) target_uv = torch.cat(candidate_uv, dim=0) scores = torch.cat(candidate_confidence, dim=0).float() quality = torch.cat(candidate_sampling_quality, dim=0).float() if color_aggregation == "texel_center": # Lexicographic selection: prefer the source sample closest to # the integer UV center, then use parser confidence only to # break effectively equal geometric samples. This recovers an # original texel from bilinear renders instead of choosing a # top/left boundary blend by scan order. selection_score = quality + scores.clamp(0.0, 1.0) * 1e-6 best_quality = selection_score.new_full( (UV_SIZE * UV_SIZE,), -torch.inf ) best_quality.scatter_reduce_( 0, target_uv, selection_score, reduce="amax", include_self=True, ) is_best = selection_score >= ( best_quality[target_uv] - 1e-7 ) candidate_indices = torch.arange( target_uv.shape[0], device=device ) first_best = torch.full( (UV_SIZE * UV_SIZE,), target_uv.shape[0], dtype=torch.long, device=device, ) first_best.scatter_reduce_( 0, target_uv[is_best], candidate_indices[is_best], reduce="amin", include_self=True, ) selected_indices = first_best[ first_best < target_uv.shape[0] ] values = values[:, selected_indices] target_uv = target_uv[selected_indices] scores = scores[selected_indices] if color_aggregation == "grid_mode": # The fitted renderer grid owns UV membership. Select the # dominant real 8-bit RGB within each layer/UV cell; UV-center # quality is only a tie-break between equally frequent colors. values, target_uv, scores = _select_exact_mode_candidates( values, target_uv, scores, color_tie_break=quality, ) if color_aggregation == "exact_mode": values, target_uv, scores = _select_exact_mode_candidates( values, target_uv, scores ) if color_aggregation == "best" and select_highest_confidence: best_scores = scores.new_full((UV_SIZE * UV_SIZE,), -torch.inf) best_scores.scatter_reduce_(0, target_uv, scores, reduce="amax", include_self=True) is_best = scores >= (best_scores[target_uv] - 1e-7) candidate_indices = torch.arange(target_uv.shape[0], device=device) first_best = torch.full( (UV_SIZE * UV_SIZE,), target_uv.shape[0], dtype=torch.long, device=device, ) first_best.scatter_reduce_( 0, target_uv[is_best], candidate_indices[is_best], reduce="amin", include_self=True, ) selected_indices = first_best[first_best < target_uv.shape[0]] values = values[:, selected_indices] target_uv = target_uv[selected_indices] scores = scores[selected_indices] accum[group, layer_index].index_add_(1, target_uv, values) counts[group, layer_index, 0].index_add_( 0, target_uv, torch.ones(target_uv.shape[0], dtype=dtype, device=device), ) confidence_map[group, layer_index, 0].scatter_reduce_( 0, target_uv, scores.to(dtype=dtype), reduce="amax", include_self=True, ) known = (counts > 0).to(dtype=dtype) aggregated = accum / counts.clamp_min(1.0) bg = rendered.new_tensor(bg_color, dtype=dtype).view(1, 1, 3, 1) / 255.0 rgb = torch.where( known.expand(-1, -1, 3, -1) > 0, aggregated[:, :, :3], bg.expand(groups, LAYER_CLASSES, 3, UV_SIZE * UV_SIZE), ) alpha = torch.where(known > 0, aggregated[:, :, 3:4], torch.zeros_like(aggregated[:, :, 3:4])) layer_channels = [rgb, alpha, known] if include_confidence: layer_channels.append(confidence_map * known) layers = torch.cat(layer_channels, dim=2).reshape( groups, -1, UV_SIZE, UV_SIZE ) return layers.clamp(0.0, 1.0) def conditioning_to_pred_uv(conditioning): """Merge two-layer parser conditioning into a preliminary RGBA skin atlas. Conditioning stores either ``RGBA + evidence`` (legacy 10-channel input) or ``RGBA + evidence + confidence`` (12-channel input) for each layer. Their Minecraft UV rectangles normally do not overlap, but outer data wins if both layers mark the same texel as observed. Unknown RGB retains the conditioning background as a useful placeholder; unknown alpha remains transparent until Minecraft base-alpha finalization. """ squeeze_batch = conditioning.dim() == 3 if squeeze_batch: conditioning = conditioning.unsqueeze(0) if conditioning.dim() != 4 or conditioning.shape[1] not in (10, 12): raise ValueError( "Expected 10- or 12-channel parser conditioning, " f"got {tuple(conditioning.shape)}." ) inner_rgba = conditioning[:, 0:4] inner_known = conditioning[:, 4:5] > 0.5 outer_offset = 6 if conditioning.shape[1] == 12 else 5 outer_rgba = conditioning[:, outer_offset : outer_offset + 4] outer_known = conditioning[:, outer_offset + 4 : outer_offset + 5] > 0.5 known = inner_known | outer_known rgba = torch.where(outer_known.expand_as(outer_rgba), outer_rgba, inner_rgba) rgba[:, 3:4] = torch.where( known, rgba[:, 3:4], torch.zeros_like(rgba[:, 3:4]), ) rgba = rgba.clamp(0.0, 1.0) return rgba[0] if squeeze_batch else rgba PART_PALETTE = ( (239, 83, 80), (255, 202, 40), (102, 187, 106), (38, 166, 154), (66, 165, 245), (171, 71, 188), ) LAYER_PALETTE = ( (72, 169, 166), (255, 179, 71), ) ROUTE_ROLE_PALETTE = ( (72, 169, 166), (255, 179, 71), (218, 112, 232), ) FACE_NAMES = ("front", "back", "left", "right", "top", "bottom") FACE_PALETTE = ( (239, 83, 80), (171, 71, 188), (66, 165, 245), (38, 166, 154), (255, 202, 40), (117, 117, 117), ) LAYER_FACE_PALETTE = ( (183, 28, 28), (106, 27, 154), (21, 101, 192), (0, 121, 107), (245, 166, 35), (84, 84, 84), (255, 128, 125), (218, 112, 232), (126, 208, 255), (102, 221, 203), (255, 224, 130), (210, 210, 210), ) def combine_layer_face(layer, face): """Encode inner/outer x six cube faces as 12 visualization classes.""" valid = (layer != IGNORE_INDEX) & (face != IGNORE_INDEX) combined = torch.full_like(layer, IGNORE_INDEX) combined[valid] = layer[valid] * FACE_CLASSES + face[valid] return combined def bg_tensor(bg_color, reference): return reference.new_tensor(bg_color).view(1, 3, 1, 1) / 255.0 def colorize_labels(labels, palette, bg_color, reference): N, H, W = labels.shape bg = bg_tensor(bg_color, reference).expand(N, 3, H, W).clone() valid = labels != IGNORE_INDEX if not valid.any(): return bg palette_tensor = reference.new_tensor(palette, dtype=reference.dtype) / 255.0 safe_labels = labels.clamp(0, len(palette) - 1) out = bg.permute(0, 2, 3, 1) out[valid] = palette_tensor[safe_labels[valid]] return bg def build_geometry_grid_debug(renderer, views, item_count, reference, bg_color=(128, 128, 128)): """Build fitted inner/outer cuboid face maps with projected UV texel boundaries.""" views = parse_views(views) if not views: raise ValueError("At least one view is required for geometry-grid debug output.") colors = [[], []] masks = [[], []] edges = [[], []] for item_index in range(item_count): static = build_static_surface_routing(renderer, views[item_index % len(views)], reference.device) for layer_index in range(LAYER_CLASSES): mask = static["masks"][layer_index] face = static["face"][layer_index] flat_uv = static["flat_uv"][layer_index] labels = combine_layer_face( torch.full_like(face, layer_index), torch.where(mask, face, torch.full_like(face, IGNORE_INDEX)), ).unsqueeze(0) color = colorize_labels(labels, LAYER_FACE_PALETTE, bg_color, reference[:1])[0] edge = torch.zeros_like(mask) different_x = (flat_uv[:, 1:] != flat_uv[:, :-1]) | (face[:, 1:] != face[:, :-1]) different_y = (flat_uv[1:, :] != flat_uv[:-1, :]) | (face[1:, :] != face[:-1, :]) edge[:, 1:] |= mask[:, 1:] & mask[:, :-1] & different_x edge[1:, :] |= mask[1:, :] & mask[:-1, :] & different_y interior = -F.max_pool2d(-mask.float().unsqueeze(0).unsqueeze(0), 3, 1, 1)[0, 0] edge |= mask & (interior < 0.5) edge_color = reference.new_tensor((0.08, 0.08, 0.08)).view(3, 1, 1) color = torch.where(edge.unsqueeze(0), edge_color, color) colors[layer_index].append(color) masks[layer_index].append(mask) edges[layer_index].append(edge) return ( torch.stack(colors[0]), torch.stack(colors[1]), torch.stack(masks[0]), torch.stack(masks[1]), torch.stack(edges[0]), torch.stack(edges[1]), ) def fill_geometry_grid_debug(rendered, foreground, layer, geometry_debug, bg_color=(128, 128, 128)): """Show only source RGB pixels actually routed to each predicted layer.""" del geometry_debug rgb = rendered[:, :3] bg = rgb.new_tensor(bg_color).view(1, 3, 1, 1) / 255.0 def fill(layer_index): classified = (foreground & (layer == layer_index)).unsqueeze(1) return torch.where(classified, rgb, bg.expand_as(rgb)) return fill(0), fill(1) def overlay_geometry_grid_debug( rendered, geometry_debug, tint_alpha=0.12, base_images=None, ): """Overlay fitted inner/outer UV texel grids on the canonical source image.""" if not 0.0 <= tint_alpha <= 1.0: raise ValueError(f"tint_alpha must be in [0, 1], got {tint_alpha}.") rgb = rendered[:, :3] if base_images is None: base_images = (rgb, rgb) if len(base_images) != LAYER_CLASSES: raise ValueError( f"base_images must contain {LAYER_CLASSES} tensors, got {len(base_images)}." ) masks = geometry_debug[2:4] edges = geometry_debug[4:6] fill_colors = ( rgb.new_tensor((0.10, 0.85, 0.95)), rgb.new_tensor((1.00, 0.65, 0.10)), ) edge_colors = ( rgb.new_tensor((0.00, 1.00, 1.00)), rgb.new_tensor((1.00, 0.25, 0.85)), ) overlays = [] for base, mask, edge, fill_color, edge_color in zip( base_images, masks, edges, fill_colors, edge_colors ): if base.shape != rgb.shape: raise ValueError( f"Geometry overlay base shape {tuple(base.shape)} does not match RGB {tuple(rgb.shape)}." ) fill_color = fill_color.view(1, 3, 1, 1) edge_color = edge_color.view(1, 3, 1, 1) tinted = base * (1.0 - tint_alpha) + fill_color * tint_alpha overlay = torch.where(mask.unsqueeze(1), tinted, base) overlay = torch.where(edge.unsqueeze(1), edge_color, overlay) overlays.append(overlay) return tuple(overlays) def colorize_surface(labels, bg_color, reference): """Colorize fixed renderer surface slots without assuming a fixed slot count.""" N, H, W = labels.shape bg = bg_tensor(bg_color, reference).expand(N, 3, H, W).clone() valid = labels != IGNORE_INDEX if not valid.any(): return bg surface = labels.clamp_min(0).to(dtype=reference.dtype) colors = torch.stack( [ (surface.mul(47.0).remainder(191.0) + 32.0) / 255.0, (surface.mul(83.0).remainder(191.0) + 32.0) / 255.0, (surface.mul(131.0).remainder(191.0) + 32.0) / 255.0, ], dim=1, ) return torch.where(valid.unsqueeze(1), colors, bg) def colorize_foreground(mask, bg_color, reference): N, H, W = mask.shape bg = bg_tensor(bg_color, reference).expand(N, 3, H, W) fg = reference.new_ones(N, 3, H, W) return torch.where(mask.unsqueeze(1), fg, bg) def colorize_uv(uv, mask, bg_color): N, _, H, W = uv.shape bg = bg_tensor(bg_color, uv).expand(N, 3, H, W) zeros = uv.new_zeros(N, 1, H, W) uv_rgb = torch.cat([uv[:, 0:1], uv[:, 1:2], zeros], dim=1) return torch.where(mask.unsqueeze(1), uv_rgb, bg)