training/pretrain/prepare_bdd100k_data.py

#!/usr/bin/env python3
"""
BDD100K数据解析与任务生成 (v3 - JSON Label + 更充分利用 Lane/Drivable)

目标
- 解析你本地的 BDD100K（Scalabel 单文件 JSON 标注 + 分目录存储）
- 生成 Stage A 可训练的单帧任务样本，并保存到:
  PROJECT_ROOT/data/dataset/pretrain/train/bdd100k_tasks.pkl

相较 v2 的主要升级
1) Label 可强制为 JSON 字符串（默认开启），便于后续自动评测 / SFT / DPO / reward 计算
2) lane 信息更充分利用：统计 style + direction + continuity + category（BDD100K lane 标注的关键属性）
3) 任务 label 中增加可追溯字段：label_json_path / drivable_id_path（若存在）
4) detection 可选附带 top-k 目标的粗空间位置（left/center/right, top/mid/bottom），默认关闭以控制 token
5) 更稳健的图片路径解析：对每个 split 建立一次图片索引（避免逐样本大量 os.path.exists）

生成 4 类任务（每图最多 4 条）：
1) bdd_attributes : weather/scene/timeofday
2) bdd_detection  : 交通要素统计摘要（可选 top-k 空间粗定位）
3) bdd_drivable   : 可行驶区域 + lane marking 摘要（direct/alternative 比例优先来自 drivable_id mask）
4) bdd_risk       : 弱监督粗风险（属性 + 交通密度 + 弱势交通参与者）

运行
- 直接运行（默认 JSON label）:
    python prepare_bdd100k_data.py
- 输出自然语言 label（不推荐）:
    python prepare_bdd100k_data.py --label_format text
- 采样（调试用）:
    python prepare_bdd100k_data.py --sample_ratio 0.02
- 额外导出 jsonl（体量较大，默认不导出）:
    python prepare_bdd100k_data.py --export_jsonl --jsonl_max_per_split 20000

"""

import argparse
import json
import pickle
import random
from collections import Counter
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
from PIL import Image
from tqdm import tqdm

# ----------------- 固定随机种子（保证可复现） -----------------
random.seed(42)


# ========================= 配置 =========================
BDD_ROOT = Path("PROJECT_ROOT/BDD-100K/bdd100k")
OUTPUT_DIR = Path("PROJECT_ROOT/data/dataset/pretrain/train")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# BDD100K detection 常见 10 类（与你当前生成逻辑一致）
DETECTION_CATEGORIES = [
    "car", "bus", "truck", "person", "rider",
    "bike", "motor", "train", "traffic light", "traffic sign"
]

# BDD100K keyframe 通常为 720p；这里用于 detection 空间粗定位（不打开图片也能估相对位置）
# 若你的 images 实际分辨率不同，可在命令行覆盖（--image_size_w/h）
DEFAULT_IMAGE_W = 1280
DEFAULT_IMAGE_H = 720

# detection 的粗空间信息（会增加 token，建议默认 False）
DEFAULT_INCLUDE_TOPK_SPATIAL = False
DEFAULT_TOPK = 10


# ========================= label 输出控制 =========================
def to_json_str(obj: Dict) -> str:
    """稳定、紧凑的 JSON 字符串（用于监督信号）"""
    return json.dumps(obj, ensure_ascii=False, separators=(",", ":"), sort_keys=True)


def wrap_label(label_obj: Dict, label_text: str, label_format: str) -> str:
    """在不改训练器的前提下，把 label 统一塞进字符串字段"""
    if label_format.lower() == "json":
        return to_json_str(label_obj)
    return label_text


def wrap_prompt(base: str, schema_hint: str, label_format: str) -> str:
    """如果 label_format=json，则强制模型输出 JSON，避免自由文本漂移"""
    if label_format.lower() != "json":
        return base
    return (
        base
        + "\nReturn ONLY a single JSON object with the following schema:\n"
        + schema_hint
        + "\nDo not add any extra text."
    )


