draw_path.py

import os
from pathlib import Path

import cv2
import tqdm
import av
import numpy as np
from torchcodec.decoders import VideoDecoder

from visualize import draw_path
from orientation import rot_from_euler, rot_from_quat


def build_idxs(max_val: float, size: int):
    """Quadratically space out points https://github.com/commaai/openpilot/blob/7d3ad941bc4ba4c923af7a1d7b48544bfc0d3e13/selfdrive/common/modeldata.h#L14-L25"""
    return np.array([max_val * ((i / (size - 1)) ** 2) for i in range(size)])


def draw(input_dir, output_dir, py, input_prefix=""):
    INPUT_H, INPUT_W = 1080, 1920
    FOCAL = 790
    K = np.array([[FOCAL, 0, INPUT_W / 2.0], [0, FOCAL, INPUT_H / 2.0], [0, 0, 1]])
    D = np.array([-0.03054249, 0.00974967, -0.00679687, 0.00238449])
    OFFSET_MS = 100
    should_rotate = False

    if "HEL" in input_dir or input_dir in [
        "all_data/CHI/snow/20250215_075823/C10668__20250215_075823-20250215_080005",
    ]:
        OFFSET_MS = 100
        INPUT_H, INPUT_W = 1280, 720
        FOCAL = 430
        K = np.array([[FOCAL, 0, INPUT_W / 2.0], [0, FOCAL, INPUT_H / 2.0], [0, 0, 1]])
        D = np.array([-0.03902752, 0.05059457, -0.05272288, 0.01720146])
        should_rotate = True

    input_dir = Path(input_dir)
    output_rel_path = input_dir.relative_to(input_prefix)
    output_rel_dir = output_rel_path.parent
    basename = output_rel_path.name
    full_output_dir = Path(output_dir) / output_rel_dir
    os.makedirs(full_output_dir, exist_ok=True)

    video_fp = input_dir / "front.mp4"
    frame_times_s = np.load(input_dir / "frame_times.npy") / 1e9 - OFFSET_MS / 1e3
    frame_positions = np.load(input_dir / "frame_positions.npy")
    frame_orientations = np.load(input_dir / "frame_orientations.npy")
    frame_paths = np.load(input_dir / "frame_paths.npy")
    frame_states = np.load(input_dir / "frame_states.npy")

    output_path = os.path.join(full_output_dir, basename + ".mp4")

    decoder = VideoDecoder(video_fp, dimension_order="NHWC")
    plan_len_seconds = 3
    plan_len_points = 10
    t_anchors = build_idxs(plan_len_seconds, plan_len_points)
    NUM_PTS = plan_len_seconds * 20

    print(f"Writing to {output_path=}")
    output = av.open(output_path, mode="w")
    # stream = output.add_stream("h264_videotoolbox", rate=20)
    stream = output.add_stream("h264", rate=20)
    stream.width = INPUT_W
    stream.height = INPUT_H
    stream.pix_fmt = "yuv420p"

    for i, frame in enumerate(decoder[:-NUM_PTS]):
        # for i, frame in enumerate(decoder[:100]):

        if should_rotate:
            frame = np.rot90(frame, k=-1)

        odom_from_local = rot_from_quat(frame_orientations[i])
        local_from_odom = odom_from_local.T
        frame_positions_local = np.einsum("ij,kj->ki", local_from_odom, frame_positions[i : i + NUM_PTS] - frame_positions[i]).astype(np.float32)

        t_local = frame_times_s[i : i + NUM_PTS] - frame_times_s[i]
        x_interp = np.interp(t_anchors, t_local, frame_positions_local[:, 0])
        y_interp = np.interp(t_anchors, t_local, frame_positions_local[:, 1])

        future_poses = np.hstack(
            [
                x_interp.reshape(-1, 1),
                y_interp.reshape(-1, 1),
                np.zeros_like(y_interp.reshape(-1, 1)),
            ]
        )
        # future_poses = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0], [5, 0, 0], [10, 0, 0]])

        future_poses = np.hstack(
            [
                frame_paths[i, :, 0:1],
                frame_paths[i, :, 1:2],
                np.zeros_like(frame_paths[i, :, 1:2]),
            ]
        )

        # project path to image pts
        # img = np.array(frame)
        if not isinstance(frame, np.ndarray):
            img = frame.numpy()
        else:
            img = frame

        p, y = py
        R = rot_from_euler([np.radians(p), -np.radians(y), 0])
        img = draw_path(future_poses, img.copy(), K=K, R=R, D=D, fill_color=(66, 237, 66), line_color=None, width=0.25)

        img = draw_status_bar(img.copy(), mode="auto" if frame_states[i] == 2 else "manual")

        av_frame = av.VideoFrame.from_ndarray(img)
        packet = stream.encode(av_frame)
        if packet:
            output.mux(packet)

    # Flush
    packet = stream.encode(None)
    if packet:
        output.mux(packet)
    output.close()


