BEV 数据处理管线

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
BEV 训练数据集 —— 读取 2.1 采集的数据，产出模型可用的张量

关键设计:
  1. lidar2img 矩阵预计算（自车系 3D → 像素，BEVFormer 直接用）
  2. 变长真值用自定义 collate_fn（每帧框数不同，不能 stack）
  3. 数据增强时图像/点云/真值必须联动变换
"""

import json
import math
import os

import numpy as np
import cv2
import torch
from torch.utils.data import Dataset


# ImageNet 统计量（骨干网络预训练时用的，必须一致）
IMG_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMG_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

# CARLA(UE4) → 标准相机 的轴重排（1.4 节推导过）
UE4_TO_CAM = np.array([
    [0,  1,  0],
    [0,  0, -1],
    [1,  0,  0],
], dtype=np.float64)

CATEGORIES = ["car", "pedestrian", "bicycle", "motorcycle"]
CAT2ID = {c: i for i, c in enumerate(CATEGORIES)}


def build_lidar2img(K, cam_translation, cam_yaw_deg):
    """
    构造「自车系 3D 点 → 图像像素」的 4×4 变换矩阵
    
    变换链（1.4 节讲过）:
        自车系 --相机外参逆--> CARLA相机系 --轴重排--> 标准相机系 --K--> 像素
    
    返回 4×4 矩阵 M，用法:
        p_img = M @ [x, y, z, 1]
        u, v = p_img[0]/p_img[2], p_img[1]/p_img[2]
        depth = p_img[2]
    """
    # --- 相机在自车系的位姿（外参）---
    yaw = math.radians(cam_yaw_deg)
    c, s = math.cos(yaw), math.sin(yaw)
    R_cam2ego = np.array([[c, -s, 0],
                          [s,  c, 0],
                          [0,  0, 1]], dtype=np.float64)
    t_cam2ego = np.array(cam_translation, dtype=np.float64)

    # --- 求逆：自车系 → 相机系 ---
    #   旋转矩阵正交，逆 = 转置
    R_ego2cam = R_cam2ego.T
    t_ego2cam = -R_ego2cam @ t_cam2ego

    # --- 组合：轴重排 · 外参逆 ---
    R_total = UE4_TO_CAM @ R_ego2cam
    t_total = UE4_TO_CAM @ t_ego2cam

    # --- 左乘内参 ---
    K = np.array(K, dtype=np.float64)
    M = np.eye(4, dtype=np.float64)
    M[:3, :3] = K @ R_total
    M[:3, 3] = K @ t_total
    return M


def scale_intrinsic(K, scale_x, scale_y):
    """
    图像缩放后，内参必须同步缩放 —— 这是最容易漏的一步
    
    忘了这步的后果：投影全部错位，且不报错。
    """
    K = np.array(K, dtype=np.float64).copy()
    K[0, :] *= scale_x     # fx, cx
    K[1, :] *= scale_y     # fy, cy
    return K


class BEVDataset(Dataset):

    def __init__(self, root, split="train", img_size=(800, 450),
                 pc_range=51.2, augment=False, val_ratio=0.2):
        """
        root      : 2.1 产出的数据集根目录
        img_size  : (W, H) 送入模型的尺寸
        augment   : 训练开启，验证关闭
        """
        self.root = root
        self.img_w, self.img_h = img_size
        self.pc_range = pc_range
        self.augment = augment

        info = json.load(open(os.path.join(root, "dataset_info.json")))
        self.cam_names = list(info["camera_rig"].keys())
        self.src_w, self.src_h = info["image_size"]

        # --- 建立样本索引 ---
        # 按场景切分 train/val，避免同场景相邻帧泄漏到验证集
        scenes = [s["name"] for s in info["scenes"]]
        n_val = max(1, int(len(scenes) * val_ratio))
        if split == "train":
            use_scenes = scenes[:-n_val] if len(scenes) > n_val else scenes
        else:
            use_scenes = scenes[-n_val:]

        self.samples = []
        self.scene_calib = {}
        for sc in use_scenes:
            sdir = os.path.join(root, sc)
            fpath = os.path.join(sdir, "frames.json")
            if not os.path.exists(fpath):
                continue
            data = json.load(open(fpath))
            self.scene_calib[sc] = json.load(
                open(os.path.join(sdir, "calib.json")))
            for fr in data["frames"]:
                self.samples.append((sc, fr))

        # --- 预计算 lidar2img（标定场景内固定，不必每次算）---
        self.l2i_cache = {}
        for sc, calib in self.scene_calib.items():
            mats = []
            sx = self.img_w / self.src_w
            sy = self.img_h / self.src_h
            for name in self.cam_names:
                c = calib["cameras"][name]
                K = scale_intrinsic(c["intrinsic"], sx, sy)
                mats.append(build_lidar2img(
                    K, c["translation"], c["rotation_yaw_deg"]))
            self.l2i_cache[sc] = np.stack(mats)      # (6, 4, 4)

        print(f"[{split}] {len(use_scenes)} 场景, {len(self.samples)} 帧")

    def __len__(self):
        return len(self.samples)

    # ---------------- 核心 ----------------
    def __getitem__(self, idx):
        scene, frame = self.samples[idx]
        sdir = os.path.join(self.root, scene)
        fid = frame["index"]

        # ---- 1) 读 6 路图像 ----
        imgs = []
        for name in self.cam_names:
            path = os.path.join(sdir, name, f"{fid:06d}.jpg")
            img = cv2.imread(path)
            if img is None:
                raise FileNotFoundError(path)
            if (img.shape[1], img.shape[0]) != (self.img_w, self.img_h):
                img = cv2.resize(img, (self.img_w, self.img_h))
            # BGR → RGB，归一化到 [0,1] 再标准化
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
            img = (img - IMG_MEAN) / IMG_STD
            imgs.append(img.transpose(2, 0, 1))       # HWC → CHW
        imgs = np.stack(imgs)                          # (6, 3, H, W)

        # ---- 2) 读点云并裁剪到 BEV 范围 ----
        lpath = os.path.join(sdir, "LIDAR_TOP", f"{fid:06d}.bin")
        pts = np.fromfile(lpath, dtype=np.float32).reshape(-1, 4)
        keep = ((np.abs(pts[:, 0]) < self.pc_range) &
                (np.abs(pts[:, 1]) < self.pc_range))
        pts = pts[keep]

        # ---- 3) 真值 → 数组 ----
        # 格式: [x, y, z, l, w, h, yaw]  共 7 维
        boxes, labels = [], []
        for b in frame["boxes"]:
            if b["category"] not in CAT2ID:
                continue
            boxes.append(b["center"] + b["size"] + [b["yaw"]])
            labels.append(CAT2ID[b["category"]])
        boxes = (np.array(boxes, dtype=np.float32)
                 if boxes else np.zeros((0, 7), dtype=np.float32))
        labels = (np.array(labels, dtype=np.int64)
                  if len(labels) else np.zeros((0,), dtype=np.int64))

        lidar2img = self.l2i_cache[scene].copy()       # (6, 4, 4)

        # ---- 4) 数据增强（图像/点云/真值/矩阵 必须联动）----
        if self.augment:
            boxes, pts, lidar2img = self._augment(boxes, pts, lidar2img)

        return {
            "imgs": torch.from_numpy(imgs),
            "lidar2img": torch.from_numpy(lidar2img.astype(np.float32)),
            "points": torch.from_numpy(pts),
            "gt_boxes": torch.from_numpy(boxes),
            "gt_labels": torch.from_numpy(labels),
            "meta": {"scene": scene, "frame": fid,
                     "ego_speed": frame.get("ego_speed", 0.0)},
        }

    # ---------------- 数据增强 ----------------
    def _augment(self, boxes, pts, lidar2img):
        """
        BEV 空间增强 —— 关键是「所有东西一起变」
        
        只变真值不变 lidar2img，会导致图像和真值对不上，训练必然失败。
        这是本节最容易踩的坑。
        """
        # --- 随机绕 z 轴旋转 ---
        angle = np.random.uniform(-np.pi / 12, np.pi / 12)   # ±15°
        c, s = np.cos(angle), np.sin(angle)
        R = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]], dtype=np.float64)

        # --- 随机缩放 ---
        scale = np.random.uniform(0.95, 1.05)

        # 组合成 4×4 变换（作用于自车系）
        T = np.eye(4, dtype=np.float64)
        T[:3, :3] = R * scale

        # 1) 变换点云
        if len(pts):
            pts = pts.copy()
            pts[:, :3] = (T[:3, :3] @ pts[:, :3].T).T

        # 2) 变换真值框
        if len(boxes):
            boxes = boxes.copy()
            boxes[:, :3] = (T[:3, :3] @ boxes[:, :3].T).T
            boxes[:, 3:6] *= scale             # 尺寸也要缩放
            boxes[:, 6] += angle                 # 朝向加上旋转量
            # 归一化 yaw 到 [-pi, pi]
            boxes[:, 6] = np.arctan2(np.sin(boxes[:, 6]),
                                       np.cos(boxes[:, 6]))

        # 3) 关键：lidar2img 要右乘 T 的逆
        #    因为 3D 点被 T 变换了，投影矩阵要反向补偿才能对应同一张图像
        T_inv = np.linalg.inv(T)
        lidar2img = lidar2img @ T_inv

        return boxes, pts, lidar2img


# ---------------- 变长真值的 collate ----------------
def bev_collate_fn(batch):
    """
    自定义批处理
    
    为什么需要：每帧的框数不同（有的 3 个，有的 20 个），
    默认 collate 会尝试 stack 成同形状张量而报错。
    解法：定长的 stack，变长的保持 list。
    """
    return {
        # 定长 → stack
        "imgs": torch.stack([b["imgs"] for b in batch]),
        "lidar2img": torch.stack([b["lidar2img"] for b in batch]),
        # 变长 → list
        "points": [b["points"] for b in batch],
        "gt_boxes": [b["gt_boxes"] for b in batch],
        "gt_labels": [b["gt_labels"] for b in batch],
        "meta": [b["meta"] for b in batch],
    }


# ---------------- 快速自测 ----------------
if __name__ == "__main__":
    import argparse
    from torch.utils.data import DataLoader

    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="./dataset")
    args = ap.parse_args()

    ds = BEVDataset(args.root, split="train", augment=True)
    dl = DataLoader(ds, batch_size=2, shuffle=True,
                    collate_fn=bev_collate_fn, num_workers=2)

    batch = next(iter(dl))
    print("\n=== 批次结构 ===")
    print(f"imgs      : {tuple(batch['imgs'].shape)}")
    print(f"lidar2img : {tuple(batch['lidar2img'].shape)}")
    for i, (p, b) in enumerate(zip(batch["points"], batch["gt_boxes"])):
        print(f"样本{i}: {p.shape[0]} 个点, {b.shape[0]} 个框")

    # 数值健康检查
    im = batch["imgs"]
    print(f"\n图像统计: 均值 {im.mean():.3f}, 标准差 {im.std():.3f}")
    print("（标准化后均值应接近 0，标准差接近 1）")

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据管线验证 —— 把 Dataset 输出画出来，确认没有静默错误

产出两张图:
  1. 六路图像 + 投影的 3D 框  → 验证 lidar2img 正确
  2. BEV 俯视图（点云 + 框）  → 验证真值坐标系正确
"""

