mapvggt / mapgs /hdmap /rasterize_map.py

Upload folder using huggingface_hub

b2efbe4 verified 4 months ago

4.69 kB

"""Render map geometry into camera views.
* :func:`rasterize_map_depth` casts each pixel ray against the ground height
  field (a fixed-point intersection) to get per-view ground depth ``D^map`` and
  the ground/lane mask ``Omega^g``. ``D^map`` is a supervision *target*
  (constant w.r.t. the Gaussians), so it is computed without gradients and
  fully vectorized over pixels — no triangle z-buffer needed.
* :func:`project_polylines` projects lane / boundary vertices into each view
  (used as the map side of the lane-chamfer term, §2.6 item 3, and for lane
  mIoU, §4.4).
* :func:`render_lane_mask` splats projected lane points into a soft raster.
"""

from __future__ import annotations

from typing import List

import torch

from mapgs.geometry.cameras import camera_centers
from mapgs.hdmap.ground_field import GroundField


@torch.no_grad()
def rasterize_map_depth(
    ground: GroundField,
    K: torch.Tensor,          # [V, 3, 3]
    cam2world: torch.Tensor,  # [V, 4, 4]
    H: int,
    W: int,
    iters: int = 16,
    near: float = 0.5,
    far: float = 200.0,
    min_descent: float = 1e-3,
    tol: float = 0.25,
):
    """Ray/height-field intersection -> ``depth [V,H,W]``, ``mask [V,H,W]`` (bool)."""
    device = K.device
    V = K.shape[0]
    dtype = K.dtype

    vv, uu = torch.meshgrid(
        torch.arange(H, device=device, dtype=dtype) + 0.5,
        torch.arange(W, device=device, dtype=dtype) + 0.5,
        indexing="ij",
    )
    ones = torch.ones_like(uu)
    pix = torch.stack([uu, vv, ones], dim=-1)  # [H, W, 3]
    Kinv = torch.inverse(K)                     # [V, 3, 3]
    r_cam = torch.einsum("vij,hwj->vhwi", Kinv, pix)  # camera ray, z-comp == 1
    R_c2w = cam2world[:, :3, :3]
    m = torch.einsum("vij,vhwj->vhwi", R_c2w, r_cam)  # world delta per unit cam-depth
    o = camera_centers(cam2world)[:, None, None, :]    # [V,1,1,3]

    mz = m[..., 2]
    descending = mz < -min_descent
    # initial guess from flat ground under the camera
    h0, _ = ground.height_at(o[..., :2].expand(V, H, W, 2))
    Z = (h0 - o[..., 2]) / mz.clamp(max=-min_descent)
    Z = Z.clamp(near, far)
    for _ in range(iters):
        xy = o[..., :2] + Z.unsqueeze(-1) * m[..., :2]
        h, _ = ground.height_at(xy)
        Z = ((h - o[..., 2]) / mz.clamp(max=-min_descent)).clamp(near, far)

    xy = o[..., :2] + Z.unsqueeze(-1) * m[..., :2]
    h, valid_xy = ground.height_at(xy)
    residual = (o[..., 2] + Z * mz - h).abs()
    mask = descending & valid_xy & (residual < tol) & (Z > near) & (Z < far)
    depth = torch.where(mask, Z, torch.zeros_like(Z))
    return depth, mask


@torch.no_grad()
def project_polylines(
    polylines,                # [P, 3] tensor or list of [Li, 3]
    K: torch.Tensor,          # [V, 3, 3]
    cam2world: torch.Tensor,  # [V, 4, 4]
    H: int,
    W: int,
) -> List[torch.Tensor]:
    """Project map polyline vertices into each view. Returns list (len V) of ``[Mi, 2]`` uv."""
    from mapgs.geometry.cameras import project_points
    from mapgs.geometry.transforms import se3_inverse

    if isinstance(polylines, (list, tuple)):
        pts = torch.cat([p for p in polylines if p.numel() > 0], 0) if len(polylines) else torch.zeros(0, 3)
    else:
        pts = polylines
    device = K.device
    pts = pts.to(device)
    V = K.shape[0]
    out: List[torch.Tensor] = []
    if pts.numel() == 0:
        return [torch.zeros(0, 2, device=device) for _ in range(V)]
    w2c = se3_inverse(cam2world)
    for v in range(V):
        uv, z = project_points(pts[None], K[v : v + 1], w2c[v : v + 1])
        uv = uv[0]
        z = z[0]
        inb = (z > 0.1) & (uv[:, 0] >= 0) & (uv[:, 0] < W) & (uv[:, 1] >= 0) & (uv[:, 1] < H)
        out.append(uv[inb])
    return out


