""" L1B 快照引擎 (Parquet 版本) 这是第一阶段 (Phase 1) 的核心 ETL 脚本。 它负责从 CS2 .dem 文件中提取 Tick 级别的快照，并将其保存为高压缩率的 Parquet 文件。 用法: python src/etl/extract_snapshots.py --demo_dir data/demos --output_dir data/processed 配置: 调整下方的参数以控制数据粒度 """ import os import argparse import pandas as pd import numpy as np from demoparser2 import DemoParser # 核心依赖 import logging import sys # ============================================================================== # ⚙️ 配置与调优参数 (可修改参数区) # ============================================================================== # [重要] 采样率 # 多久截取一次快照？ # 较低值 = 数据更多，精度更高，处理更慢。 # 较高值 = 数据更少，处理更快。 SNAPSHOT_INTERVAL_SECONDS = 2 # 👈 建议值: 1-5秒 (默认: 2s) # [重要] 回合过滤器 # 包含哪些回合？ # 'clutch_only': 仅保留发生残局 (<= 3v3) 的回合。 # 'all': 保留所有回合 (数据集会非常巨大)。 FILTER_MODE = 'clutch_only' # 👈 选项: 'all' | 'clutch_only' # [重要] 残局定义 # 什么样的局面算作“残局”？ MAX_PLAYERS_PER_TEAM = 2 # 👈 建议值: 2 (意味着 <= 2vX 或 Xv2) # 字段选择 (用于优化) # 仅从 demo 中提取这些字段以节省内存 WANTED_FIELDS = [ "game_time", # 游戏时间 "team_num", # 队伍编号 "player_name", # 玩家昵称 "steamid", # Steam ID "X", "Y", "Z", # 坐标位置 "view_X", "view_Y", # 视角角度 "health", # 生命值 "armor_value", # 护甲值 "has_defuser", # 是否有拆弹钳 "has_helmet", # 是否有头盔 "active_weapon_name", # 当前手持武器 "flash_duration", # 致盲持续时间 (是否被白) "is_alive", # 是否存活 "balance" # [NEW] 剩余金钱 (Correct field name) ] # ============================================================================== # 配置结束 # ============================================================================== logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') def is_clutch_situation(ct_alive, t_alive): """ 检查当前状态是否符合“残局”场景。 条件: 至少有一方队伍的存活人数 <= MAX_PLAYERS_PER_TEAM。 (例如: 2v5 对于剩2人的那队来说就是残局) """ if ct_alive == 0 or t_alive == 0: return False # 用户需求: "对面有几个人都无所谓，只要一方剩两个人" # 含义: 如果 CT <= N 或者 T <= N，即视为残局。 is_ct_clutch = (ct_alive <= MAX_PLAYERS_PER_TEAM) is_t_clutch = (t_alive <= MAX_PLAYERS_PER_TEAM) return is_ct_clutch or is_t_clutch def process_demo(demo_path, output_dir, delete_source=False): """ 解析单个 .dem 文件并将快照导出为 Parquet 格式。 """ demo_name = os.path.basename(demo_path).replace('.dem', '') output_path = os.path.join(output_dir, f"{demo_name}.parquet") if os.path.exists(output_path): logging.info(f"跳过 {demo_name}, 文件已存在。") if delete_source: try: os.remove(demo_path) logging.info(f"已删除源文件 (因为已存在处理结果): {demo_path}") except Exception as e: logging.warning(f"删除源文件失败: {e}") return logging.info(f"正在处理: {demo_name}") try: parser = DemoParser(demo_path) # 1. 解析元数据 (地图, 头部信息) header = parser.parse_header() map_name = header.get("map_name", "unknown") # 2. 提取事件 (回合开始/结束, 炸弹) 以识别回合边界 # [修复] 解析 round_start 事件以获取 round 信息，解决 KeyError: 'round' # [新增] 解析 round_end 事件以获取 round_winner 信息 # [新增] 解析 bomb 事件以获取 is_bomb_planted 和 bomb_site event_names = ["round_start", "round_end", "bomb_planted", "bomb_defused", "bomb_exploded"] parsed_events = parser.parse_events(event_names) round_df = None winner_df = None bomb_events = [] # parse_events 返回 [(event_name, df), ...] for event_name, event_data in parsed_events: if event_name == "round_start": round_df = event_data elif event_name == "round_end": winner_df = event_data elif event_name in ["bomb_planted", "bomb_defused", "bomb_exploded"]: # 统一处理炸弹事件 # bomb_planted 有 site 字段 # 其他可能没有，需要填充 temp_df = event_data.copy() temp_df['event_type'] = event_name if 'site' not in temp_df.columns: temp_df['site'] = 0 bomb_events.append(temp_df[['tick', 'event_type', 'site']]) # 3. 提取玩家状态 (繁重的工作) # 我们先获取所有 Tick 的数据，然后再进行过滤 df = parser.parse_ticks(WANTED_FIELDS) # [修复] 将 Round 信息合并到 DataFrame if round_df is not None and not round_df.empty: # 确保按 tick 排序 round_df = round_df.sort_values('tick') df = df.sort_values('tick') # 使用 merge_asof 将最近的 round_start 匹配给每个 tick # direction='backward' 意味着找 tick <= 当前tick 的最近一次 round_start df = pd.merge_asof(df, round_df[['tick', 'round']], on='tick', direction='backward') # 填充 NaN (比赛开始前的 tick) 为 0 df['round'] = df['round'].fillna(0).astype(int) else: logging.warning(f"在 {demo_name} 中未找到 round_start 事件，默认为第 1 回合") df['round'] = 1 # [新增] 将 Winner 信息合并到 DataFrame if winner_df is not None and not winner_df.empty: # winner_df 包含 'round' 和 'winner' # 这里的 'round' 是结束的回合号。 # 我们直接将 winner 映射到 df 中的 round 列 # 清洗 winner 数据 (T -> 0, CT -> 1) # 注意: demoparser2 返回的 winner 可能是 int (2/3) 也可能是 str ('T'/'CT') # 我们先统一转为字符串处理 winner_map = df[['round']].copy().drop_duplicates() # 建立 round -> winner 字典 # 过滤无效的 winner valid_winners = winner_df.dropna(subset=['winner']) round_winner_dict = {} for _, row in valid_winners.iterrows(): r = row['round'] w = row['winner'] if w == 'T' or w == 2: round_winner_dict[r] = 0 # T wins elif w == 'CT' or w == 3: round_winner_dict[r] = 1 # CT wins # 映射到主 DataFrame df['round_winner'] = df['round'].map(round_winner_dict) # 移除没有结果的回合 (例如 warmup 或未结束的回合) # df = df.dropna(subset=['round_winner']) # 暂时保留，由后续步骤决定是否丢弃 else: logging.warning(f"在 {demo_name} 中未找到 round_end 事件，无法标记胜者") df['round_winner'] = None # [新增] 合并炸弹状态 (is_bomb_planted) if bomb_events: bomb_df = pd.concat(bomb_events).sort_values('tick') # 逻辑： # bomb_planted -> is_planted=1, site=X # bomb_defused/exploded -> is_planted=0, site=0 # round_start/end -> 也可以作为重置点 (state=0)，但我们没有把它们放入 bomb_events # 我们假设 round_start 时炸弹肯定没下，但 merge_asof 会延续上一个状态 # 所以我们需要把 round_start 也加入作为重置事件 if round_df is not None: reset_df = round_df[['tick']].copy() reset_df['event_type'] = 'reset' reset_df['site'] = 0 bomb_df = pd.concat([bomb_df, reset_df]).sort_values('tick') # 计算状态 # 1 = Planted, 0 = Not Planted bomb_df['is_bomb_planted'] = bomb_df['event_type'].apply(lambda x: 1 if x == 'bomb_planted' else 0) # site 已经在 bomb_planted 事件中有值，其他为 0 # 使用 merge_asof 传播状态 # 注意：bomb_df 可能有同一 tick 多个事件，merge_asof 取最后一个 # 所以我们要确保排序正确 (reset 应该在 planted 之前？不，reset 是 round_start，肯定在 planted 之前) # 只需要 tick, is_bomb_planted, site state_df = bomb_df[['tick', 'is_bomb_planted', 'site']].copy() df = pd.merge_asof(df, state_df, on='tick', direction='backward') # 填充 NaN 为 0 (未下包) df['is_bomb_planted'] = df['is_bomb_planted'].fillna(0).astype(int) df['site'] = df['site'].fillna(0).astype(int) else: df['is_bomb_planted'] = 0 df['site'] = 0 # 4. 数据清洗与优化 # 将 team_num 转换为 int (CT=3, T=2) df['team_num'] = df['team_num'].fillna(0).astype(int) # 5. 应用采样间隔过滤器 # 我们不需要每一帧 (128/s)，而是每 N 秒取一帧 # 近似计算: tick_rate 大约是 64 或 128。 # 我们使用 'game_time' 来过滤。 df['time_bin'] = (df['game_time'] // SNAPSHOT_INTERVAL_SECONDS).astype(int) # [修复] 采样逻辑优化：找出每个 (round, time_bin) 的起始 tick，保留该 tick 的所有玩家数据 # 旧逻辑 groupby().first() 会丢失其他玩家数据 bin_start_ticks = df.groupby(['round', 'time_bin'])['tick'].min() selected_ticks = bin_start_ticks.values # 提取快照 (包含被选中 tick 的所有玩家行) snapshot_df = df[df['tick'].isin(selected_ticks)].copy() # 6. 应用残局逻辑过滤器 if FILTER_MODE == 'clutch_only': # 我们需要计算每一帧各队的存活人数 # snapshot_df 已经是采样后的数据 (每个 tick 包含所有玩家) # 高效的存活人数计算: alive_counts = snapshot_df[snapshot_df['is_alive'] == True].groupby(['round', 'time_bin', 'team_num']).size().unstack(fill_value=0) # 确保列存在 (2=T, 3=CT) if 2 not in alive_counts.columns: alive_counts[2] = 0 if 3 not in alive_counts.columns: alive_counts[3] = 0 # 过滤出满足残局条件的帧 # alive_counts 的索引是 (round, time_bin) clutch_mask = [is_clutch_situation(row[3], row[2]) for index, row in alive_counts.iterrows()] valid_indices = alive_counts[clutch_mask].index # 过滤主 DataFrame # 构建一个复合键用于快速过滤 snapshot_df['frame_id'] = list(zip(snapshot_df['round'], snapshot_df['time_bin'])) valid_frame_ids = set(valid_indices) snapshot_df = snapshot_df[snapshot_df['frame_id'].isin(valid_frame_ids)].copy() snapshot_df.drop(columns=['frame_id'], inplace=True) if snapshot_df.empty: logging.warning(f"在 {demo_name} 中未找到有效快照 (过滤器: {FILTER_MODE})") return # 7. 添加元数据 snapshot_df['match_id'] = demo_name snapshot_df['map_name'] = map_name # [优化] 数据类型降维与压缩 # 这一步能显著减少内存占用和文件体积 # Float64 -> Float32 float_cols = ['X', 'Y', 'Z', 'view_X', 'view_Y', 'game_time', 'flash_duration'] for col in float_cols: if col in snapshot_df.columns: snapshot_df[col] = snapshot_df[col].astype('float32') # Int64 -> Int8/Int16 # team_num: 2 or 3 -> int8 snapshot_df['team_num'] = snapshot_df['team_num'].astype('int8') # health, armor: 0-100 -> int16 (uint8 也可以但 pandas 对 uint 支持有时候有坑) for col in ['health', 'armor_value', 'balance', 'site']: if col in snapshot_df.columns: snapshot_df[col] = snapshot_df[col].fillna(0).astype('int16') # round, tick: int32 (enough for millions) snapshot_df['round'] = snapshot_df['round'].astype('int16') snapshot_df['tick'] = snapshot_df['tick'].astype('int32') # Booleans -> int8 or bool bool_cols = ['is_alive', 'has_defuser', 'has_helmet', 'is_bomb_planted'] for col in bool_cols: if col in snapshot_df.columns: snapshot_df[col] = snapshot_df[col].astype('int8') # 0/1 is better for ML sometimes # Drop redundant columns if 'time_bin' in snapshot_df.columns: snapshot_df.drop(columns=['time_bin'], inplace=True) # 8. 保存为 Parquet (L1B 层) # 使用 zstd 压缩算法，通常比 snappy 压缩率高 30-50% snapshot_df.to_parquet(output_path, index=False, compression='zstd') logging.info(f"已保存 {len(snapshot_df)} 条快照到 {output_path} (压缩模式: ZSTD)") # [NEW] 删除源文件逻辑 if delete_source: try: os.remove(demo_path) logging.info(f"处理成功，已删除源文件: {demo_path}") except Exception as e: logging.warning(f"删除源文件失败: {e}") except Exception as e: logging.error(f"处理失败 {demo_name}: {str(e)}") # 如果是 Source 1 错误，给予明确提示 if "Source1" in str(e): logging.error("❌ 这是一个 CS:GO (Source 1) 的 Demo，本系统仅支持 CS2 (Source 2) Demo。") sys.exit(1) def main(): parser = argparse.ArgumentParser(description="L1B 快照引擎") parser.add_argument('--demo_dir', type=str, default='data/demos', help='输入 .dem 文件的目录') parser.add_argument('--file', type=str, help='处理单个 .dem 文件 (如果指定此参数，将忽略 --demo_dir)') parser.add_argument('--output_dir', type=str, default='data/processed', help='输出 .parquet 文件的目录') parser.add_argument('--delete-source', action='store_true', help='处理成功后删除源文件') args = parser.parse_args() if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) # 模式 1: 单文件处理 if args.file: if not os.path.exists(args.file): logging.error(f"文件不存在: {args.file}") return if not args.file.endswith('.dem'): logging.error(f"无效的文件扩展名: {args.file}") return process_demo(args.file, args.output_dir, delete_source=args.delete_source) return # 模式 2: 目录批处理 if not os.path.exists(args.demo_dir): logging.warning(f"目录不存在: {args.demo_dir}") return demo_files = [os.path.join(args.demo_dir, f) for f in os.listdir(args.demo_dir) if f.endswith('.dem')] if not demo_files: logging.warning(f"在 {args.demo_dir} 中未找到 .dem 文件。请添加 demo 文件。") return for demo_path in demo_files: process_demo(demo_path, args.output_dir, delete_source=args.delete_source) if __name__ == "__main__": main()