training/PRETRAIN/prepare_pretrain_data.py

#!/usr/bin/env python3
"""
预训练数据准备脚本
生成三个任务的训练数据：
1. 环境描述（天气、道路、光照）
2. 单帧事故判断
3. 序列事故预测和描述
"""

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

random.seed(42)

# ============ 配置 ============
PRETRAIN_ROOT = Path("PROJECT_ROOT/data/dataset/pretrain")
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


# ============ 数据加载 ============
def load_all_annotations():
    """加载所有annotation.json"""
    all_data = []
    
    # 加载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)
                all_data.append(data)
    
    # 加载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)
    
    print(f"加载 {len(all_data)} 案例")
    return all_data


def split_data(all_data):
    """划分train/val/test"""
    random.shuffle(all_data)
    n = len(all_data)
    n_train = int(n * TRAIN_RATIO)
    n_val = int(n * VAL_RATIO)
    
    train_data = all_data[:n_train]
    val_data = all_data[n_train:n_train + n_val]
    test_data = all_data[n_train + n_val:]
    
    print(f"训练: {len(train_data)}, 验证: {len(val_data)}, 测试: {len(test_data)}")
    return train_data, val_data, test_data


# ============ 任务1: 环境描述 ============
def prepare_task1_environment(data_split, split_name):
    """单帧环境描述: weather, road_type, light"""
    samples = []
    
    for data in data_split:
        case_dir = Path(data["case_dir"])
        frames = sorted([f for f in case_dir.glob("*.jpg")])
        
        if len(frames) == 0:
            continue
        
        # 每视频采3-5帧
        n_samples = random.randint(3, 5)
        sampled = random.sample(frames, min(n_samples, len(frames)))
        
        for frame_path in sampled:
            if data["dataset"] == "nexar":
                weather = data.get("weather", "Unknown")
                road = data.get("road_type", "Unknown")
                light = data.get("light_conditions", "Unknown")
            else:
                weather = data.get("weather", "Unknown")
                road = data.get("road_type", "Unknown")
                light = data.get("time_of_day", "Unknown")
            
            label = f"Weather: {weather}, Road: {road}, Light: {light}"
            
            samples.append({
                "task": "environment",
                "image_path": str(frame_path),
                "label": label,
                "metadata": {
                    "case_id": data["id"],
                    "dataset": data["dataset"]
                }
            })
    
    print(f"[{split_name}] 任务1: {len(samples)} 样本")
    return samples


# ============ 任务2: 单帧事故判断 ============
def prepare_task2_accident(data_split, split_name):
    """单帧判断是否事故"""
    samples = []
    
    for data in data_split:
        case_dir = Path(data["case_dir"])
        frames = sorted([f for f in case_dir.glob("*.jpg")])
        
        if len(frames) == 0:
            continue
        
        has_accident = data.get("accident", False)
        if isinstance(has_accident, str):
            has_accident = has_accident.lower() == "true"
        
        accident_time = data.get("accident_time")
        
        # 转换字符串为int
        if isinstance(accident_time, str):
            try:
                accident_time = int(accident_time)
            except ValueError:
                accident_time = None
        
        # 采3-5帧
        n_samples = random.randint(3, 5)
        sampled_idx = random.sample(range(len(frames)), min(n_samples, len(frames)))
        
        for idx in sampled_idx:
            frame_path = frames[idx]
            frame_num = int(frame_path.stem)
            
            # 事故前后1秒内为事故帧
            is_accident_frame = False
            if has_accident and accident_time is not None and accident_time > 0:
                if abs(frame_num - accident_time) <= 20:  # 20fps
                    is_accident_frame = True
            
            label = "Yes" if is_accident_frame else "No"
            
            samples.append({
                "task": "accident_detection",
                "image_path": str(frame_path),
                "label": label,
                "metadata": {
                    "case_id": data["id"],
                    "dataset": data["dataset"],
                    "frame_num": frame_num
                }
            })
    
    print(f"[{split_name}] 任务2: {len(samples)} 样本")
    return samples