# ========================= BDD100K标注解析器 =========================
class BDD100KParser:
    """BDD100K JSON标注解析器（适配 Scalabel 单文件标注 + 分目录存储）"""

    def __init__(self, bdd_root: Path):
        self.bdd_root = bdd_root
        self.images_dir = bdd_root / "images" / "100k"
        self.labels_dir = bdd_root / "labels" / "100k"
        self.drivable_dir = bdd_root / "drivable_maps" / "labels"

        # split -> { filename: path, stem: path }
        self._image_index: Dict[str, Dict[str, Path]] = {}

    # ---------- path helpers ----------
    @staticmethod
    def _ensure_jpg(name: str) -> str:
        if not name:
            return ""
        return name if name.lower().endswith(".jpg") else f"{name}.jpg"

    def build_image_index(self, split: str) -> None:
        """
        为每个 split 建立一次索引：在 images/100k/{split} 下递归扫 *.jpg
        70K 扫描一次开销可接受，换来后续 resolve_image_path 的稳定性与速度。
        """
        if split in self._image_index:
            return

        root = self.images_dir / split
        idx: Dict[str, Path] = {}
        if not root.exists():
            self._image_index[split] = idx
            return

        for p in root.rglob("*.jpg"):
            idx[p.name] = p
            idx[p.stem] = p  # 允许用不带 .jpg 的 name 查
        self._image_index[split] = idx

    def resolve_image_path(self, split: str, name: str) -> Optional[Path]:
        """
        优先用索引查；若索引未建则 fallback 两种常见路径：
        1) images/100k/{split}/{name}.jpg
        2) images/100k/{split}/{stem[:4]}/{name}.jpg
        """
        name_jpg = self._ensure_jpg(name)
        stem = Path(name_jpg).stem
        prefix = stem[:4]

        # index lookup
        if split in self._image_index:
            idx = self._image_index[split]
            if name_jpg in idx:
                return idx[name_jpg]
            if stem in idx:
                return idx[stem]

        candidates = [
            self.images_dir / split / name_jpg,
            self.images_dir / split / prefix / name_jpg,
        ]
        for p in candidates:
            if p.exists():
                return p
        return None

    def resolve_drivable_map_path(self, split: str, name: str) -> Optional[Path]:
        """
        drivable_maps/labels/{split}/{prefix}/{stem}_drivable_id.png
        prefix 通常是 stem 的前4位
        """
        name_jpg = self._ensure_jpg(name)
        stem = Path(name_jpg).stem
        prefix = stem[:4]
        candidates = [
            self.drivable_dir / split / f"{stem}_drivable_id.png",
            self.drivable_dir / split / prefix / f"{stem}_drivable_id.png",
        ]
        for p in candidates:
            if p.exists():
                return p
        return None

    # ---------- label loading ----------
    def load_labels(self, split: str) -> List[Dict]:
        """
        递归扫描 labels/100k/{split}/**/*.json
        返回每个 json 的 dict（frame）
        """
        label_dir = self.labels_dir / split
        if not label_dir.exists():
            print(f"⚠️  标注目录不存在: {label_dir}")
            return []

        print(f"加载 {split} 标注: {label_dir}")
        json_files = sorted(label_dir.rglob("*.json"))
        print(f"  找到 {len(json_files)} 个标注文件")

        frames: List[Dict] = []
        for json_file in tqdm(json_files, desc=f"Loading {split}"):
            try:
                frame = json.loads(json_file.read_text())
                # name 可能不带 .jpg
                if "name" in frame and isinstance(frame["name"], str):
                    frame["name"] = self._ensure_jpg(frame["name"])
                else:
                    frame["name"] = self._ensure_jpg(json_file.stem)

                frame["_label_path"] = str(json_file)  # 便于追溯/调试
                frames.append(frame)
            except Exception as e:
                print(f"  ⚠️  读取失败 {json_file}: {e}")
                continue

        print(f"  成功加载 {len(frames)} 个标注")
        return frames

    # ---------- parsers ----------
    def parse_attributes(self, frame: Dict) -> Dict:
        attrs = frame.get("attributes", {}) or {}
        return {
            "weather": attrs.get("weather", "undefined"),
            "scene": attrs.get("scene", "undefined"),
            "timeofday": attrs.get("timeofday", "undefined"),
        }

    @staticmethod
    def _get_objects(frame: Dict) -> List[Dict]:
        """
        Scalabel（BDD100K）常见结构：
          {"frames":[{"objects":[...]}], "attributes":{...}, "name":...}
        """
        frames = frame.get("frames", [])
        if isinstance(frames, list) and frames and isinstance(frames[0], dict):
            objs = frames[0].get("objects", [])
            if isinstance(objs, list):
                return objs
        objs = frame.get("objects", [])
        return objs if isinstance(objs, list) else []

    def parse_detections(self, frame: Dict) -> List[Dict]:
        detections: List[Dict] = []
        for obj in self._get_objects(frame):
            category = obj.get("category", "")
            if category not in DETECTION_CATEGORIES:
                continue
            box2d = obj.get("box2d")
            if not box2d:
                continue
            detections.append(
                {
                    "category": category,
                    "box2d": box2d,
                    "attributes": obj.get("attributes", {}) or {},
                }
            )
        return detections

    def parse_drivable_area(self, frame: Dict, split: str) -> Dict:
        """
        优先用 drivable_id.png 统计 direct/alternative 比例；
        若不存在，则根据 poly2d 的 area/* 类别粗略判断。
        """
        info: Dict = {
            "has_direct": False,
            "has_alternative": False,
            "direct_ratio": None,
            "alternative_ratio": None,
            "background_ratio": None,
            "source": "none",
            "num_polygons": 0,
            "drivable_id_path": None,
        }

        name = frame.get("name", "")
        mask_path = self.resolve_drivable_map_path(split, name)
        if mask_path is not None:
            try:
                mask = np.array(Image.open(mask_path))
                if mask.ndim == 3:
                    mask = mask[:, :, 0]
                total = float(mask.size)
                if total > 0:
                    bg = float(np.sum(mask == 0)) / total
                    direct = float(np.sum(mask == 1)) / total
                    alt = float(np.sum(mask == 2)) / total
                    info.update(
                        {
                            "has_direct": direct > 1e-4,
                            "has_alternative": alt > 1e-4,
                            "direct_ratio": direct,
                            "alternative_ratio": alt,
                            "background_ratio": bg,
                            "source": "drivable_id",
                            "drivable_id_path": str(mask_path),
                        }
                    )
                    return info
            except Exception:
                pass  # 失败则 fallback

        # fallback: poly2d categories
        for obj in self._get_objects(frame):
            cat = obj.get("category", "")
            poly = obj.get("poly2d")
            if not poly:
                continue
            if isinstance(cat, str) and cat.startswith("area/"):
                info["num_polygons"] += 1
                # 兼容不同命名
                if "drivable" in cat:
                    info["has_direct"] = True
                if "alternative" in cat:
                    info["has_alternative"] = True

        if info["num_polygons"] > 0:
            info["source"] = "poly2d"
        return info

    def parse_lanes(self, frame: Dict) -> Dict:
        """
        lane poly2d 的类别通常为 lane/*。
        BDD100K lane 标注（论文）强调三属性：direction、continuity、category（并常带 style/continuity 等字段）。
        我们尽可能从 attributes 中抽取：style/direction/continuity，并统计 category（lane/ 后缀）。 
        """
        style = Counter()
        direction = Counter()
        continuity = Counter()
        category = Counter()
        num_markings = 0

        for obj in self._get_objects(frame):
            cat = obj.get("category", "")
            poly = obj.get("poly2d")
            if not poly:
                continue
            if not (isinstance(cat, str) and cat.startswith("lane/")):
                continue

            num_markings += 1
            category[cat.split("/", 1)[1]] += 1
            attrs = obj.get("attributes", {}) or {}

            # style：你的样例里用 attributes.style；部分工具会写 laneStyle
            v_style = (attrs.get("style") or attrs.get("laneStyle") or "unknown")
            style[str(v_style).lower()] += 1

            # direction：BDD100K lane 标注常用 parallel/perpendicular（论文），工具里也可能用 vertical
            v_dir = (attrs.get("direction") or attrs.get("laneDirection") or "unknown")
            direction[str(v_dir).lower()] += 1

            # continuity：full/dashed 或类似字段；不同导出可能键名不同
            v_cont = (attrs.get("continuity") or attrs.get("laneContinuity") or "unknown")
            continuity[str(v_cont).lower()] += 1

        return {
            "has_lanes": num_markings > 0,
            "num_markings": num_markings,
            "style": dict(style),
            "direction": dict(direction),
            "continuity": dict(continuity),
            "category": dict(category),
        }