def render_lane_mask(uv: torch.Tensor, H: int, W: int, radius: float = 2.0) -> torch.Tensor:
    """Splat projected lane points ``[M, 2]`` into a soft mask ``[H, W]`` in [0,1]."""
    device = uv.device
    if uv.numel() == 0:
        return torch.zeros(H, W, device=device)
    vv, uu = torch.meshgrid(
        torch.arange(H, device=device, dtype=torch.float32),
        torch.arange(W, device=device, dtype=torch.float32),
        indexing="ij",
    )
    # distance from each pixel to nearest lane point (chunked to bound memory)
    grid = torch.stack([uu, vv], dim=-1).reshape(-1, 2)  # [HW, 2]
    min_d2 = torch.full((grid.shape[0],), 1e9, device=device)
    for i in range(0, uv.shape[0], 4096):
        chunk = uv[i : i + 4096]
        d2 = (grid[:, None, :] - chunk[None, :, :]).pow(2).sum(-1).min(dim=1).values
        min_d2 = torch.minimum(min_d2, d2)
    mask = torch.exp(-min_d2 / (2 * radius ** 2))
    return mask.reshape(H, W)

	"""Render map geometry into camera views.

	* :func:`rasterize_map_depth` casts each pixel ray against the ground height
	field (a fixed-point intersection) to get per-view ground depth ``D^map`` and
	the ground/lane mask ``Omega^g``. ``D^map`` is a supervision target
	(constant w.r.t. the Gaussians), so it is computed without gradients and
	fully vectorized over pixels — no triangle z-buffer needed.
	* :func:`project_polylines` projects lane / boundary vertices into each view
	(used as the map side of the lane-chamfer term, §2.6 item 3, and for lane
	mIoU, §4.4).
	* :func:`render_lane_mask` splats projected lane points into a soft raster.
	"""

	from __future__ import annotations

	from typing import List

	import torch

	from mapgs.geometry.cameras import camera_centers
	from mapgs.hdmap.ground_field import GroundField


	@torch.no_grad()
	def rasterize_map_depth(
	ground: GroundField,
	K: torch.Tensor, # [V, 3, 3]
	cam2world: torch.Tensor, # [V, 4, 4]
	H: int,
	W: int,
	iters: int = 16,
	near: float = 0.5,
	far: float = 200.0,
	min_descent: float = 1e-3,
	tol: float = 0.25,
	):
	"""Ray/height-field intersection -> ``depth [V,H,W]``, ``mask [V,H,W]`` (bool)."""
	device = K.device
	V = K.shape[0]
	dtype = K.dtype

	vv, uu = torch.meshgrid(
	torch.arange(H, device=device, dtype=dtype) + 0.5,
	torch.arange(W, device=device, dtype=dtype) + 0.5,
	indexing="ij",
	)
	ones = torch.ones_like(uu)
	pix = torch.stack([uu, vv, ones], dim=-1) # [H, W, 3]
	Kinv = torch.inverse(K) # [V, 3, 3]
	r_cam = torch.einsum("vij,hwj->vhwi", Kinv, pix) # camera ray, z-comp == 1
	R_c2w = cam2world[:, :3, :3]
	m = torch.einsum("vij,vhwj->vhwi", R_c2w, r_cam) # world delta per unit cam-depth
	o = camera_centers(cam2world)[:, None, None, :] # [V,1,1,3]

	mz = m[..., 2]
	descending = mz < -min_descent
	# initial guess from flat ground under the camera
	h0, _ = ground.height_at(o[..., :2].expand(V, H, W, 2))
	Z = (h0 - o[..., 2]) / mz.clamp(max=-min_descent)
	Z = Z.clamp(near, far)
	for _ in range(iters):
	xy = o[..., :2] + Z.unsqueeze(-1) * m[..., :2]
	h, _ = ground.height_at(xy)
	Z = ((h - o[..., 2]) / mz.clamp(max=-min_descent)).clamp(near, far)

	xy = o[..., :2] + Z.unsqueeze(-1) * m[..., :2]
	h, valid_xy = ground.height_at(xy)
	residual = (o[..., 2] + Z * mz - h).abs()
	mask = descending & valid_xy & (residual < tol) & (Z > near) & (Z < far)
	depth = torch.where(mask, Z, torch.zeros_like(Z))
	return depth, mask


	@torch.no_grad()
	def project_polylines(
	polylines, # [P, 3] tensor or list of [Li, 3]
	K: torch.Tensor, # [V, 3, 3]
	cam2world: torch.Tensor, # [V, 4, 4]
	H: int,
	W: int,
	) -> List[torch.Tensor]:
	"""Project map polyline vertices into each view. Returns list (len V) of ``[Mi, 2]`` uv."""
	from mapgs.geometry.cameras import project_points
	from mapgs.geometry.transforms import se3_inverse

	if isinstance(polylines, (list, tuple)):
	pts = torch.cat([p for p in polylines if p.numel() > 0], 0) if len(polylines) else torch.zeros(0, 3)
	else:
	pts = polylines
	device = K.device
	pts = pts.to(device)
	V = K.shape[0]
	out: List[torch.Tensor] = []
	if pts.numel() == 0:
	return [torch.zeros(0, 2, device=device) for _ in range(V)]
	w2c = se3_inverse(cam2world)
	for v in range(V):
	uv, z = project_points(pts[None], K[v : v + 1], w2c[v : v + 1])
	uv = uv[0]
	z = z[0]
	inb = (z > 0.1) & (uv[:, 0] >= 0) & (uv[:, 0] < W) & (uv[:, 1] >= 0) & (uv[:, 1] < H)
	out.append(uv[inb])
	return out


	def render_lane_mask(uv: torch.Tensor, H: int, W: int, radius: float = 2.0) -> torch.Tensor:
	"""Splat projected lane points ``[M, 2]`` into a soft mask ``[H, W]`` in [0,1]."""
	device = uv.device
	if uv.numel() == 0:
	return torch.zeros(H, W, device=device)
	vv, uu = torch.meshgrid(
	torch.arange(H, device=device, dtype=torch.float32),
	torch.arange(W, device=device, dtype=torch.float32),
	indexing="ij",
	)
	# distance from each pixel to nearest lane point (chunked to bound memory)
	grid = torch.stack([uu, vv], dim=-1).reshape(-1, 2) # [HW, 2]
	min_d2 = torch.full((grid.shape[0],), 1e9, device=device)
	for i in range(0, uv.shape[0], 4096):
	chunk = uv[i : i + 4096]
	d2 = (grid[:, None, :] - chunk[None, :, :]).pow(2).sum(-1).min(dim=1).values
	min_d2 = torch.minimum(min_d2, d2)
	mask = torch.exp(-min_d2 / (2 * radius ** 2))
	return mask.reshape(H, W)