mapgs/data/scene.py

"""Procedural driving-scene generator (for synthetic data, §2.1).

The ground-truth world is built as a set of 3D Gaussians (road surface + lane
stripes + curbs + roadside structures + moving vehicles). Rendering it with the
same rasterizer yields multi-view-consistent GT images and GT depth, and the HD
map (ground height field + lane / boundary polylines) is read off the same
scene (optionally with injected noise for the map-robustness study, §4.6).

World frame: x = lateral (right), y = forward, z = up.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional

import math
import numpy as np
import torch

from mapgs.geometry.cameras import look_at_pose
from mapgs.hdmap.ground_field import grid_field_from_points, GridGroundField
from mapgs.render.gaussians import Gaussians, GROUP_DYNAMIC


@dataclass
class MapNoise:
    height_std: float = 0.0     # m, additive ground-height noise
    lane_offset_std: float = 0.0  # m, lateral lane shift
    drop_prob: float = 0.0      # fraction of polyline segments dropped


@dataclass
class ProceduralScene:
    static: Gaussians
    dyn_canonical: List[Gaussians]
    box_centers: torch.Tensor   # [I, F, 3]
    box_rots: torch.Tensor      # [I, F, 3, 3]
    box_size: torch.Tensor      # [I, 3]
    canon_idx: torch.Tensor     # [I]
    ground: GridGroundField
    lanes: List[torch.Tensor]
    boundaries: List[torch.Tensor]
    K: torch.Tensor             # [V, 3, 3]
    cam2world: torch.Tensor     # [F, V, 4, 4]
    F: int
    V: int
    fps: int


def _box_gaussians(g, center, size, color, n, opacity=0.9, jitter_scale=0.06):
    pts = (g_rand(g, (n, 3)) - 0.5) * torch.tensor(size)
    means = pts + torch.tensor(center)
    scales = torch.full((n, 3), jitter_scale)
    quats = torch.zeros(n, 4); quats[:, 0] = 1
    opac = torch.full((n,), opacity)
    cols = torch.tensor(color)[None].repeat(n, 1) + 0.03 * g_rand(g, (n, 3))
    return means.float(), scales.float(), quats.float(), opac.float(), cols.clamp(0, 1).float()


def g_rand(g, shape):
    return torch.rand(shape, generator=g)


def _centerline_x(y, curve_amp, curve_freq):
    return curve_amp * torch.sin(y * curve_freq)


def generate_scene(
    seed: int = 0,
    F: int = 20,
    fps: int = 10,
    H: int = 256,
    W: int = 448,
    n_dynamic: int = 2,
    fov_deg: float = 70.0,
    curve_amp: float = 2.0,
    curve_freq: float = 0.02,
    length: float = 70.0,
    map_noise: Optional[MapNoise] = None,
) -> ProceduralScene:
    g = torch.Generator().manual_seed(seed)
    map_noise = map_noise or MapNoise()
    V = 3
    dt = 1.0 / fps

    # ---- ground surface samples + height field ----
    ys = torch.linspace(-5, length, 200)
    xs = torch.linspace(-12, 12, 60)
    GX, GY = torch.meshgrid(xs, ys, indexing="xy")
    slope = 0.01
    base_z = slope * GY + 0.3 * torch.sin(GY * 0.05)
    gpts = torch.stack([GX.reshape(-1), GY.reshape(-1), base_z.reshape(-1)], -1)

    def ground_z(x, y):
        return slope * y + 0.3 * torch.sin(y * 0.05)

    # map ground field (optionally noisy)
    map_pts = gpts.clone().numpy()
    if map_noise.height_std > 0:
        map_pts[:, 2] += np.random.RandomState(seed).randn(map_pts.shape[0]) * map_noise.height_std
    ground = grid_field_from_points(map_pts, spacing=0.7)

    # ---- lanes & boundaries (polylines along the curved centerline) ----
    yline = torch.linspace(0, length, 120)
    cx = _centerline_x(yline, curve_amp, curve_freq)
    lane_off = [-3.5, 0.0, 3.5]
    bnd_off = [-7.0, 7.0]
    lanes, boundaries = [], []
    rs = np.random.RandomState(seed + 1)
    for off in lane_off:
        loff = off + (rs.randn() * map_noise.lane_offset_std)
        pl = torch.stack([cx + loff, yline, ground_z(cx + loff, yline)], -1)
        if map_noise.drop_prob > 0:  # drop a contiguous chunk
            keep = torch.rand(pl.shape[0], generator=g) > map_noise.drop_prob
            pl = pl[keep] if keep.any() else pl
        lanes.append(pl.float())
    for off in bnd_off:
        pl = torch.stack([cx + off, yline, ground_z(cx + off, yline)], -1)
        boundaries.append(pl.float())

    # ---- static GT gaussians: road + stripes + curbs + buildings/poles ----
    parts = []
    # road surface gaussians
    ny = 140
    yy = torch.linspace(0, length, ny)
    cxx = _centerline_x(yy, curve_amp, curve_freq)
    for lane_w in torch.linspace(-6, 6, 25):
        xw = cxx + lane_w
        means = torch.stack([xw, yy, ground_z(xw, yy)], -1)
        scales = torch.full((ny, 3), 0.18); scales[:, 2] = 0.02
        quats = torch.zeros(ny, 4); quats[:, 0] = 1
        opac = torch.full((ny,), 0.95)
        gray = 0.32 + 0.04 * g_rand(g, (ny, 1))
        cols = gray.repeat(1, 3)
        parts.append((means, scales, quats, opac, cols))
    # lane stripes (dashed white)
    for off in lane_off:
        xw = cxx + off
        dash = (torch.arange(ny) % 6 < 3)
        means = torch.stack([xw, yy, ground_z(xw, yy) + 0.01], -1)[dash]
        m = dash.sum()
        scales = torch.full((m, 3), 0.10); scales[:, 2] = 0.02
        quats = torch.zeros(m, 4); quats[:, 0] = 1
        opac = torch.full((m,), 0.98)
        cols = torch.full((m, 3), 0.9)
        parts.append((means, scales, quats, opac, cols))
    # roadside structures (buildings/poles) -> vertical content for L_vert
    n_build = 10
    for i in range(n_build):
        side = 1 if g_rand(g, (1,)).item() > 0.5 else -1
        y0 = float(g_rand(g, (1,)).item()) * length
        x0 = side * (8 + 5 * float(g_rand(g, (1,)).item()))
        h = 2 + 6 * float(g_rand(g, (1,)).item())
        col = (0.3 + 0.5 * g_rand(g, (3,))).tolist()
        parts.append(_box_gaussians(g, [x0 + _centerline_x(torch.tensor(y0), curve_amp, curve_freq).item(), y0, ground_z(torch.tensor(x0), torch.tensor(y0)).item() + h / 2],
                                    [3.0, 3.0, h], col, n=200))

    static = Gaussians(
        means=torch.cat([p[0] for p in parts]),
        scales=torch.cat([p[1] for p in parts]),
        quats=torch.cat([p[2] for p in parts]),
        opacities=torch.cat([p[3] for p in parts]),
        colors=torch.cat([p[4] for p in parts]),
    )

    # ---- dynamic vehicles ----
    dyn_canonical: List[Gaussians] = []
    box_centers = torch.zeros(n_dynamic, F, 3)
    box_rots = torch.eye(3).view(1, 1, 3, 3).repeat(n_dynamic, F, 1, 1)
    box_size = torch.zeros(n_dynamic, 3)
    canon_idx = torch.zeros(n_dynamic, dtype=torch.long)
    for i in range(n_dynamic):
        lane_x = lane_off[i % len(lane_off)]
        y0 = 10 + 15 * float(g_rand(g, (1,)).item())
        speed = 3 + 4 * float(g_rand(g, (1,)).item())  # m/s
        size = [4.2, 1.9, 1.5]
        box_size[i] = torch.tensor(size)
        col = (0.2 + 0.6 * g_rand(g, (3,))).tolist()
        for f in range(F):
            y = y0 + speed * f * dt
            x = _centerline_x(torch.tensor(y), curve_amp, curve_freq).item() + lane_x
            z = ground_z(torch.tensor(x), torch.tensor(y)).item() + size[2] / 2
            box_centers[i, f] = torch.tensor([x, y, z])
        # canonical gaussians placed at frame-0 box pose
        c0 = box_centers[i, 0].tolist()
        dyn_canonical.append(Gaussians(*_box_gaussians(g, c0, size, col, n=150)))

    # ---- cameras: ego along lane 0, 3-cam front rig ----
    fx = 0.5 * W / math.tan(math.radians(fov_deg) / 2)
    K = torch.tensor([[fx, 0, W / 2], [0, fx, H / 2], [0, 0, 1]]).float()[None].repeat(V, 1, 1)
    cam_yaws = [math.radians(22), 0.0, math.radians(-22)]  # left, center, right
    cam_lat = [-0.4, 0.0, 0.4]
    cam2world = torch.zeros(F, V, 4, 4)
    ego_speed = 6.0
    for f in range(F):
        y_ego = 2 + ego_speed * f * dt
        x_ego = _centerline_x(torch.tensor(y_ego), curve_amp, curve_freq).item()
        z_ego = ground_z(torch.tensor(x_ego), torch.tensor(y_ego)).item() + 1.5
        fwd = torch.tensor([_centerline_x(torch.tensor(y_ego + 1), curve_amp, curve_freq).item() - x_ego, 1.0, 0.0])
        fwd = fwd / fwd.norm()
        heading = math.atan2(fwd[0].item(), fwd[1].item())
        for v in range(V):
            yaw = heading + cam_yaws[v]
            eye = torch.tensor([x_ego + cam_lat[v] * math.cos(heading), y_ego, z_ego])
            tgt = eye + torch.tensor([math.sin(yaw), math.cos(yaw), -0.05]) * 10.0
            cam2world[f, v] = look_at_pose(eye, tgt)

    return ProceduralScene(
        static=static, dyn_canonical=dyn_canonical, box_centers=box_centers, box_rots=box_rots,
        box_size=box_size, canon_idx=canon_idx, ground=ground, lanes=lanes, boundaries=boundaries,
        K=K, cam2world=cam2world, F=F, V=V, fps=fps,
    )