# ========================= 任务生成器 =========================
class BDDTaskGenerator:
    def __init__(
        self,
        parser: BDD100KParser,
        label_format: str = "json",
        include_topk_spatial: bool = DEFAULT_INCLUDE_TOPK_SPATIAL,
        topk: int = DEFAULT_TOPK,
        image_size: Tuple[int, int] = (DEFAULT_IMAGE_W, DEFAULT_IMAGE_H),
    ):
        self.parser = parser
        self.label_format = label_format
        self.include_topk_spatial = include_topk_spatial
        self.topk = max(1, int(topk))
        self.image_w, self.image_h = int(image_size[0]), int(image_size[1])

    # ---------- small helpers ----------
    def _common_metadata(self, frame: Dict, name: str) -> Dict:
        md = {
            "frame_name": name,
            "dataset": "bdd100k",
            "label_json_path": frame.get("_label_path"),
        }
        return md

    def _bbox_to_region(self, box2d: Dict) -> Dict:
        """
        把 bbox 映射成粗空间区域（不打开图片）
        - x: left/center/right by bbox center x
        - y: top/mid/bottom by bbox center y
        """
        x1, y1, x2, y2 = box2d.get("x1", 0), box2d.get("y1", 0), box2d.get("x2", 0), box2d.get("y2", 0)
        cx = (float(x1) + float(x2)) / 2.0
        cy = (float(y1) + float(y2)) / 2.0

        rx = cx / max(1.0, float(self.image_w))
        ry = cy / max(1.0, float(self.image_h))

        if rx < 1/3:
            x_bin = "left"
        elif rx < 2/3:
            x_bin = "center"
        else:
            x_bin = "right"

        if ry < 1/3:
            y_bin = "top"
        elif ry < 2/3:
            y_bin = "middle"
        else:
            y_bin = "bottom"

        area = max(0.0, (float(x2) - float(x1))) * max(0.0, (float(y2) - float(y1)))
        area_ratio = area / max(1.0, float(self.image_w * self.image_h))
        return {"x": x_bin, "y": y_bin, "area_ratio": round(area_ratio, 4)}

    # ---------- tasks ----------
    def generate_task1_attributes(self, frame: Dict, split: str) -> Optional[Dict]:
        name = frame.get("name", "")
        image_path = self.parser.resolve_image_path(split, name)
        if image_path is None:
            return None

        attrs = self.parser.parse_attributes(frame)
        if all(v == "undefined" for v in attrs.values()):
            return None

        label_obj = {"weather": attrs["weather"], "scene": attrs["scene"], "timeofday": attrs["timeofday"]}
        label_text = f"Weather: {attrs['weather']}; Scene: {attrs['scene']}; Time: {attrs['timeofday']}"

        return {
            "task": "bdd_attributes",
            "subtask": "scene_attributes",
            "image_path": str(image_path),
            "user_prompt": wrap_prompt(
                "Describe the driving scene attributes: weather, scene type, and time of day.",
                '{"weather":"...","scene":"...","timeofday":"..."}',
                self.label_format,
            ),
            "label": wrap_label(label_obj, label_text, self.label_format),
            "difficulty": "easy",
            "metadata": {**self._common_metadata(frame, name), **attrs},
        }

    def generate_task2_detection_summary(self, frame: Dict, split: str) -> Optional[Dict]:
        name = frame.get("name", "")
        image_path = self.parser.resolve_image_path(split, name)
        if image_path is None:
            return None

        detections = self.parser.parse_detections(frame)
        if len(detections) == 0:
            return None

        category_counts = Counter(d["category"] for d in detections)

        # 便于人读的 text label（当 label_format=text）
        parts = []
        for cat, count in category_counts.most_common():
            if count == 1:
                parts.append(f"1 {cat}")
            else:
                plural = cat + "s" if cat not in ["person", "traffic light", "traffic sign"] else (
                    "people" if cat == "person" else cat + "s"
                )
                parts.append(f"{count} {plural}")
        summary = f"There is {parts[0]}." if len(parts) == 1 else f"There are {', '.join(parts[:-1])}, and {parts[-1]}."

        label_obj: Dict = {
            "counts": dict(category_counts),
            "num_objects": len(detections),
        }

        if self.include_topk_spatial:
            # 选面积最大的 top-k（近似“更重要”）
            det_sorted = sorted(
                detections,
                key=lambda d: max(0.0, (float(d["box2d"].get("x2", 0)) - float(d["box2d"].get("x1", 0))))
                              * max(0.0, (float(d["box2d"].get("y2", 0)) - float(d["box2d"].get("y1", 0)))),
                reverse=True
            )
            topk = det_sorted[: self.topk]
            label_obj["topk_objects"] = [
                {"category": d["category"], **self._bbox_to_region(d["box2d"])}
                for d in topk
            ]

        return {
            "task": "bdd_detection",
            "subtask": "traffic_elements",
            "image_path": str(image_path),
            "user_prompt": wrap_prompt(
                "Summarize the traffic elements in this image (vehicles, pedestrians, traffic lights/signs).",
                '{"counts":{"car":3,"person":1,...},"num_objects":N'
                + (', "topk_objects":[{"category":"car","x":"left|center|right","y":"top|middle|bottom","area_ratio":0.0123},...]' if self.include_topk_spatial else "")
                + "}",
                self.label_format,
            ),
            "label": wrap_label(label_obj, summary, self.label_format),
            "difficulty": "easy",
            "metadata": {
                **self._common_metadata(frame, name),
                "num_objects": len(detections),
                "categories": dict(category_counts),
            },
        }

    def generate_task3_drivable_area(self, frame: Dict, split: str) -> Optional[Dict]:
        name = frame.get("name", "")
        image_path = self.parser.resolve_image_path(split, name)
        if image_path is None:
            return None

        drivable_info = self.parser.parse_drivable_area(frame, split)
        lane_info = self.parser.parse_lanes(frame)

        # 如果什么都解析不到，就跳过（避免噪声）
        if (
            not drivable_info.get("has_direct")
            and not drivable_info.get("has_alternative")
            and not lane_info.get("has_lanes")
            and drivable_info.get("num_polygons", 0) == 0
            and drivable_info.get("source") == "none"
        ):
            return None

        parts = []
        if drivable_info.get("source") == "drivable_id" and drivable_info.get("direct_ratio") is not None:
            d = drivable_info["direct_ratio"] * 100
            a = drivable_info["alternative_ratio"] * 100
            parts.append(f"Drivable area coverage: direct {d:.1f}%, alternative {a:.1f}%")
        else:
            if drivable_info.get("has_direct"):
                parts.append("Direct drivable path available")
            if drivable_info.get("has_alternative"):
                parts.append("Alternative drivable region exists")
            if not drivable_info.get("has_direct") and not drivable_info.get("has_alternative"):
                parts.append("Drivable area present")

        if lane_info.get("has_lanes"):
            # style 里常见 solid/dashed/full/unknown，尽量人读
            s = lane_info.get("style", {})
            if s.get("solid", 0) and s.get("dashed", 0):
                parts.append("Lane markings: mixed solid and dashed")
            elif s.get("solid", 0):
                parts.append("Lane markings: solid")
            elif s.get("dashed", 0):
                parts.append("Lane markings: dashed")
            else:
                parts.append("Lane markings present")

        label_text = "; ".join(parts) + "."

        # JSON label：尽量结构化，数值做 round 便于稳定训练
        dri = dict(drivable_info)
        if dri.get("direct_ratio") is not None:
            dri["direct_ratio"] = round(float(dri["direct_ratio"]), 6)
        if dri.get("alternative_ratio") is not None:
            dri["alternative_ratio"] = round(float(dri["alternative_ratio"]), 6)
        if dri.get("background_ratio") is not None:
            dri["background_ratio"] = round(float(dri["background_ratio"]), 6)

        label_obj = {"drivable": dri, "lanes": lane_info}

        return {
            "task": "bdd_drivable",
            "subtask": "drivable_description",
            "image_path": str(image_path),
            "user_prompt": wrap_prompt(
                "Describe the drivable area and lane marking structure in this scene.",
                '{"drivable":{"source":"drivable_id|poly2d|none","direct_ratio":0.0,"alternative_ratio":0.0,"background_ratio":0.0,'
                '"has_direct":true,"has_alternative":false,"drivable_id_path":"...|null"},'
                '"lanes":{"has_lanes":true,"num_markings":N,"style":{...},"direction":{...},"continuity":{...},"category":{...}}}',
                self.label_format,
            ),
            "label": wrap_label(label_obj, label_text, self.label_format),
            "difficulty": "medium",
            "metadata": {**self._common_metadata(frame, name), **drivable_info, **lane_info},
        }

    def generate_task4_risk_level(self, frame: Dict, split: str) -> Optional[Dict]:
        name = frame.get("name", "")
        image_path = self.parser.resolve_image_path(split, name)
        if image_path is None:
            return None

        attrs = self.parser.parse_attributes(frame)
        detections = self.parser.parse_detections(frame)

        risk_score = 0
        risk_factors: List[str] = []

        # weather
        weather = attrs.get("weather", "undefined")
        if weather in ["rainy", "snowy", "foggy"]:
            risk_score += 2
            risk_factors.append(f"{weather} weather")
        elif weather == "overcast":
            risk_score += 1
            risk_factors.append("overcast conditions")

        # time
        timeofday = attrs.get("timeofday", "undefined")
        if timeofday == "night":
            risk_score += 2
            risk_factors.append("nighttime")
        elif timeofday in ["dawn/dusk", "dawn", "dusk"]:
            risk_score += 1
            risk_factors.append("low light")

        # scene
        scene = attrs.get("scene", "undefined")
        if scene in ["city street", "residential"]:
            risk_score += 1
            risk_factors.append("urban area")

        # traffic density
        num_vehicles = sum(1 for d in detections if d["category"] in ["car", "bus", "truck"])
        if num_vehicles >= 10:
            risk_score += 2
            risk_factors.append("high traffic density")
        elif num_vehicles >= 5:
            risk_score += 1
            risk_factors.append("moderate traffic")

        # vulnerable users
        num_vulnerable = sum(1 for d in detections if d["category"] in ["person", "rider", "bike"])
        if num_vulnerable > 0:
            risk_score += 1
            risk_factors.append(f"{num_vulnerable} vulnerable road user(s)")

        if risk_score >= 5:
            risk_level = "High risk"
        elif risk_score >= 3:
            risk_level = "Medium risk"
        elif risk_score >= 1:
            risk_level = "Low-medium risk"
        else:
            risk_level = "Low risk"

        label_text = f"{risk_level}: {', '.join(risk_factors)}" if risk_factors else f"{risk_level}: normal conditions"

        label_obj = {
            "risk_level": risk_level,
            "risk_score": int(risk_score),
            "risk_factors": risk_factors,
            "evidence": {
                "weather": attrs.get("weather", "undefined"),
                "scene": attrs.get("scene", "undefined"),
                "timeofday": attrs.get("timeofday", "undefined"),
                "num_vehicles": int(num_vehicles),
                "num_vulnerable": int(num_vulnerable),
                "num_objects": int(len(detections)),
            },
        }

        return {
            "task": "bdd_risk",
            "subtask": "risk_assessment",
            "image_path": str(image_path),
            "user_prompt": wrap_prompt(
                "Assess the risk level of this driving scenario based on environmental and traffic conditions.",
                '{"risk_level":"Low risk|Low-medium risk|Medium risk|High risk","risk_score":N,"risk_factors":["..."],'
                '"evidence":{"weather":"...","scene":"...","timeofday":"...","num_vehicles":N,"num_vulnerable":N,"num_objects":N}}',
                self.label_format,
            ),
            "label": wrap_label(label_obj, label_text, self.label_format),
            "difficulty": "medium",
            "metadata": {**self._common_metadata(frame, name), "risk_score": int(risk_score), "risk_factors": risk_factors, **attrs},
        }