def draw_status_bar(img, mode, pad=12):
    """
    Draw bottom-right status: green dot + 'auto' or red dot + 'manual'.
    img: BGR uint8 image (OpenCV)
    mode: "auto" or "manual"
    """
    text = "auto" if mode == "auto" else "manual"
    color = (0, 255, 0) if mode == "auto" else (255, 0, 0)  # BGR

    font = cv2.FONT_HERSHEY_SIMPLEX
    scale = 1
    thickness = 2

    (tw, th), baseline = cv2.getTextSize(text, font, scale, thickness)

    # Layout
    radius = 8
    gap = 10  # gap between dot and text
    box_w = radius * 2 + gap + tw + pad * 2
    box_h = max(th + baseline, radius * 2) + pad * 2

    h, w = img.shape[:2]
    x2, y2 = w - pad, h - pad
    x1, y1 = x2 - box_w, y2 - box_h

    # Background box (semi-transparent)
    overlay = img.copy()
    cv2.rectangle(overlay, (x1, y1), (x2, y2), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.45, img, 0.55, 0, img)

    # Dot
    dot_center = (x1 + pad + radius, y1 + box_h // 2)
    cv2.circle(img, dot_center, radius, color, -1, lineType=cv2.LINE_AA)

    # Text (vertically centered to dot)
    cy = dot_center[1]
    text_org = (dot_center[0] + radius + gap, cy + th // 2)
    cv2.putText(img, text, text_org, font, scale, color, thickness, cv2.LINE_AA)

    return img


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Draw path on video frames.")
    parser.add_argument("--input", type=str, help="Directory containing input data.")
    parser.add_argument("--output", type=str, help="Directory to save output video.")
    parser.add_argument("--py", type=float, nargs=2, default=(0, 0), help="Pitch and yaw angles in degrees.")

    args = parser.parse_args()

    if not os.path.exists(args.input):
        raise FileNotFoundError(f"Input directory {args.input} does not exist.")

    # walk through the input directory
    dirs_to_process = []
    for root, dirs, files in os.walk(args.input):
        if "HEL" in root:
            continue
        if "C10668__20250215_075823-20250215_080005" in root:
            continue
        done = [
            "C11093__20250425_104914-20250425_104955",
            "C10527__20250424_150906-20250424_150959",
            "C11093__20250428_162730-20250428_162741",
            "C10815__20250427_172858-20250427_172917",
            "C11016__20250404_151053-20250404_151201",
            "C10498__20250424_165142-20250424_165221",
            "C11093__20250425_105334-20250425_105413",
            "C10668__20250215_075823-20250215_080005",
        ]
        already = False
        for d in done:
            if d in root:
                already = True
                break
        if already:
            continue

        for fn in files:
            if fn == "front.mp4":
                dirs_to_process.append(root)
                break

    for input_dir in tqdm.tqdm(dirs_to_process):
        # print(f"Processing {input_dir}...")
        draw(input_dir, args.output, args.py, input_prefix=args.input)
    # draw(args.input, args.output, args.py)