training/pretrain/prepare_pretrain_data_adaptive.py

#!/usr/bin/env python3
"""
自适应Prompt的预训练数据准备
策略：根据annotation长度调整prompt难度，而不是修改annotation本身
"""

import json
import os
import pickle
import random
import cv2
from pathlib import Path
from typing import Dict, List, Tuple
from collections import defaultdict

random.seed(42)

# ============ 配置 ============
PRETRAIN_ROOT = Path("PROJECT_ROOT/data/dataset/pretrain")
DAD_ROOT = Path("PROJECT_ROOT/DAD/videos/training")
OUTPUT_DIR = PRETRAIN_ROOT / "train"
OUTPUT_DIR.mkdir(exist_ok=True)

NEXAR_ROOT = PRETRAIN_ROOT / "nexar"
DADA_ROOT = PRETRAIN_ROOT / "DADA-2000"

TRAIN_RATIO = 0.7
VAL_RATIO = 0.15
TEST_RATIO = 0.15

# 标注质量阈值
ANNOTATION_SHORT_THRESHOLD = 20  # 少于20字符为简单标注


# ============ Prompt Templates ============
class AdaptivePrompts:
    """自适应Prompt生成器"""
    
    # 简单标注 - 简单prompt (只要求识别对象/类型)
    SHORT_ANNOTATION_PROMPTS = [
        "What object or vehicle was involved in this accident?",
        "Identify the main entity in this traffic incident.",
        "What type of collision is shown?",
        "Briefly describe what is involved in this accident.",
    ]
    
    # 详细标注 - 详细prompt (要求完整描述)
    DETAILED_ANNOTATION_PROMPTS = [
        "Describe the accident in this image. What happened and why?",
        "Provide a detailed description of the traffic incident.",
        "Explain what led to this accident and what occurred.",
        "Describe this accident scenario in detail.",
    ]
    
    # 序列任务 - 根据标注长度调整
    SHORT_SEQUENCE_PROMPTS = [
        "Analyze this driving sequence. What type of incident occurred?",
        "What is the main object involved in this traffic sequence?",
    ]
    
    DETAILED_SEQUENCE_PROMPTS = [
        "Analyze this driving video sequence. Describe the accident: what happened, when, and why?",
        "Based on this video sequence, provide a detailed description of the accident.",
    ]
    
    @staticmethod
    def get_accident_prompt(annotation: str, is_sequence: bool = False):
        """
        根据annotation长度选择合适的prompt
        
        Args:
            annotation: accident_type标注
            is_sequence: 是否为序列任务
        
        Returns:
            user_prompt, difficulty_level
        """
        if not annotation or annotation.lower() in ['null', 'none', 'unknown', '']:
            annotation = ""
        
        char_count = len(annotation.strip())
        
        # 判断标注是否简单
        is_short = char_count < ANNOTATION_SHORT_THRESHOLD
        
        if is_sequence:
            prompts = AdaptivePrompts.SHORT_SEQUENCE_PROMPTS if is_short else AdaptivePrompts.DETAILED_SEQUENCE_PROMPTS
        else:
            prompts = AdaptivePrompts.SHORT_ANNOTATION_PROMPTS if is_short else AdaptivePrompts.DETAILED_ANNOTATION_PROMPTS
        
        prompt = random.choice(prompts)
        difficulty = "medium" if is_short else "hard"
        
        return prompt, difficulty, is_short


