E_project_3d_bboxes_to_frame.py

"""Tutorial: project 3D bounding boxes from a VCD label file onto camera frames.

For every requested camera and frame, this projects each object's 3D cuboid corners
into the camera's image plane using the calibration already stored in the VCD file,
draws the resulting wireframe box, and saves the annotated image. By default it
processes every camera and every frame found in the VCD file.

CLI usage (run from the repository root)
---------
    uv run code_examples/plot_data/E_project_3d_bboxes_to_frame.py
    uv run code_examples/plot_data/E_project_3d_bboxes_to_frame.py \
        --vcd-path /path/to/vcd_3d_data.json \
        --cameras cam1 cam2 \
        --start-frame 0 --end-frame 10 \
        --output-dir /path/to/output

Flags
-----
    --vcd-path         Path to a VCD JSON label file.
                       Default: bundled sample (input/vcd_3d_data.json).
    --cameras          One or more camera names to project onto, e.g. cam1 cam3.
                       Default: every camera found in the VCD.
    --start-frame      int. First frame to process, inclusive (default: first frame).
    --end-frame        int. Last frame to process, inclusive (default: last frame).
    --output-dir       Directory to save annotated images
                       (default: output_annotations in the project root).
"""

import argparse
from pathlib import Path

import cv2
from helpers_3d import (
    get_bbox3d_edge_corners,
    get_camera_poses,
    project_points3d_to_image,
    read_bbox3d_labels,
)
from image_io import read_rgb_image
from vcd import core

# Edges connecting the 8 cuboid corners: top face, bottom face, then vertical edges.
BBOX_EDGES = [
    (0, 1), (1, 2), (2, 3), (3, 0),
    (4, 5), (5, 6), (6, 7), (7, 4),
    (0, 4), (1, 5), (2, 6), (3, 7),
]  # fmt: skip

# Bundled sample dataset shipped with this repo.
DEFAULT_VCD_PATH = Path(__file__).resolve().parents[2] / "input" / "vcd_3d_data.json"
DEFAULT_OUTPUT_DIR = Path(__file__).resolve().parents[2] / "output_annotations"


def get_projected_bboxes(vcd, camera_name, n_frame):
    """Return the projected 2D corners of every 3D bounding box visible in a frame."""
    frame_data = vcd.get_frame(n_frame)
    if not frame_data:
        return []

    uids_list = frame_data["objects"].keys()
    bboxes_list = read_bbox3d_labels(vcd, uids_list, n_frame)

    projected_bboxes = []
    for bbox3d in bboxes_list:
        corners = get_bbox3d_edge_corners(bbox3d)
        projected_bboxes.append(project_points3d_to_image(vcd, corners, camera_name))
    return projected_bboxes


def plot_projected_bboxes(image, projected_bboxes, color=(0, 255, 0)):
    """Draw every projected bounding box wireframe onto an image."""
    for bbox in projected_bboxes:
        if not all(point[3] for point in bbox):
            continue  # skip boxes with corners outside the camera's field of view

        points = [(point[1], point[2]) for point in bbox]
        for start, end in BBOX_EDGES:
            cv2.line(image, points[start], points[end], color, 2)
    return image


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Project the 3D bounding boxes of a VCD file onto camera images."
    )
    parser.add_argument(
        "--vcd-path",
        type=Path,
        default=DEFAULT_VCD_PATH,
        help="Path to a VCD JSON label file (default: bundled sample dataset).",
    )
    parser.add_argument(
        "--cameras",
        nargs="+",
        default=None,
        help="Camera names to project onto, e.g. cam1 cam3 (default: every camera in the VCD).",
    )
    parser.add_argument(
        "--start-frame",
        type=int,
        default=None,
        help="First frame to process, inclusive (default: the VCD's first frame).",
    )
    parser.add_argument(
        "--end-frame",
        type=int,
        default=None,
        help="Last frame to process, inclusive (default: the VCD's last frame).",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help="Directory to save annotated images (default: 'output_annotations' in the project root).",
    )
    return parser.parse_args()


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

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

    camera_names = args.cameras or sorted(get_camera_poses(vcd))
    if not camera_names:
        raise SystemExit("No cameras found in the VCD file.")

    last_frame = vcd.get_frame_intervals().fis_dict[-1]["frame_end"]
    start_frame = 0 if args.start_frame is None else args.start_frame
    end_frame = last_frame if args.end_frame is None else args.end_frame

    output_root = args.output_dir / "bbox3D"

    for camera_name in camera_names:
        output_folder = output_root / camera_name
        output_folder.mkdir(parents=True, exist_ok=True)

        for n_frame in range(start_frame, end_frame + 1):
            image = read_rgb_image(args.vcd_path, camera_name, n_frame)
            if image is None:
                continue

            projected_bboxes = get_projected_bboxes(vcd, camera_name, n_frame)
            image = plot_projected_bboxes(image, projected_bboxes)

            output_path = output_folder / f"bbox3d_{n_frame:06d}.jpeg"
            cv2.imwrite(str(output_path), image)
            print(f"Image saved in: {output_path}")