Function render_background_scene_flow

paper/render_point_cloud.py:170–230

(
    device: torch.device,
    birds: LoadedScene,
    proj: Callable,
    background_mask: Bool[Tensor, "height width"],
)

def render_background_scene_flow(
    device: torch.device,
    birds: LoadedScene,
    proj: Callable,
    background_mask: Bool[Tensor, "height width"],
) -> None:
    # Render the background points and lines.
    mask = background_mask & (~birds.highlight_mask)
    xyz_background_original = birds.xyz_camera_space[0][mask]
    num_points = xyz_background_original.shape[0]
    xyz_background_other = einsum(
        birds.exports.extrinsics[0, SECOND_FRAME].inverse(),
        homogenize_points(xyz_background_original),
        "i j, p j -> p i",
    )[..., :3]
    xyz_background = torch.cat((xyz_background_original, xyz_background_other), dim=0)
    is_other_background = (
        torch.zeros((num_points,), dtype=torch.bool, device=device),
        torch.ones((num_points,), dtype=torch.bool, device=device),
    )
    is_other_background = torch.cat(is_other_background, dim=0)
    xy_background, ordering_background = proj(xyz_background)
    lines_background = rearrange(xy_background, "(e l) xy -> l e xy", e=2)
    xy_background = xy_background[ordering_background]
    is_other_background = is_other_background[ordering_background]

    # Export to SVG.
    fig = svg.SVG(
        width=CANVAS_SIZE,
        height=CANVAS_SIZE,
        elements=[],
        viewBox=svg.ViewBoxSpec(0, 0, CANVAS_SIZE, CANVAS_SIZE),
    )

    # Draw the connecting lines.
    for start, end in lines_background.tolist():
        line = svg.Line(
            x1=start[0],
            y1=start[1],
            x2=end[0],
            y2=end[1],
            stroke="#cccccc",
            stroke_width=4,
            stroke_linecap="round",
        )
        fig.elements.append(line)

    # Draw the endpoints.
    for (x, y), is_background in zip(xy_background.tolist(), is_other_background):
        line = svg.Line(
            x1=x,
            y1=y,
            x2=x,
            y2=y,
            stroke=DISTINCT_COLORS[1 if is_background else 0],
            stroke_width=12,
            stroke_linecap="round",
        )