# ========================= 主流程 =========================
def process_split(
    parser: BDD100KParser,
    generator: BDDTaskGenerator,
    split: str,
    sample_ratio: float = 1.0,
) -> Dict[str, List[Dict]]:
    print(f"\n{'='*70}")
    print(f"处理 {split.upper()} Split")
    print("=" * 70)

    parser.build_image_index(split)
    frames = parser.load_labels(split)
    if len(frames) == 0:
        print(f"⚠️  {split} split没有数据")
        return {"bdd_attributes": [], "bdd_detection": [], "bdd_drivable": [], "bdd_risk": []}

    if sample_ratio < 1.0:
        n_sample = max(1, int(len(frames) * sample_ratio))
        frames = random.sample(frames, n_sample)
        print(f"  采样: {len(frames)} / {int(len(frames)/sample_ratio)}")

    tasks: Dict[str, List[Dict]] = {"bdd_attributes": [], "bdd_detection": [], "bdd_drivable": [], "bdd_risk": []}

    print("\n生成任务样本...")
    for frame in tqdm(frames, desc=f"Generating {split}"):
        s1 = generator.generate_task1_attributes(frame, split)
        if s1:
            tasks["bdd_attributes"].append(s1)

        s2 = generator.generate_task2_detection_summary(frame, split)
        if s2:
            tasks["bdd_detection"].append(s2)

        s3 = generator.generate_task3_drivable_area(frame, split)
        if s3:
            tasks["bdd_drivable"].append(s3)

        s4 = generator.generate_task4_risk_level(frame, split)
        if s4:
            tasks["bdd_risk"].append(s4)

    print(f"\n{split.upper()} 任务统计:")
    for task_name, samples in tasks.items():
        print(f"  {task_name}: {len(samples)} 样本")
    print(f"  总计: {sum(len(v) for v in tasks.values())} 样本")

    return tasks