import argparse
import os

import numpy as np
import cv2

from bev_dataset import BEVDataset, IMG_MEAN, IMG_STD, CATEGORIES


BOX_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)]

COLORS = [(80,220,120), (100,160,255),
          (255,200,80), (255,120,200)]


def box_to_corners(box):
    """[x,y,z,l,w,h,yaw] → (8,3) 角点（自车系）"""
    x, y, z, l, w, h, yaw = box
    # 局部系下 8 个角点
    xs = l / 2 * np.array([1,1,1,1,-1,-1,-1,-1])
    ys = w / 2 * np.array([1,-1,-1,1,1,-1,-1,1])
    zs = h / 2 * np.array([1,1,-1,-1,1,1,-1,-1])
    corners = np.stack([xs, ys, zs])              # (3,8)

    c, s = np.cos(yaw), np.sin(yaw)
    R = np.array([[c,-s,0],[s,c,0],[0,0,1]])
    corners = R @ corners
    corners += np.array([[x],[y],[z]])
    return corners.T                              # (8,3)


def denormalize(img_chw):
    """标准化的张量 → 可显示的 BGR 图"""
    img = img_chw.transpose(1, 2, 0)
    img = img * IMG_STD + IMG_MEAN
    img = np.clip(img * 255, 0, 255).astype(np.uint8)
    return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)