# ============ 任务3: 序列预测 ============
def prepare_task3_sequence(data_split, split_name):
    """序列判断事故+描述"""
    samples = []
    
    for data in data_split:
        case_dir = Path(data["case_dir"])
        frames = sorted([f for f in case_dir.glob("*.jpg")])
        
        if len(frames) < 8:  # 至少需要8帧才能采样
            continue
        
        # 处理risky_time
        risky_time = data.get("risky_time")
        
        # 转换字符串为int
        if isinstance(risky_time, str):
            try:
                risky_time = int(risky_time)
            except ValueError:
                risky_time = None
        
        # 判断是否有事故
        has_accident = data.get("accident", False)
        if isinstance(has_accident, str):
            has_accident = has_accident.lower() == "true"
        
        accident_type = data.get("accident_type", "No accident")
        if accident_type is None or accident_type == "null":
            accident_type = "No accident"
        
        # 确定采样起始点
        if risky_time is not None and risky_time > 0 and has_accident:
            # 有事故且有risky_time: 从risky_time前0.2秒开始
            start_frame = max(0, risky_time - 8)
        else:
            # 无事故或无risky_time: 随机选择起始点
            # 确保至少能采样到2帧
            max_start = len(frames) - 16  # 至少留8帧(2个采样点)
            if max_start <= 0:
                start_frame = 0
            else:
                start_frame = random.randint(0, max_start)
        
        # 每4帧选1帧
        # sequence = []
        # for i in range(start_frame, len(frames), 4):
        #     if i < len(frames):
        #         sequence.append(str(frames[i]))
        STRIDE = 8           # 20fps → 8 帧 = 0.4s
        T_MAX  = 16          # 建议上限（可改 16）；不改变任务，只控显存

        # 先按 0.4s 间隔取全程
        seq_full = list(range(start_frame, len(frames), STRIDE))
        seq_full = [str(frames[i]) for i in seq_full if i < len(frames)]

        # 再把超长的均匀采到 T_MAX
        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
        
        # 至少需要2帧
        if len(sequence) < 2:
            continue
        
        # 构造标签
        accident_label = "Yes" if has_accident else "No"
        label = f"Accident: {accident_label}. Description: {accident_type}"
        
        samples.append({
            "task": "sequence_prediction",
            "image_sequence": sequence,
            "label": label,
            "metadata": {
                "case_id": data["id"],
                "dataset": data["dataset"],
                "sequence_length": len(sequence),
                "has_accident": has_accident,
                "start_frame": start_frame
            }
        })
    
    print(f"[{split_name}] 任务3: {len(samples)} 样本")
    return samples


# ============ 主流程 ============
def main():
    print("=" * 50)
    print("准备预训练数据")
    print("=" * 50)
    
    # 加载数据
    all_data = load_all_annotations()
    
    # 划分数据
    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处理 {split_name}...")
        
        task1 = prepare_task1_environment(data_split, split_name)
        task2 = prepare_task2_accident(data_split, split_name)
        task3 = prepare_task3_sequence(data_split, split_name)
        
        results[split_name] = {
            "task1_environment": task1,
            "task2_accident_detection": task2,
            "task3_sequence_prediction": task3,
            "total_cases": len(data_split)
        }
    
    # 保存
    print("\n" + "=" * 50)
    print("保存数据...")
    
    output_file = OUTPUT_DIR / "pretrain_data.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": len(results[split]["task1_environment"]),
            "task2": len(results[split]["task2_accident_detection"]),
            "task3": len(results[split]["task3_sequence_prediction"])
        }
    
    output_json = OUTPUT_DIR / "pretrain_summary.json"
    with open(output_json, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"✓ 统计: {output_json}")
    
    print("\n" + "=" * 50)
    print("统计:")
    for split in ["train", "val", "test"]:
        print(f"\n{split.upper()}: {summary[split]['cases']} 案例")
        print(f"  任务1: {summary[split]['task1']}")
        print(f"  任务2: {summary[split]['task2']}")
        print(f"  任务3: {summary[split]['task3']}")
    
    print("\n✅ 完成!")


if __name__ == "__main__":
    main()