# ============ DAD视频处理 ============
class DADProcessor:
    """DAD数据集处理器"""
    
    def __init__(self, dad_root: Path, output_dir: Path):
        self.dad_root = dad_root
        self.output_dir = output_dir / "dad_frames"
        self.output_dir.mkdir(parents=True, exist_ok=True)
    
    def extract_frames(self, video_path: Path, output_case_dir: Path, 
                       fps: int = 20, max_frames: int = 300):
        """从视频中提取帧"""
        output_case_dir.mkdir(parents=True, exist_ok=True)
        
        cap = cv2.VideoCapture(str(video_path))
        if not cap.isOpened():
            print(f"无法打开视频: {video_path}")
            return 0
        
        original_fps = cap.get(cv2.CAP_PROP_FPS)
        
        if fps >= original_fps:
            frame_interval = 1
        else:
            frame_interval = int(original_fps / fps)
        
        frame_count = 0
        saved_count = 0
        
        while True:
            ret, frame = cap.read()
            if not ret or saved_count >= max_frames:
                break
            
            if frame_count % frame_interval == 0:
                frame_path = output_case_dir / f"{saved_count:06d}.jpg"
                cv2.imwrite(str(frame_path), frame)
                saved_count += 1
            
            frame_count += 1
        
        cap.release()
        return saved_count
    
    def process_dad_dataset(self):
        """处理DAD数据集"""
        dad_data = []
        
        for split in ["positive", "negative"]:
            split_dir = self.dad_root / split
            if not split_dir.exists():
                print(f"DAD {split} 目录不存在: {split_dir}")
                continue
            
            video_files = list(split_dir.glob("*.mp4")) + list(split_dir.glob("*.avi"))
            
            print(f"\n处理 DAD {split}: {len(video_files)} 视频")
            
            for vid_file in video_files:
                case_id = f"dad_{split}_{vid_file.stem}"
                case_dir = self.output_dir / case_id
                
                n_frames = self.extract_frames(vid_file, case_dir)
                
                if n_frames == 0:
                    continue
                
                # 生成annotation - 简单标注
                annotation = {
                    "id": case_id,
                    "dataset": "dad",
                    "source_video": str(vid_file),
                    "accident": (split == "positive"),
                    "accident_type": "accident" if split == "positive" else "normal",  # 简单标注
                    "weather": "Unknown",
                    "road_type": "Unknown",
                    "time_of_day": "Unknown",
                    "risky_time": None,
                    "accident_time": None,
                    "n_frames": n_frames,
                    "fps": 20
                }
                
                with open(case_dir / "annotation.json", 'w') as f:
                    json.dump(annotation, f, indent=2)
                
                annotation["case_dir"] = str(case_dir)
                dad_data.append(annotation)
                
                if len(dad_data) % 20 == 0:
                    print(f"已处理: {len(dad_data)} DAD视频...")
        
        print(f"\n✓ DAD数据集处理完成: {len(dad_data)} cases")
        return dad_data


# ============ 数据加载 ============
def load_all_annotations(include_dad: bool = True):
    """加载所有annotation"""
    all_data = []
    
    # 加载NEXAR
    print("加载 NEXAR...")
    for split in ["positive", "negative"]:
        split_dir = NEXAR_ROOT / split
        if not split_dir.exists():
            continue
        for case_dir in sorted(split_dir.iterdir()):
            if not case_dir.is_dir():
                continue
            anno_file = case_dir / "annotation.json"
            if not anno_file.exists():
                continue
            
            with open(anno_file) as f:
                data = json.load(f)
                data["dataset"] = "nexar"
                data["case_dir"] = str(case_dir)
                if "id" not in data:
                    data["id"] = case_dir.name
                all_data.append(data)
    
    # 加载DADA-2000
    print("加载 DADA-2000...")
    for case_dir in sorted(DADA_ROOT.iterdir()):
        if not case_dir.is_dir():
            continue
        anno_file = case_dir / "annotation.json"
        if not anno_file.exists():
            continue
        
        with open(anno_file) as f:
            data = json.load(f)
            data["dataset"] = "dada"
            data["case_dir"] = str(case_dir)
            data["id"] = case_dir.name
            all_data.append(data)
    
    # 加载DAD
    if include_dad:
        print("处理 DAD数据集...")
        dad_processor = DADProcessor(DAD_ROOT, OUTPUT_DIR.parent)
        dad_data = dad_processor.process_dad_dataset()
        all_data.extend(dad_data)
    
    print(f"\n总计: {len(all_data)} 案例")
    
    # 统计
    stats = defaultdict(int)
    for d in all_data:
        stats[d["dataset"]] += 1
    print("数据集分布:")
    for ds, count in stats.items():
        print(f"  {ds}: {count}")
    
    return all_data


def split_data(all_data):
    """按数据集分层划分"""
    by_dataset = defaultdict(list)
    for data in all_data:
        by_dataset[data["dataset"]].append(data)
    
    train_data = []
    val_data = []
    test_data = []
    
    for dataset, items in by_dataset.items():
        random.shuffle(items)
        n = len(items)
        n_train = int(n * TRAIN_RATIO)
        n_val = int(n * VAL_RATIO)
        
        train_data.extend(items[:n_train])
        val_data.extend(items[n_train:n_train + n_val])
        test_data.extend(items[n_train + n_val:])
    
    print(f"\n数据划分:")
    print(f"  训练: {len(train_data)}")
    print(f"  验证: {len(val_data)}")
    print(f"  测试: {len(test_data)}")
    
    return train_data, val_data, test_data