def draw_cameras(sample, cam_names, out_path):
    """六路图像 + 投影 3D 框，验证 lidar2img 是否正确"""
    imgs = sample["imgs"].numpy()
    l2i = sample["lidar2img"].numpy()
    boxes = sample["gt_boxes"].numpy()
    labels = sample["gt_labels"].numpy()

    panels = []
    for ci in range(len(cam_names)):
        img = denormalize(imgs[ci])
        h, w = img.shape[:2]
        M = l2i[ci]
        n_drawn = 0

        for bi, box in enumerate(boxes):
            corners = box_to_corners(box)               # (8,3)
            pts = np.hstack([corners, np.ones((8,1))]).T  # (4,8)
            proj = M @ pts                              # (4,8)
            depth = proj[2]

            # 全部角点必须在相机前方
            if np.any(depth < 0.5):
                continue
            uv = proj[:2] / depth

            # 至少部分在画面内
            if (uv[0].max() < 0 or uv[0].min() > w or
                uv[1].max() < 0 or uv[1].min() > h):
                continue

            color = COLORS[int(labels[bi]) % len(COLORS)]
            for i, j in BOX_EDGES:
                p1 = (int(uv[0,i]), int(uv[1,i]))
                p2 = (int(uv[0,j]), int(uv[1,j]))
                cv2.line(img, p1, p2, color, 2, cv2.LINE_AA)
            n_drawn += 1

        cv2.rectangle(img, (0,0), (250,26), (0,0,0), -1)
        cv2.putText(img, f"{cam_names[ci]} ({n_drawn})", (6,18),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1)
        panels.append(img)

    # 拼成 2×3 网格
    while len(panels) < 6:
        panels.append(np.zeros_like(panels[0]))
    top = np.hstack(panels[:3])
    bot = np.hstack(panels[3:6])
    cv2.imwrite(out_path, np.vstack([top, bot]))
    print(f"  相机视图 → {out_path}")


