File size: 5,440 Bytes
e6f80f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Optional scene regularizers, distinct from matched prediction-vs-GT GIoU.
Collision is normalized upright OBB intersection volume. Declared support pairs
are exempt only from volumetric overlap regularization; runtime verifies actual
contact. Boundary uses a convex input polygon's halfplanes; concave or unknown
boundaries fail explicitly when this training term is requested.
"""
from __future__ import annotations
import torch
from fastfill.v2.boxes import intersection_area
from fastfill.v2.geometry import bev_corners
from fastfill.v2.validation import effective_support_requests
def _convex_polygon(vertices):
from shapely.geometry import Polygon
polygon = Polygon(vertices)
if not polygon.is_valid or polygon.area <= 0 or abs(polygon.convex_hull.area - polygon.area) > 1e-8:
raise ValueError("boundary training regularizer requires a known convex polygon")
area = sum(a[0] * b[1] - a[1] * b[0] for a, b in zip(vertices, vertices[1:] + vertices[:1]))
return vertices if area > 0 else list(reversed(vertices))
def _collision(p, s, y, q, t, z):
original = p
# Compute all volume factors in float64 before multiplying/dividing. A
# float32 box side can be finite while its volume overflows or underflows.
# Metal has no float64 kernels; CPU transfers remain part of autograd.
values = (p, s, y, q, t, z)
p, s, y, q, t, z = (value.cpu().double() if original.device.type == "mps" else value.double()
for value in values)
offset = p.detach()
p, q = p - offset, q - offset
zero = sum((value * 0).sum() for value in (p, s, y, q, t, z))
height = (torch.minimum(p[2] + s[2], q[2] + t[2]) - torch.maximum(p[2], q[2])).clamp_min(0)
if not bool(height.detach() > 0):
# Multiply before reducing: finite large coordinates can overflow a
# sum even though the correct disjoint-box contribution is exactly zero.
result = zero
else:
volume = intersection_area(p, s, y, q, t, z) * height
denominator = (s.prod() + t.prod()) / 2
if not bool(torch.isfinite(denominator) & (denominator > 0)):
raise ValueError("collision normalization requires finite positive float64 box volume")
result = volume / denominator + zero
# Separate the CPU dtype conversion and transfer: a combined `.to` can
# attempt a forbidden float64 conversion on Metal during backward.
return result.to(dtype=original.dtype).to(original.device) if original.device.type == "mps" else result
def _boundary(p, s, y, room):
if room.get("boundary_known", True) is not True:
raise ValueError("boundary regularizer cannot use unknown/incomplete room boundaries")
poly = p.new_tensor(_convex_polygon(room["floor_polygon_xy_m"]))
edge = poly.roll(-1, 0) - poly
offsets = bev_corners(p, s, y)[:, None] - poly[None]
cross = edge[None, :, 0] * offsets[..., 1] - edge[None, :, 1] * offsets[..., 0]
distance = cross / edge.norm(dim=-1).clamp_min(torch.finfo(p.dtype).tiny)
return torch.relu(-distance).mean()
def regularizer_specs(condition, config):
"""One immutable enumeration for loss terms and eligibility counting."""
objects, room = effective_support_requests(condition), condition["room"]
if config.boundary and objects:
if room.get("boundary_known", True) is not True:
raise ValueError("boundary regularizer cannot use unknown/incomplete room boundaries")
_convex_polygon(room["floor_polygon_xy_m"])
pairs, fixed_pairs = [], []
if config.collision:
for i, obj in enumerate(objects):
pairs.extend((i, j) for j in range(i)
if obj.get("support_parent") != objects[j]["id"]
and objects[j].get("support_parent") != obj["id"])
fixed_pairs.extend((i, fixed) for fixed in room.get("fixed_objects", [])
if obj.get("support_parent") != fixed["id"])
return objects, room, tuple(pairs), tuple(fixed_pairs)
def scene_regularizers(positions, sizes, yaws, batch, config, global_counts=True):
from fastfill.v2.losses import _mean
zero = (positions[batch["slot_mask"]] * 0).sum()
collision, boundary = [], []
if "conditions" not in batch:
raise ValueError("scene regularizers require original conditions")
for b, condition in enumerate(batch["conditions"]):
objects, room, pairs, fixed_pairs = regularizer_specs(condition, config)
if config.boundary:
boundary.extend(_boundary(positions[b, i], sizes[b, i], yaws[b, i], room) for i in range(len(objects)))
collision.extend(_collision(positions[b, i], sizes[b, i], yaws[b, i],
positions[b, j], sizes[b, j], yaws[b, j]) for i, j in pairs)
for i, fixed in fixed_pairs:
q, t, z = [positions.new_tensor(fixed[k]) for k in ("bottom_center_m", "size_local_m", "yaw_rad")]
collision.append(_collision(positions[b, i], sizes[b, i], yaws[b, i], q, t, z))
terms = {"collision": collision, "boundary": boundary}
return {**{key: _mean(sum(values, zero), len(values), global_counts)[0] for key, values in terms.items()},
"sums": {key: sum((v.detach().float() for v in values), torch.zeros((), device=zero.device)) for key, values in terms.items()},
"counts": {key: len(values) for key, values in terms.items()}}
|