Download SKING_DDJ_v101c/runtime/dense_uv_parser/utils.py from EntropyDrop/Sking: direct link, hf CLI and curl.
- Browser
- Download file 186 kB
-
https://huggingface.co/EntropyDrop/Sking/resolve/main/SKING_DDJ_v101c/runtime/dense_uv_parser/utils.py
- Command line
-
hf download hf://EntropyDrop/Sking/SKING_DDJ_v101c/runtime/dense_uv_parser/utils.py
-
curl -L -o utils.py https://huggingface.co/EntropyDrop/Sking/resolve/main/SKING_DDJ_v101c/runtime/dense_uv_parser/utils.py
186 kB
| 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) | |