# ============ 任务3: 事故描述（自适应prompt）============
def prepare_task3_accident_description_adaptive(data_split, split_name):
    """
    事故描述任务 - 根据annotation长度自适应调整prompt
    
    短标注 → 简单prompt (识别对象)
    长标注 → 详细prompt (完整描述)
    """
    samples = []
    
    annotation_stats = {
        'short': 0,
        'detailed': 0
    }
    
    for data in data_split:
        has_accident = data.get("accident", False)
        if isinstance(has_accident, str):
            has_accident = has_accident.lower() == "true"
        
        # 只使用有事故的cases
        if not has_accident:
            continue
        
        case_dir = Path(data["case_dir"])
        frames = sorted([f for f in case_dir.glob("*.jpg")])
        
        if len(frames) < 3:
            continue
        
        # 获取原始标注
        accident_type = data.get("accident_type", "")
        if not accident_type or accident_type.lower() in ["null", "none"]:
            accident_type = "Traffic incident"
        
        # 根据标注长度选择prompt
        user_prompt, difficulty, is_short = AdaptivePrompts.get_accident_prompt(
            accident_type, 
            is_sequence=False
        )
        
        # 统计
        annotation_stats['short' if is_short else 'detailed'] += 1
        
        accident_time = data.get("accident_time")
        risky_time = data.get("risky_time")
        
        if isinstance(accident_time, str):
            try:
                accident_time = int(accident_time)
            except:
                accident_time = None
        
        if isinstance(risky_time, str):
            try:
                risky_time = int(risky_time)
            except:
                risky_time = None
        
        # 找到事故相关帧
        accident_frames = []
        for idx, frame in enumerate(frames):
            frame_num = int(frame.stem)
            
            is_accident_frame = False
            if accident_time and abs(frame_num - accident_time) <= 15:
                is_accident_frame = True
            elif risky_time and abs(frame_num - risky_time) <= 20:
                is_accident_frame = True
            
            if is_accident_frame:
                accident_frames.append(idx)
        
        # 采样2-3个事故帧
        if accident_frames:
            n_samples = min(2 if is_short else 3, len(accident_frames))
            sampled = random.sample(accident_frames, n_samples)
            
            for idx in sampled:
                samples.append({
                    "task": "accident_description",
                    "subtask": "adaptive_description",
                    "image_path": str(frames[idx]),
                    "user_prompt": user_prompt,  # 自适应prompt
                    "label": accident_type,  # 保持原始标注不变
                    "difficulty": difficulty,
                    "metadata": {
                        "case_id": data["id"],
                        "dataset": data["dataset"],
                        "frame_num": int(frames[idx].stem),
                        "annotation_length": len(accident_type),
                        "is_short_annotation": is_short
                    }
                })
    
    print(f"[{split_name}] 任务3-事故描述 (自适应): {len(samples)} 样本")
    print(f"  短标注: {annotation_stats['short']} (简单prompt)")
    print(f"  详细标注: {annotation_stats['detailed']} (详细prompt)")
    
    # 统计数据集分布
    from collections import Counter
    dataset_dist = Counter(s["metadata"]["dataset"] for s in samples)
    print(f"  数据集分布:")
    for ds, count in dataset_dist.items():
        print(f"    {ds}: {count} 样本")
    
    return samples


