#!/usr/bin/env python3 """ 自适应Prompt的数据集加载器 使用数据中的user_prompt字段，而不是固定的prompt模板 """ import pickle from pathlib import Path from typing import Dict, List, Optional, Tuple import torch from torch.utils.data import Dataset from PIL import Image import random class AdaptivePretrainDataset(Dataset): """ 自适应Prompt预训练数据集 每个样本都有自己的user_prompt，根据annotation长度定制 Args: data_file: pretrain_data_adaptive.pkl路径 split: 'train', 'val', 或 'test' tasks: 任务列表 curriculum_stage: 0=easy, 1=medium, 2=hard, 3=all use_system_prompt: 是否使用system prompt """ # System prompts (任务级别) SYSTEM_PROMPTS = { "scene_understanding": "You are an expert driving scene analyzer. Describe the environment accurately.", "binary_detection": "You are a traffic safety AI. Detect abnormal driving situations.", "accident_description": "You are an accident analysis AI. Answer based on the question asked.", # 更通用 "sequence_prediction": "You are a temporal driving AI. Analyze video sequences for accident prediction." } def __init__( self, data_file: str, split: str = "train", tasks: List[str] = None, curriculum_stage: int = 3, use_system_prompt: bool = True ): self.split = split self.tasks = tasks or ["task1", "task2", "task3", "task4"] self.curriculum_stage = curriculum_stage self.use_system_prompt = use_system_prompt # 加载数据 with open(data_file, "rb") as f: all_data = pickle.load(f) split_data = all_data[split] # 收集样本 self.samples = [] task_map = { "task1": "task1_scene_understanding", "task2": "task2_binary_detection", "task3": "task3_accident_description", "task4": "task4_sequence_prediction" } for task in self.tasks: if task in task_map: task_samples = split_data.get(task_map[task], []) # Curriculum filtering if curriculum_stage < 3: difficulty_map = {0: "easy", 1: "medium", 2: "hard"} target_difficulty = difficulty_map[curriculum_stage] task_samples = [ s for s in task_samples if s.get("difficulty", "easy") == target_difficulty ] self.samples.extend(task_samples) # Shuffle if split == "train": random.shuffle(self.samples) print(f"{'='*70}") print(f"数据集加载: {split}") print(f"Curriculum Stage: {curriculum_stage} ({['easy', 'medium', 'hard', 'all'][curriculum_stage]})") print(f"任务: {tasks}") print(f"样本数: {len(self.samples)}") # 统计 from collections import Counter task_dist = Counter(s["task"] for s in self.samples) print(f"\n任务分布:") for task, count in task_dist.items(): print(f" {task}: {count}") # 统计短/长标注 if curriculum_stage == 3 and ("task3" in self.tasks or "task4" in self.tasks): short_count = sum( 1 for s in self.samples if s["task"] in ["accident_description", "sequence_prediction"] and s["metadata"].get("is_short_annotation", False) ) detailed_count = sum( 1 for s in self.samples if s["task"] in ["accident_description", "sequence_prediction"] and not s["metadata"].get("is_short_annotation", False) ) if short_count + detailed_count > 0: print(f"\nAnnotation分布 (任务3&4):") print(f" 短标注 (<20字符): {short_count}") print(f" 详细标注 (>=20字符): {detailed_count}") # 难度分布 if curriculum_stage == 3: diff_dist = Counter(s.get("difficulty", "unknown") for s in self.samples) print(f"\n难度分布:") for diff, count in diff_dist.items(): print(f" {diff}: {count}") print("=" * 70) def __len__(self): return len(self.samples) def __getitem__(self, idx): sample = self.samples[idx] task_type = sample["task"] # 获取system prompt (任务级别) system_prompt = self.SYSTEM_PROMPTS[task_type] if self.use_system_prompt else "" # 使用样本中的user_prompt (自适应) user_prompt = sample.get("user_prompt", "") if task_type in ["scene_understanding", "binary_detection", "accident_description"]: # 单帧任务 image = Image.open(sample["image_path"]).convert("RGB") return { "task": task_type, "subtask": sample.get("subtask", task_type), "image": image, "system_prompt": system_prompt, "user_prompt": user_prompt, # 自适应prompt "label": sample["label"], "difficulty": sample.get("difficulty", "unknown"), "metadata": sample["metadata"] } elif task_type == "sequence_prediction": # 序列任务 images = [] for img_path in sample["image_sequence"]: img = Image.open(img_path).convert("RGB") images.append(img) return { "task": task_type, "subtask": sample.get("subtask", task_type), "image_sequence": images, "system_prompt": system_prompt, "user_prompt": user_prompt, # 自适应prompt "label": sample["label"], "difficulty": sample.get("difficulty", "unknown"), "metadata": sample["metadata"] } else: raise ValueError(f"未知任务类型: {task_type}") def collate_fn_adaptive(batch): """ 自适应collate函数 每个样本有自己的user_prompt """ single_frame_batch = [] sequence_batch = [] for item in batch: if item["task"] in ["scene_understanding", "binary_detection", "accident_description"]: single_frame_batch.append(item) elif item["task"] == "sequence_prediction": sequence_batch.append(item) result = {} # 单帧任务 if single_frame_batch: result["single_frame"] = { "task": [x["task"] for x in single_frame_batch], "subtask": [x["subtask"] for x in single_frame_batch], "images": [x["image"] for x in single_frame_batch], "system_prompts": [x["system_prompt"] for x in single_frame_batch], "user_prompts": [x["user_prompt"] for x in single_frame_batch], # 每个样本不同 "labels": [x["label"] for x in single_frame_batch], "difficulties": [x["difficulty"] for x in single_frame_batch], "metadata": [x["metadata"] for x in single_frame_batch] } # 序列任务 if sequence_batch: result["sequence"] = { "task": [x["task"] for x in sequence_batch], "subtask": [x["subtask"] for x in sequence_batch], "image_sequences": [x["image_sequence"] for x in sequence_batch], "system_prompts": [x["system_prompt"] for x in sequence_batch], "user_prompts": [x["user_prompt"] for x in sequence_batch], # 每个样本不同 "labels": [x["label"] for x in sequence_batch], "difficulties": [x["difficulty"] for x in sequence_batch], "metadata": [x["metadata"] for x in sequence_batch] } return result # ============ 测试代码 ============ if __name__ == "__main__": from torch.utils.data import DataLoader data_file = "PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data_adaptive.pkl" print("\n" + "=" * 70) print("测试自适应Prompt数据集") print("=" * 70) # 创建数据集 dataset = AdaptivePretrainDataset( data_file=data_file, split="train", tasks=["task1", "task2", "task3", "task4"], curriculum_stage=3 ) loader = DataLoader( dataset, batch_size=4, shuffle=False, num_workers=0, collate_fn=collate_fn_adaptive ) # 测试一个batch batch = next(iter(loader)) print("\n" + "=" * 70) print("Batch示例") print("=" * 70) if "single_frame" in batch: sf = batch["single_frame"] print(f"\n单帧任务: {len(sf['images'])} 样本") for i in range(len(sf['task'])): print(f"\n样本 {i+1}:") print(f" 任务: {sf['task'][i]}") print(f" 难度: {sf['difficulties'][i]}") print(f" System: {sf['system_prompts'][i][:60]}...") print(f" User Prompt: {sf['user_prompts'][i]}") # 注意每个都不同 print(f" Label: {sf['labels'][i][:60]}...") # 如果是事故描述任务，显示annotation长度 if sf['task'][i] == 'accident_description': is_short = sf['metadata'][i].get('is_short_annotation', False) anno_len = sf['metadata'][i].get('annotation_length', 0) print(f" Annotation: {'短' if is_short else '详细'} ({anno_len}字符)") if "sequence" in batch: seq = batch["sequence"] print(f"\n序列任务: {len(seq['image_sequences'])} 样本") for i in range(len(seq['task'])): print(f"\n样本 {i+1}:") print(f" 序列长度: {len(seq['image_sequences'][i])}") print(f" 难度: {seq['difficulties'][i]}") print(f" User Prompt: {seq['user_prompts'][i]}") # 注意每个都不同 print(f" Label: {seq['labels'][i][:60]}...") is_short = seq['metadata'][i].get('is_short_annotation', False) anno_len = seq['metadata'][i].get('annotation_length', 0) print(f" Annotation: {'短' if is_short else '详细'} ({anno_len}字符)") print("\n✅ 数据集测试完成!")