def export_jsonl(all_data: Dict, out_path: Path, max_per_split: int = 0) -> None:
    """
    导出 jsonl（便于人工检查/交给其他模型写 code）
    max_per_split=0 表示全量导出；>0 表示每个 split 只导出最多 N 条（建议调试时用）
    """
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", encoding="utf-8") as f:
        for split, task_dict in all_data.items():
            # 混合导出，保留 task 字段即可筛
            samples = []
            for arr in task_dict.values():
                samples.extend(arr)
            if max_per_split and len(samples) > max_per_split:
                samples = random.sample(samples, max_per_split)
            for s in samples:
                rec = dict(s)
                rec["split"] = split
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")


def main():
    ap = argparse.ArgumentParser(description="Prepare BDD100K Stage-A tasks")
    ap.add_argument("--bdd_root", type=str, default=str(BDD_ROOT))
    ap.add_argument("--output_dir", type=str, default=str(OUTPUT_DIR))
    ap.add_argument("--sample_ratio", type=float, default=1.0)
    ap.add_argument("--label_format", type=str, default="json", choices=["json", "text"])
    ap.add_argument("--include_topk_spatial", action="store_true", help="bdd_detection label 中加入 top-k 粗空间位置（会增加 token）")
    ap.add_argument("--topk", type=int, default=DEFAULT_TOPK)
    ap.add_argument("--image_size_w", type=int, default=DEFAULT_IMAGE_W)
    ap.add_argument("--image_size_h", type=int, default=DEFAULT_IMAGE_H)
    ap.add_argument("--export_jsonl", action="store_true", help="额外导出 jsonl（体量较大）")
    ap.add_argument("--jsonl_max_per_split", type=int, default=0, help="jsonl 每个 split 最大导出条数（0=全量）")
    args = ap.parse_args()

    bdd_root = Path(args.bdd_root)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    print("=" * 70)
    print("BDD100K数据处理 (v3)")
    print("JSON label + lane/drivable 更充分利用 + 可追溯 metadata")
    print("=" * 70)

    parser = BDD100KParser(bdd_root)
    generator = BDDTaskGenerator(
        parser,
        label_format=args.label_format,
        include_topk_spatial=args.include_topk_spatial,
        topk=args.topk,
        image_size=(args.image_size_w, args.image_size_h),
    )

    if not parser.images_dir.exists():
        print(f"❌ 图像目录不存在: {parser.images_dir}")
        return
    if not parser.labels_dir.exists():
        print(f"❌ 标注目录不存在: {parser.labels_dir}")
        return

    print(f"✓ BDD根目录: {bdd_root}")
    print(f"✓ 图像目录: {parser.images_dir}")
    print(f"✓ 标注目录: {parser.labels_dir}")
    print(f"✓ Drivable目录: {parser.drivable_dir}")
    print(f"✓ Label Format: {args.label_format}")
    print(f"✓ include_topk_spatial: {bool(args.include_topk_spatial)} (topk={args.topk})")
    print(f"✓ image_size (for spatial binning): {args.image_size_w}x{args.image_size_h}")
    print(f"✓ sample_ratio: {args.sample_ratio}")

    all_data: Dict[str, Dict[str, List[Dict]]] = {}
    for split in ["train", "val", "test"]:
        all_data[split] = process_split(parser, generator, split, sample_ratio=args.sample_ratio)

    # 保存 pkl
    print("\n" + "=" * 70)
    print("保存数据...")
    output_file = output_dir / "bdd100k_tasks.pkl"
    with output_file.open("wb") as f:
        pickle.dump(all_data, f)
    print(f"✓ 保存到: {output_file}")

    # summary json
    summary: Dict[str, Dict] = {}
    for split in ["train", "val", "test"]:
        summary[split] = {k: len(v) for k, v in all_data[split].items()}
        summary[split]["total"] = sum(summary[split].values())

    summary_file = output_dir / "bdd100k_summary.json"
    summary_file.write_text(json.dumps(summary, ensure_ascii=False, indent=2))
    print(f"✓ 统计: {summary_file}")

    # 可选导出 jsonl
    if args.export_jsonl:
        jsonl_path = output_dir / "bdd100k_tasks.jsonl"
        print(f"导出 JSONL 到: {jsonl_path} (max_per_split={args.jsonl_max_per_split})")
        export_jsonl(all_data, jsonl_path, max_per_split=args.jsonl_max_per_split)
        print("✓ JSONL 导出完成")

    print("\n" + "=" * 70)
    print("数据处理完成 - 统计:")
    print("=" * 70)
    for split in ["train", "val", "test"]:
        print(f"\n{split.upper()}:")
        for task_name in ["bdd_attributes", "bdd_detection", "bdd_drivable", "bdd_risk"]:
            print(f"  {task_name}: {summary[split].get(task_name, 0)}")
        print("  ─────────────────")
        print(f"  总计: {summary[split]['total']}")

    print("\n✅ 完成!")
    print("\n下一步:")
    print("1. 运行 python prepare_stage_a_data.py 整合Stage A数据")
    print("2. 运行 python train_two_stage.py --stage A 开始Stage A训练")


if __name__ == "__main__":
    main()