# ============ 任务4: 序列预测（自适应prompt）============
def prepare_task4_sequence_adaptive(data_split, split_name):
    """序列预测 - 自适应prompt"""
    samples = []
    
    annotation_stats = {
        'short': 0,
        'detailed': 0
    }
    
    for data in data_split:
        case_dir = Path(data["case_dir"])
        frames = sorted([f for f in case_dir.glob("*.jpg")])
        
        if len(frames) < 12:
            continue
        
        has_accident = data.get("accident", False)
        if isinstance(has_accident, str):
            has_accident = has_accident.lower() == "true"
        
        accident_type = data.get("accident_type", "")
        if not accident_type or accident_type.lower() in ["null", "none"]:
            accident_type = "Normal driving" if not has_accident else "Traffic incident"
        
        # 根据标注长度选择prompt
        user_prompt, difficulty, is_short = AdaptivePrompts.get_accident_prompt(
            accident_type,
            is_sequence=True
        )
        
        annotation_stats['short' if is_short else 'detailed'] += 1
        
        risky_time = data.get("risky_time")
        if isinstance(risky_time, str):
            try:
                risky_time = int(risky_time)
            except:
                risky_time = None
        
        # 采样起始点
        if risky_time and risky_time > 0 and has_accident:
            start_frame = max(0, risky_time - 20)
        else:
            max_start = len(frames) - 24
            start_frame = random.randint(0, max(0, max_start))
        
        # 采样序列
        STRIDE = 8
        T_MAX = 16
        
        seq_full = list(range(start_frame, len(frames), STRIDE))
        seq_full = [str(frames[i]) for i in seq_full if i < len(frames)]
        
        if len(seq_full) > T_MAX:
            import numpy as np
            idx = np.linspace(0, len(seq_full) - 1, T_MAX).round().astype(int).tolist()
            sequence = [seq_full[j] for j in idx]
        else:
            sequence = seq_full
        
        if len(sequence) < 4:
            continue
        
        # 构造完整答案
        if has_accident:
            label = f"Accident detected. {accident_type}"
        else:
            label = f"Normal driving. {accident_type}"
        
        samples.append({
            "task": "sequence_prediction",
            "subtask": "adaptive_sequence",
            "image_sequence": sequence,
            "user_prompt": user_prompt,  # 自适应prompt
            "label": label,
            "difficulty": difficulty,
            "metadata": {
                "case_id": data["id"],
                "dataset": data["dataset"],
                "sequence_length": len(sequence),
                "has_accident": has_accident,
                "annotation_length": len(accident_type),
                "is_short_annotation": is_short,
                "start_frame": start_frame
            }
        })
    
    print(f"[{split_name}] 任务4-序列预测 (自适应): {len(samples)} 样本")
    print(f"  短标注: {annotation_stats['short']} (简单prompt)")
    print(f"  详细标注: {annotation_stats['detailed']} (详细prompt)")
    
    # 统计数据集分布
    from collections import Counter
    dataset_dist = Counter(s["metadata"]["dataset"] for s in samples)
    print(f"  数据集分布:")
    for ds, count in dataset_dist.items():
        print(f"    {ds}: {count} 样本")
    
    return samples


# ============ 其他任务保持不变 ============
def prepare_task1_scene_understanding(data_split, split_name):
    """
    场景理解任务
    注意: DAD数据集没有天气/道路标注，因此不参与此任务
    """
    samples = []
    
    for data in data_split:
        # 跳过DAD数据集（缺少环境标注）
        if data.get("dataset") == "dad":
            continue
        
        case_dir = Path(data["case_dir"])
        frames = sorted([f for f in case_dir.glob("*.jpg")])
        
        if len(frames) == 0:
            continue
        
        # 检查是否有有效的环境信息
        weather = data.get("weather", "Unknown")
        road = data.get("road_type", "Unknown")
        light = data.get("time_of_day", "") or data.get("light_conditions", "Unknown")
        
        # 如果所有信息都是Unknown，跳过
        if all(x == "Unknown" for x in [weather, road, light]):
            continue
        
        n_samples = random.randint(4, 6)
        sampled = random.sample(frames, min(n_samples, len(frames)))
        
        for frame_path in sampled:
            env_label = f"Weather: {weather}, Road: {road}, Light: {light}"
            
            has_accident = data.get("accident", False)
            if isinstance(has_accident, str):
                has_accident = has_accident.lower() == "true"
            
            risk_level = "High risk" if has_accident else "Normal"
            label = f"{env_label}. Risk: {risk_level}"
            
            samples.append({
                "task": "scene_understanding",
                "subtask": "environment",
                "image_path": str(frame_path),
                "user_prompt": "Analyze this driving scene. What are the weather, road, and lighting conditions? What is the risk level?",
                "label": label,
                "difficulty": "easy",
                "metadata": {
                    "case_id": data["id"],
                    "dataset": data["dataset"],
                    "has_accident": has_accident
                }
            })
    
    print(f"[{split_name}] 任务1-场景理解: {len(samples)} 样本")
    print(f"  ✓ DADA-2000 + NEXAR (环境信息完整)")
    print(f"  ✗ 已跳过DAD数据集 (环境信息全为Unknown)")
    
    # 统计数据集分布
    from collections import Counter
    dataset_dist = Counter(s["metadata"]["dataset"] for s in samples)
    for ds, count in dataset_dist.items():
        print(f"    {ds}: {count} 样本")
    
    return samples


