D_visualize_3d_scene_frame.py

"""Tutorial: read 3D bounding boxes and skeletons from a VCD label file and plot them with Open3D.

Bounding boxes are stored as [center(3), rotation(3), extent(3)] cuboids and skeletons as a flat
list of joint positions; this script loads a single frame and renders both annotation types in
an interactive Open3D viewer, alongside the matching RGB frame in a separate OpenCV window.

Press 1,2,3,4... to visualize different camera views.

CLI usage (run from the repository root)
---------
    uv run code_examples/plot_data/D_visualize_3d_scene_frame.py
    uv run code_examples/plot_data/D_visualize_3d_scene_frame.py \
        --vcd-path /path/to/vcd_3d_data.json --frame 100

Flags
-----
    --vcd-path         Path to a VCD JSON label file.
                       Default: bundled sample (input/vcd_3d_data.json).
    --skeleton-yaml    Path to the skeleton bone-connections YAML description.
                       Default: bundled sample (code_examples/plot_data/skeleton_description.yaml).
    --frame            int. Frame number to visualize (default: 402).
"""

import argparse
from pathlib import Path

import cv2
import numpy as np
import open3d as o3d
import yaml
from helpers_3d import (
    assign_skeleton_color,
    build_camera_views,
    get_bbox3d_corners_world_coordinates,
    read_bbox3d_labels,
    read_skeleton3d_labels,
)
from image_io import read_rgb_image
from vcd import core

RGB_WINDOW_NAME = "RGB view"

# Bundled sample dataset shipped with this repo.
DEFAULT_SKELETON_YAML = Path(__file__).resolve().parent / "skeleton_description.yaml"


def get_3d_bboxes_geometries(vcd, n_frame):
    frame_data = vcd.get_frame(n_frame)
    if not frame_data:
        return []

    geometries = []
    uids_list = frame_data["objects"].keys()
    bboxes_list = read_bbox3d_labels(vcd, uids_list, n_frame)
    for bbox3d in bboxes_list:
        # Color-code by object type so pedestrians/benches stand out from generic objects
        if "Bench" in bbox3d[1]["name"]:
            color = (255, 0, 0)
        elif "Pedestrian" in bbox3d[1]["name"]:
            color = (0, 255, 255)
        else:
            color = (0, 0, 255)
        o3d_cuboid, _ = get_bbox3d_corners_world_coordinates(bbox3d, color)
        geometries.append(o3d_cuboid)

    return geometries


def get_skeleton3d_o3d(skeleton_dict, bone_connections, color=[0.0, 1.0, 0.0]):
    def create_bone_mapping(bone_connections):
        return {bone: index for index, bone in enumerate(bone_connections.keys())}

    bone_order = create_bone_mapping(bone_connections)

    # Change y axis. O3D uses the y in the opposite direction
    bones_positions = np.array(skeleton_dict[0]) * (1, -1, 1)

    lines = o3d.geometry.LineSet()
    lines.points = o3d.utility.Vector3dVector(bones_positions)

    bone_lines = [
        [bone_order[bone_name], bone_order[connected_bone]]
        for bone_name, data in bone_connections.items()
        for connected_bone in data["connections"]
    ]
    lines.lines = o3d.utility.Vector2iVector(bone_lines)
    lines.colors = o3d.utility.Vector3dVector([color for _ in bone_lines])

    # Draw a sphere at the head joint (assumed to be the first point) for visibility
    head_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.1)
    head_sphere.translate(bones_positions[0])
    head_sphere.paint_uniform_color(color)

    return [lines, head_sphere]


def get_3d_skeleton_geometries(vcd, bone_connections, n_frame):
    frame_data = vcd.get_frame(n_frame)
    if not frame_data:
        return []

    geometries = []
    uids_list = list(frame_data["objects"].keys())
    pedestrian_uids = [
        uid for uid in uids_list if vcd.get_object(uid)["type"] == "Pedestrian"
    ]
    skeletons_list = read_skeleton3d_labels(vcd, pedestrian_uids, n_frame)

    for skeleton3d in skeletons_list:
        uid_name = skeleton3d[1]["name"]
        color = assign_skeleton_color(uid_name)
        geometries.extend(get_skeleton3d_o3d(skeleton3d, bone_connections, color))

    return geometries


def show_camera_rgb_image(vcd_path, camera_name, n_frame):
    """Load and display the RGB frame matching a camera view in an OpenCV window."""
    image = read_rgb_image(vcd_path, camera_name, n_frame)
    if image is None:
        return
    cv2.imshow(RGB_WINDOW_NAME, image)
    cv2.setWindowTitle(RGB_WINDOW_NAME, f"RGB view: {camera_name}")
    cv2.waitKey(1)


def pump_rgb_window_events(visualizer):
    """Animation callback that keeps the OpenCV RGB window responsive; no O3D redraw needed."""
    cv2.waitKey(1)
    return False


def switch_camera_view(vis, camera_views, camera_name, vcd_path, n_frame):
    """Point the O3D viewer at a camera's pose and show its matching RGB frame."""
    vis.get_view_control().convert_from_pinhole_camera_parameters(
        camera_views[camera_name], allow_arbitrary=True
    )
    print(f"Switched to camera view: {camera_name}")
    show_camera_rgb_image(vcd_path, camera_name, n_frame)


def register_camera_view_callbacks(vis, camera_views, vcd_path, n_frame):
    """Bind number keys (1-9) so pressing one jumps the viewer to that camera's pose."""

    def make_callback(camera_name):
        def _switch_camera(visualizer):
            switch_camera_view(visualizer, camera_views, camera_name, vcd_path, n_frame)
            return False

        return _switch_camera

    for index, camera_name in enumerate(sorted(camera_views)[:9]):
        vis.register_key_callback(ord(str(index + 1)), make_callback(camera_name))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Plot the 3D bounding boxes and skeletons of a single VCD frame with Open3D."
    )
    parser.add_argument(
        "--vcd-path",
        type=Path,
        required=True,
        help="Path to a 3D VCD JSON label file.",
    )
    parser.add_argument(
        "--skeleton-yaml",
        type=Path,
        default=DEFAULT_SKELETON_YAML,
        help="Path to the skeleton bone-connections YAML description.",
    )
    parser.add_argument(
        "--frame",
        type=int,
        default=102,
        help="Frame number to visualize (default: 102).",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    with open(args.skeleton_yaml) as f:
        bone_connections = yaml.safe_load(f)["bones"]

    vcd = core.OpenLABEL()
    vcd.load_from_file(str(args.vcd_path))

    geometries = get_3d_bboxes_geometries(vcd, args.frame)
    geometries += get_3d_skeleton_geometries(vcd, bone_connections, args.frame)

    camera_views = build_camera_views(vcd)

    vis = o3d.visualization.VisualizerWithKeyCallback()
    vis.create_window(window_name="VCD 3D scene (press 1-9 to switch camera view)")
    for geometry in geometries:
        vis.add_geometry(geometry)

    if camera_views:
        register_camera_view_callbacks(vis, camera_views, args.vcd_path, args.frame)
        # Keep the OpenCV RGB window responsive to redraws while the O3D loop runs.
        vis.register_animation_callback(pump_rgb_window_events)
        keys = ", ".join(
            f"{i + 1}={name}" for i, name in enumerate(sorted(camera_views)[:9])
        )
        print(f"Press a number key to jump to a camera view: {keys}")
        switch_camera_view(
            vis, camera_views, sorted(camera_views)[0], args.vcd_path, args.frame
        )

    vis.run()
    vis.destroy_window()
    cv2.destroyAllWindows()