EntropyDrop's picture
Release SKING_DDJ_v101c portable inference bundle
76d7a6e verified
Raw History Blame Contribute Delete
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)