def prepare_task2_binary_detection(data_split, split_name):
    """二分类检测"""
    samples = []
    
    accident_cases = []
    normal_cases = []
    
    for data in data_split:
        has_accident = data.get("accident", False)
        if isinstance(has_accident, str):
            has_accident = has_accident.lower() == "true"
        
        if has_accident:
            accident_cases.append(data)
        else:
            normal_cases.append(data)
    
    # 从有事故的cases采样
    for data in accident_cases:
        case_dir = Path(data["case_dir"])
        frames = sorted([f for f in case_dir.glob("*.jpg")])
        
        if len(frames) == 0:
            continue
        
        accident_time = data.get("accident_time")
        risky_time = data.get("risky_time")
        
        if isinstance(accident_time, str):
            try:
                accident_time = int(accident_time)
            except ValueError:
                accident_time = None
        
        if isinstance(risky_time, str):
            try:
                risky_time = int(risky_time)
            except ValueError:
                risky_time = None
        
        n_accident_samples = 3
        n_normal_samples = 2
        
        accident_samples = []
        normal_samples = []
        
        for idx in range(len(frames)):
            frame_num = int(frames[idx].stem)
            
            is_accident = False
            if accident_time and abs(frame_num - accident_time) <= 20:
                is_accident = True
            elif risky_time and abs(frame_num - risky_time) <= 30:
                is_accident = True
            
            if is_accident:
                accident_samples.append(idx)
            elif risky_time and frame_num < risky_time - 60:
                normal_samples.append(idx)
        
        if accident_samples:
            sampled_acc = random.sample(accident_samples, 
                                       min(n_accident_samples, len(accident_samples)))
            for idx in sampled_acc:
                samples.append({
                    "task": "binary_detection",
                    "subtask": "accident_classification",
                    "image_path": str(frames[idx]),
                    "user_prompt": "Is there an accident or traffic incident in this image? Answer: 'Accident detected' or 'Normal driving'.",
                    "label": "Accident detected",
                    "difficulty": "medium",
                    "metadata": {
                        "case_id": data["id"],
                        "dataset": data["dataset"],
                        "frame_num": int(frames[idx].stem),
                        "is_positive": True
                    }
                })
        
        if normal_samples:
            sampled_norm = random.sample(normal_samples,
                                        min(n_normal_samples, len(normal_samples)))
            for idx in sampled_norm:
                samples.append({
                    "task": "binary_detection",
                    "subtask": "accident_classification",
                    "image_path": str(frames[idx]),
                    "user_prompt": "Is there an accident or traffic incident in this image? Answer: 'Accident detected' or 'Normal driving'.",
                    "label": "Normal driving",
                    "difficulty": "medium",
                    "metadata": {
                        "case_id": data["id"],
                        "dataset": data["dataset"],
                        "frame_num": int(frames[idx].stem),
                        "is_positive": False
                    }
                })
    
    # 从无事故的cases采样
    for data in normal_cases:
        case_dir = Path(data["case_dir"])
        frames = sorted([f for f in case_dir.glob("*.jpg")])
        
        if len(frames) == 0:
            continue
        
        n_samples = random.randint(3, 4)
        sampled = random.sample(range(len(frames)), min(n_samples, len(frames)))
        
        for idx in sampled:
            samples.append({
                "task": "binary_detection",
                "subtask": "accident_classification",
                "image_path": str(frames[idx]),
                "user_prompt": "Is there an accident or traffic incident in this image? Answer: 'Accident detected' or 'Normal driving'.",
                "label": "Normal driving",
                "difficulty": "easy",
                "metadata": {
                    "case_id": data["id"],
                    "dataset": data["dataset"],
                    "frame_num": int(frames[idx].stem),
                    "is_positive": False
                }
            })
    
    print(f"[{split_name}] 任务2-二分类检测: {len(samples)} 样本")
    
    positive = sum(1 for s in samples if s["metadata"]["is_positive"])
    negative = len(samples) - positive
    print(f"  正样本: {positive}, 负样本: {negative}")
    
    # 统计数据集分布
    from collections import Counter
    dataset_dist = Counter(s["metadata"]["dataset"] for s in samples)
    print(f"  数据集分布:")
    for ds, count in dataset_dist.items():
        print(f"    {ds}: {count} 样本")
    
    return samples