def draw_bev(sample, out_path, pc_range=51.2, size=800):
    """BEV 俯视图，验证真值坐标系是否正确"""
    canvas = np.full((size, size, 3), 18, dtype=np.uint8)
    scale = size / (2 * pc_range)

    def to_px(x, y):
        # 自车系 x 向前 → 画面向上；y 向右 → 画面向右
        u = int(size / 2 + y * scale)
        v = int(size / 2 - x * scale)
        return u, v

    # 距离环
    for r in [10, 20, 30, 40, 50]:
        cv2.circle(canvas, (size//2, size//2),
                   int(r*scale), (45,52,63), 1)
        cv2.putText(canvas, f"{r}m",
                    (size//2+4, size//2-int(r*scale)+14),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.35, (80,90,105), 1)

    # 点云
    pts = sample["points"].numpy()
    for p in pts[::3]:                    # 抽稀加速
        u, v = to_px(p[0], p[1])
        if 0 <= u < size and 0 <= v < size:
            canvas[v, u] = (70, 80, 95)

    # 真值框（俯视 → 只画底面 4 角）
    boxes = sample["gt_boxes"].numpy()
    labels = sample["gt_labels"].numpy()
    for bi, box in enumerate(boxes):
        corners = box_to_corners(box)[:4]      # 顶面 4 角
        poly = np.array([to_px(c[0], c[1]) for c in corners], np.int32)
        color = COLORS[int(labels[bi]) % len(COLORS)]
        cv2.polylines(canvas, [poly], True, color, 2, cv2.LINE_AA)
        # 朝向指示线（从中心指向车头）
        cx, cy = to_px(box[0], box[1])
        hx = box[0] + np.cos(box[6]) * box[3] / 2
        hy = box[1] + np.sin(box[6]) * box[3] / 2
        cv2.line(canvas, (cx,cy), to_px(hx,hy), color, 1, cv2.LINE_AA)

    # 自车
    cv2.rectangle(canvas,
                  (size//2-7, size//2-14),
                  (size//2+7, size//2+14), (255,190,60), -1)

    # 信息
    m = sample["meta"]
    cv2.putText(canvas, f"{m['scene']} frame {m['frame']} | "
                f"{len(boxes)} boxes | {len(pts)} pts",
                (10, 22), cv2.FONT_HERSHEY_SIMPLEX,
                0.5, (220,230,240), 1)
    cv2.putText(canvas, "UP = vehicle forward (+x)",
                (10, size-14), cv2.FONT_HERSHEY_SIMPLEX,
                0.42, (110,120,135), 1)

    cv2.imwrite(out_path, canvas)
    print(f"  BEV 视图  → {out_path}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="./dataset")
    ap.add_argument("--out", default="./verify")
    ap.add_argument("--n", type=int, default=3)
    ap.add_argument("--augment", action="store_true")
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    ds = BEVDataset(args.root, split="train", augment=args.augment)

    for i in range(min(args.n, len(ds))):
        idx = i * max(1, len(ds) // max(args.n, 1))
        print(f"\n样本 {idx}:")
        sample = ds[idx]
        draw_cameras(sample, ds.cam_names,
                     os.path.join(args.out, f"{i:02d}_cameras.jpg"))
        draw_bev(sample, os.path.join(args.out, f"{i:02d}_bev.jpg"))

    print("\n检查要点:")
    print("  ① 相机图上的框应精确套住车辆 → lidar2img 正确")
    print("  ② BEV 图上框应贴合点云簇     → 真值坐标系正确")
    print("  ③ 朝向线应指向车头方向       → yaw 定义正确")
    print("  ④ 开 --augment 后两图仍应对齐 → 增强联动正确")

conda activate ad

# ① Dataset 自测
python bev_dataset.py --root ./dataset

# ② 可视化验证（不带增强）
python verify_pipeline.py --root ./dataset --out ./verify

# ③ 验证增强联动是否正确 —— 这一步很关键
python verify_pipeline.py --root ./dataset --out ./verify_aug --augment

# ④ 拉回本地看图
# scp -r l20:~/work/verify* ./