# ============ 主流程 ============
def main():
    """主函数"""
    print("=" * 70)
    print("自适应Prompt预训练数据准备")
    print("策略: 根据annotation长度调整prompt，保持原始标注不变")
    print("=" * 70)
    
    # 加载数据
    all_data = load_all_annotations(include_dad=True)
    
    # 划分数据
    train_data, val_data, test_data = split_data(all_data)
    
    # 准备各任务
    results = {}
    
    for split_name, data_split in [("train", train_data), 
                                     ("val", val_data), 
                                     ("test", test_data)]:
        print(f"\n{'='*70}")
        print(f"处理 {split_name.upper()} Split")
        print("=" * 70)
        
        task1 = prepare_task1_scene_understanding(data_split, split_name)
        task2 = prepare_task2_binary_detection(data_split, split_name)
        task3 = prepare_task3_accident_description_adaptive(data_split, split_name)  # 自适应
        task4 = prepare_task4_sequence_adaptive(data_split, split_name)  # 自适应
        
        results[split_name] = {
            "task1_scene_understanding": task1,
            "task2_binary_detection": task2,
            "task3_accident_description": task3,
            "task4_sequence_prediction": task4,
            "total_cases": len(data_split)
        }
    
    # 保存
    print("\n" + "=" * 70)
    print("保存数据...")
    
    output_file = OUTPUT_DIR / "pretrain_data_adaptive.pkl"
    with open(output_file, "wb") as f:
        pickle.dump(results, f)
    print(f"✓ 保存到: {output_file}")
    
    # 统计
    summary = {}
    for split in ["train", "val", "test"]:
        summary[split] = {
            "cases": results[split]["total_cases"],
            "task1_scene": len(results[split]["task1_scene_understanding"]),
            "task2_binary": len(results[split]["task2_binary_detection"]),
            "task3_description": len(results[split]["task3_accident_description"]),
            "task4_sequence": len(results[split]["task4_sequence_prediction"]),
            "total_samples": (
                len(results[split]["task1_scene_understanding"]) +
                len(results[split]["task2_binary_detection"]) +
                len(results[split]["task3_accident_description"]) +
                len(results[split]["task4_sequence_prediction"])
            )
        }
    
    output_json = OUTPUT_DIR / "pretrain_summary_adaptive.json"
    with open(output_json, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"✓ 统计: {output_json}")
    
    # 打印总结
    print("\n" + "=" * 70)
    print("数据准备完成 - 统计:")
    print("=" * 70)
    for split in ["train", "val", "test"]:
        print(f"\n{split.upper()}: {summary[split]['cases']} cases")
        print(f"  任务1 (场景理解):   {summary[split]['task1_scene']}")
        print(f"  任务2 (二分类):     {summary[split]['task2_binary']}")
        print(f"  任务3 (事故描述):   {summary[split]['task3_description']}")
        print(f"  任务4 (序列预测):   {summary[split]['task4_sequence']}")
        print(f"  ───────────────────────────────")
        print(f"  总样本数:          {summary[split]['total_samples']}")
    
    print("\n✅ 完成!")
    
    # DAD数据集使用总结
    print("\n" + "=" * 70)
    print("DAD数据集使用总结:")
    print("=" * 70)
    print("✗ 任务1 (场景理解):   未使用 - 环境信息全为Unknown")
    print("✓ 任务2 (二分类):     已使用 - 提供大量normal driving样本")
    print("✓ 任务3 (事故描述):   已使用 - positive样本参与训练")
    print("✓ 任务4 (序列预测):   已使用 - 所有样本参与训练")
    print("\n策略: DAD数据集专注于二分类和序列理解任务")
    
    # 统计DAD在训练集中的占比
    dad_count_task2 = sum(
        1 for s in results["train"]["task2_binary_detection"]
        if s["metadata"]["dataset"] == "dad"
    )
    dad_count_task3 = sum(
        1 for s in results["train"]["task3_accident_description"]
        if s["metadata"]["dataset"] == "dad"
    )
    dad_count_task4 = sum(
        1 for s in results["train"]["task4_sequence_prediction"]
        if s["metadata"]["dataset"] == "dad"
    )
    
    total_dad = dad_count_task2 + dad_count_task3 + dad_count_task4
    total_samples = summary["train"]["total_samples"]
    
    print(f"\nDAD在训练集中的样本分布:")
    print(f"  任务2: {dad_count_task2} 样本")
    print(f"  任务3: {dad_count_task3} 样本")
    print(f"  任务4: {dad_count_task4} 样本")
    print(f"  总计: {total_dad} 样本 ({total_dad/total_samples*100:.1f}% of 训练集)")
    
    print("\n下一步:")
    print("1. 运行 python test_adaptive_data.py 验证数据")
    print("2. 使用 train_pretrain_adaptive.py 开始训练")


if __name__ == "__main__":
    main()