{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## 导入依赖与配置路径","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport json\nimport gc\nimport os\nimport glob\nimport joblib\nimport itertools\nimport warnings\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\n# 设置显示选项\npd.set_option('display.max_columns', None)\npd.set_option('display.width', 1000)\nwarnings.filterwarnings('ignore')\n\n# ================= 配置类 (CFG) =================\nclass CFG:\n    # 比赛数据的标准路径 (Kaggle 环境)\n    train_path = '/kaggle/input/MABe-mouse-behavior-detection/train.csv'\n    test_path = '/kaggle/input/MABe-mouse-behavior-detection/test.csv'\n    \n    # 追踪数据 (Tracking Data) 的路径\n    train_tracking_path = '/kaggle/input/MABe-mouse-behavior-detection/train_tracking'\n    test_tracking_path = '/kaggle/input/MABe-mouse-behavior-detection/test_tracking'\n    \n    # 标注数据 (Annotations) 的路径\n    train_annotation_path = '/kaggle/input/MABe-mouse-behavior-detection/train_annotation'\n    \n    model_name = 'two_stage_model'  # 模型保存文件夹\n    seed = 42\n\nprint(\"Step 1: 环境配置完成。\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T05:53:24.836676Z","iopub.execute_input":"2025-12-25T05:53:24.836930Z","iopub.status.idle":"2025-12-25T05:53:29.069095Z","shell.execute_reply.started":"2025-12-25T05:53:24.836903Z","shell.execute_reply":"2025-12-25T05:53:29.068270Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 加载元数据并查看","metadata":{}},{"cell_type":"code","source":"# 1. 读取元数据 CSV\nprint(\"Step 2: 正在读取 CSV 文件...\")\n\n# 检查文件是否存在，防止路径错误\nif not os.path.exists(CFG.train_path):\n    print(f\"错误: 找不到文件 {CFG.train_path}，请检查路径配置。\")\nelse:\n    train = pd.read_csv(CFG.train_path)\n    test = pd.read_csv(CFG.test_path)\n\n    # 2. 基础预处理\n    # 计算每一行有多少只老鼠 (总数4 - 缺失的只数)\n    train['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\n\n    # 过滤掉 'MABe22' 实验室的数据 (通常不用于训练)\n    train = train.query(\"~lab_id.str.startswith('MABe22_')\").reset_index(drop=True)\n\n    # 3. 提取所有出现的“身体部位追踪配置”\n    body_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\n    # 定义需要丢弃的冗余身体部位\n    drop_body_parts = [\n        'headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', \n        'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft', 'headpiece_topfrontright', \n        'spine_1', 'spine_2', 'tail_middle_1', 'tail_middle_2', 'tail_midpoint'\n    ]\n\n    print(f\"训练集形状 (过滤后): {train.shape}\")\n    print(f\"测试集形状: {test.shape}\")\n    print(\"\\n[预览] 训练集前 3 行:\")\n    display(train.head(3))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T05:53:29.070578Z","iopub.execute_input":"2025-12-25T05:53:29.070992Z","iopub.status.idle":"2025-12-25T05:53:29.265034Z","shell.execute_reply.started":"2025-12-25T05:53:29.070967Z","shell.execute_reply":"2025-12-25T05:53:29.264238Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 定义数学与物理工具函数","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport json\n\ndef robustify(submission, dataset, traintest, traintest_directory=None):\n    \"\"\"\n    清洗提交文件，处理重叠帧，并为缺失视频填充默认值\n    \"\"\"\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n    \n    # 1. 确保 start/stop frame 是整数\n    # 处理可能的字符串或浮点数，并将 NaN 填充为 0\n    submission['start_frame'] = pd.to_numeric(submission['start_frame'], errors='coerce').fillna(0).astype(int)\n    submission['stop_frame'] = pd.to_numeric(submission['stop_frame'], errors='coerce').fillna(0).astype(int)\n    \n    # 2. 过滤无效帧 (Start >= Stop)\n    old_len = len(submission)\n    submission = submission[submission.start_frame < submission.stop_frame]\n    if len(submission) != old_len:\n        print(f\"Robstify: Dropped {old_len - len(submission)} rows with start >= stop\")\n\n    # 3. 处理重叠帧 (同个视频同个老鼠不能同时做两个动作)\n    # 按 video_id + agent_id + target_id 分组\n    group_list = []\n    for _, group in submission.groupby(['video_id', 'agent_id', 'target_id']):\n        group = group.sort_values('start_frame')\n        mask = np.ones(len(group), dtype=bool)\n        last_stop_frame = -1\n        \n        for i, (_, row) in enumerate(group.iterrows()):\n            # 如果当前动作的开始时间 小于 上一个动作的结束时间，说明重叠了\n            if row['start_frame'] < last_stop_frame:\n                mask[i] = False # 丢弃重叠的后一个动作\n            else:\n                last_stop_frame = row['stop_frame']\n        group_list.append(group[mask])\n        \n    submission = pd.concat(group_list) if group_list else submission\n\n    # 4. 兜底填充 (处理没有任何预测结果的视频)\n    # Kaggle 评分系统要求每个视频至少有一条预测，否则会报错\n    s_list = []\n    \n    # 获取已经有预测的视频列表\n    predicted_videos = set(submission.video_id.unique())\n    \n    for idx, row in dataset.iterrows():\n        video_id = row['video_id']\n        lab_id = row['lab_id']\n        \n        # 如果是 MABe22 数据或已有预测，跳过\n        if lab_id.startswith('MABe22') or video_id in predicted_videos:\n            continue\n            \n        # 如果 behaviors_labeled 为空，跳过\n        if not isinstance(row.behaviors_labeled, str):\n            continue\n\n        # print(f\"Warning: Video {video_id} has no predictions. Filling with dummy data.\")\n        \n        # 读取该视频的元数据，获取最大帧数\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        if not os.path.exists(path): continue\n            \n        try:\n            vid = pd.read_parquet(path)\n            start_frame = vid.video_frame.min()\n            stop_frame = vid.video_frame.max() + 1\n        except:\n            # 如果读取失败，随便给个默认值\n            start_frame = 0\n            stop_frame = 100\n    \n        # 解析该视频需要预测哪些行为\n        try:\n            vid_behaviors = json.loads(row['behaviors_labeled'])\n            vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in vid_behaviors}))\n            vid_behaviors = [b.split(',') for b in vid_behaviors] # [agent, target, action]\n            vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action'])\n        except:\n            continue\n    \n        # 为每对 agent-target 均匀生成一段 dummy 预测\n        # 比如把整个视频时长平分给所有可能的动作\n        for (agent, target), actions in vid_behaviors.groupby(['agent', 'target']):\n            if len(actions) == 0: continue\n            \n            batch_length = int(np.ceil((stop_frame - start_frame) / len(actions)))\n            for i, (_, action_row) in enumerate(actions.iterrows()):\n                batch_start = start_frame + i * batch_length\n                batch_stop = min(batch_start + batch_length, stop_frame)\n                s_list.append((video_id, agent, target, action_row['action'], batch_start, batch_stop))\n\n    if len(s_list) > 0:\n        print(f\"Robustify: Filled {len(s_list)} empty segments for missing videos.\")\n        fill_df = pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n        submission = pd.concat([submission, fill_df], ignore_index=True)\n\n    submission = submission.reset_index(drop=True)\n    return submission\n\nprint(\"Robustify 函数定义完成。\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T05:53:29.265979Z","iopub.execute_input":"2025-12-25T05:53:29.266216Z","iopub.status.idle":"2025-12-25T05:53:29.280235Z","shell.execute_reply.started":"2025-12-25T05:53:29.266197Z","shell.execute_reply":"2025-12-25T05:53:29.279502Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ================= 基础工具函数 =================\n\ndef safe_rolling(series, window, func, min_periods=None):\n    \"\"\"安全的滚动窗口计算，处理 NaN\"\"\"\n    if min_periods is None:\n        min_periods = max(1, window // 4)\n    return series.rolling(window, min_periods=min_periods, center=True).apply(func, raw=True)\n\ndef _scale(n_frames_at_30fps, fps, ref=30.0):\n    \"\"\"\n    FPS 自适应缩放\n    [cite_start]将 '30FPS下的帧数' 转换为 '当前视频FPS下的帧数' [cite: 38]\n    \"\"\"\n    return max(1, int(round(n_frames_at_30fps * float(fps) / ref)))\n\ndef _scale_signed(n_frames_at_30fps, fps, ref=30.0):\n    \"\"\"带符号的 FPS 缩放 (用于 shift 操作)\"\"\"\n    if n_frames_at_30fps == 0: return 0\n    s = 1 if n_frames_at_30fps > 0 else -1\n    mag = max(1, int(round(abs(n_frames_at_30fps) * float(fps) / ref)))\n    return s * mag\n\ndef _fps_from_meta(meta_df, fallback_lookup, default_fps=30.0):\n    \"\"\"从元数据中获取 FPS\"\"\"\n    if 'frames_per_second' in meta_df.columns and pd.notnull(meta_df['frames_per_second']).any():\n        return float(meta_df['frames_per_second'].iloc[0])\n    vid = meta_df['video_id'].iloc[0]\n    return float(fallback_lookup.get(vid, default_fps))\n\n# ================= 几何计算函数 =================\n\ndef _angle_between(v1, v2):\n    \"\"\"计算两个向量之间的夹角 (弧度)\"\"\"\n    dot = v1['x'] * v2['x'] + v1['y'] * v2['y']\n    norm = np.sqrt(v1['x']**2 + v1['y']**2) * np.sqrt(v2['x']**2 + v2['y']**2)\n    # 加上 1e-6 防止除以 0\n    cos_angle = dot / (norm + 1e-6)\n    # clip 防止数值误差导致超出 [-1, 1]\n    return np.arccos(np.clip(cos_angle, -1.0, 1.0))\n\ndef _triangle_area(p1, p2, p3):\n    \"\"\"计算三角形面积 (Shoelace公式)，用于衡量身体形态\"\"\"\n    return 0.5 * np.abs(\n        p1['x'] * (p2['y'] - p3['y']) +\n        p2['x'] * (p3['y'] - p1['y']) +\n        p3['x'] * (p1['y'] - p2['y'])\n    )\n\nprint(\"Step 3: 工具函数定义完成。\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T05:53:29.281268Z","iopub.execute_input":"2025-12-25T05:53:29.281570Z","iopub.status.idle":"2025-12-25T05:53:29.304598Z","shell.execute_reply.started":"2025-12-25T05:53:29.281542Z","shell.execute_reply":"2025-12-25T05:53:29.303899Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 核心特征提取与数据生成","metadata":{}},{"cell_type":"code","source":"# ================= 辅助特征函数 =================\n\ndef add_arena_features(X, cx, cy, arena_dims=None):\n    \"\"\"添加场地相关特征 (趋触性/Wall-hugging)\"\"\"\n    # 如果没有场地大小信息，返回原特征\n    if arena_dims is None or np.isnan(arena_dims).any():\n        return X\n    \n    w, h = arena_dims\n    # 距离场地中心的距离\n    dist_center = np.sqrt((cx - w/2)**2 + (cy - h/2)**2)\n    X['dist_center'] = dist_center\n    \n    # 距离最近墙壁的距离\n    dist_wall_x = np.minimum(cx, w - cx)\n    dist_wall_y = np.minimum(cy, h - cy)\n    X['dist_wall'] = np.minimum(dist_wall_x, dist_wall_y)\n    \n    return X\n\n# 为了保持代码独立性，把之前用到的 add_* 函数也在这里快速定义一下\n# (在实际完整脚本中，这些通常作为独立函数存在，这里为了方便运行整合在一起)\ndef add_curvature_features(X, center_x, center_y, fps):\n    vel_x = center_x.diff().fillna(0); vel_y = center_y.diff().fillna(0)\n    acc_x = vel_x.diff().fillna(0); acc_y = vel_y.diff().fillna(0)\n    cross_prod = vel_x * acc_y - vel_y * acc_x\n    vel_mag = np.sqrt(vel_x**2 + vel_y**2)\n    curvature = np.abs(cross_prod) / (vel_mag**3 + 1e-6)\n    for w in [25, 50]:\n        ws = _scale(w, fps)\n        X[f'curv_mean_{w}'] = curvature.rolling(ws, min_periods=1).mean()\n    return X\n\ndef add_multiscale_features(X, center_x, center_y, fps):\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * fps\n    for scale in [20, 60]:\n        ws = _scale(scale, fps)\n        X[f'sp_m{scale}'] = speed.rolling(ws, min_periods=1).mean()\n        X[f'sp_s{scale}'] = speed.rolling(ws, min_periods=1).std()\n    return X\n\n# ================= 单体特征提取 (Single) =================\n\ndef transform_single(single_mouse, body_parts_tracked, fps, arena_dims=None):\n    \"\"\"\n    输入: 单只老鼠的坐标 DataFrame (index=frame, columns=bodyparts)\n    输出: 特征 DataFrame\n    \"\"\"\n    # 获取这只老鼠有哪些身体部位被追踪了 (例如 'nose', 'tail_base')\n    available_body_parts = single_mouse.columns.get_level_values(0)\n    \n    # 1. 基础点对点距离 (形态学)\n    # 计算所有部位两两之间的欧几里得距离平方\n    X = pd.DataFrame({\n        f\"{p1}+{p2}\": np.square(single_mouse[p1] - single_mouse[p2]).sum(axis=1)\n        for p1, p2 in itertools.combinations(body_parts_tracked, 2)\n        if p1 in available_body_parts and p2 in available_body_parts\n    })\n    \n    # 2. 几何特征 (面积与角度)\n    # 头部面积 (Nose-EarL-EarR)\n    if all(p in available_body_parts for p in ['nose', 'ear_left', 'ear_right']):\n        X['head_area'] = _triangle_area(single_mouse['nose'], single_mouse['ear_left'], single_mouse['ear_right'])\n    \n    # 脊柱弯曲角 (Nose-Center-Tail)\n    if all(p in available_body_parts for p in ['nose', 'body_center', 'tail_base']):\n        v1 = single_mouse['nose'] - single_mouse['body_center']\n        v2 = single_mouse['tail_base'] - single_mouse['body_center']\n        X['spine_angle'] = _angle_between(v1, v2)\n\n    # 3. 运动学特征 (基于身体中心)\n    if 'body_center' in available_body_parts:\n        cx = single_mouse['body_center']['x']\n        cy = single_mouse['body_center']['y']\n        \n        # 场地位置 (Wall hugging)\n        X = add_arena_features(X, cx, cy, arena_dims)\n\n        for w in [5, 15, 30]:\n            ws = _scale(w, fps) # 自适应窗口\n            roll = dict(min_periods=1, center=True)\n            \n            # 位置均值 (平滑轨迹)\n            X[f'cx_m{w}'] = cx.rolling(ws, **roll).mean()\n            X[f'cy_m{w}'] = cy.rolling(ws, **roll).mean()\n            \n            # 速度与高频震动 (Jitter)\n            vel_x = cx.diff().fillna(0); vel_y = cy.diff().fillna(0)\n            speed = np.sqrt(vel_x**2 + vel_y**2) * fps\n            \n            X[f'speed_mean_{w}'] = speed.rolling(ws, **roll).mean()\n            \n            # Jitter: 加速度的大小 (Grooming时身体中心会有微小高频震动)\n            acc_mag = np.sqrt(vel_x.diff()**2 + vel_y.diff()**2).fillna(0)\n            X[f'jitter_{w}'] = acc_mag.rolling(ws, **roll).mean()\n            \n            # 静止比率 (Resting Ratio): 窗口内有多少比例的时间速度接近0\n            is_resting = (speed < 1.0).astype(float)\n            X[f'rest_ratio_{w}'] = is_resting.rolling(ws, **roll).mean()\n\n        # 调用其他特征函数\n        X = add_curvature_features(X, cx, cy, fps)\n        X = add_multiscale_features(X, cx, cy, fps)\n\n    elite_keywords = [\n        'speed',        # 速度绝对是核心\n        'jitter',       # 抓挠/理毛的关键\n        'facing',       # 社交朝向变化\n        'spine_angle',  # 身体弯曲（转身/蜷缩）\n        'head_area'     # 头部伸缩（探测）\n    ]\n    \n    # 只保留同时满足：1. 在当前X中存在 2. 属于精英特征 的列\n    target_cols = [c for c in X.columns if any(k in c for k in elite_keywords)]\n    \n    # 2. 减少时间步长：如果内存还不够，只保留一个步长（比如只用 10）\n    # 去掉 [5, 15]，改成 [10] 可以减少一半的新增列\n    steps = [6] # 取个中间值，约 0.2秒\n    \n    for step in steps:\n        X[[f'{c}_lag{step}' for c in target_cols]] = X[target_cols].shift(step).fillna(method='bfill')\n        X[[f'{c}_lead{step}' for c in target_cols]] = X[target_cols].shift(-step).fillna(method='ffill')\n    \n    # 3. 差分只保留给 \"speed\" (加速度)\n    # 仅计算速度的变化，其他特征的差分暂时舍弃以省内存\n    speed_cols = [c for c in target_cols if 'speed' in c]\n    for c in speed_cols:\n        X[f'{c}_diff_6'] = X[c] - X[c].shift(6).fillna(0)\n    return X.astype(np.float16)\n\n# ================= 双体特征提取 (Pair) =================\n\ndef transform_pair(mouse_pair, body_parts_tracked, fps):\n    \"\"\"\n    输入: 两只老鼠的联合 DataFrame (columns=['A', 'B'])\n    输出: 交互特征 DataFrame\n    \"\"\"\n    avail_A = mouse_pair['A'].columns.get_level_values(0)\n    avail_B = mouse_pair['B'].columns.get_level_values(0)\n\n    # 1. 跨个体距离 (Inter-mouse distances)\n    # 计算 A 的所有部位 到 B 的所有部位 的距离\n    X = pd.DataFrame({\n        f\"12+{p1}+{p2}\": np.square(mouse_pair['A'][p1] - mouse_pair['B'][p2]).sum(axis=1)\n        for p1, p2 in itertools.product(body_parts_tracked, repeat=2)\n        if p1 in avail_A and p2 in avail_B\n    })\n\n    # 2. 社交几何 (Social Geometry) - 关键改进点\n    has_vec_parts_A = all(p in avail_A for p in ['nose', 'body_center'])\n    has_vec_parts_B = all(p in avail_B for p in ['nose', 'body_center'])\n    \n    if has_vec_parts_A and has_vec_parts_B:\n        cA = mouse_pair['A']['body_center']\n        cB = mouse_pair['B']['body_center']\n        noseA = mouse_pair['A']['nose']\n        \n        # 1. 构建 A 的身体向量 (Center -> Nose)\n        vec_body_A = noseA - cA\n        norm_body_A = np.sqrt(vec_body_A['x']**2 + vec_body_A['y']**2) + 1e-6\n        \n        # 单位方向向量 (Heading Vector)\n        u_A_x = vec_body_A['x'] / norm_body_A\n        u_A_y = vec_body_A['y'] / norm_body_A\n        \n        # 2. 构建 A -> B 的向量\n        vec_AB = cB - cA\n        \n        # 3. 投影计算 (Egocentric Coordinates)\n        # 前向距离 (Forward): A到B的向量 在 A朝向上的投影\n        # 正值表示 B 在 A 前方，负值表示在后方\n        X['dist_forward'] = vec_AB['x'] * u_A_x + vec_AB['y'] * u_A_y\n        \n        # 侧向距离 (Lateral): A到B的向量 在 A朝向的法线上的投影\n        # 垂直向量 (-y, x)\n        X['dist_lateral'] = vec_AB['x'] * (-u_A_y) + vec_AB['y'] * u_A_x\n        \n        # 4. 相对角度 (A看B的角度，-pi 到 pi)\n        # 使用 arctan2 计算更精确的角度\n        X['angle_A_to_B'] = np.arctan2(X['dist_lateral'], X['dist_forward'])\n    # 3. 动态交互 (Dynamic Interaction)\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        # 两只老鼠中心点的距离\n        dist_full = np.sqrt(np.square(mouse_pair['A']['body_center'] - mouse_pair['B']['body_center']).sum(axis=1))\n        \n        for w in [10, 30]:\n            ws = _scale(w, fps)\n            roll = dict(min_periods=1, center=True)\n            \n            # 接近速度 (Approach Speed): 距离变小的速率\n            # diff < 0 表示距离变小(接近)，取负号变正值\n            approach_speed = -dist_full.diff().fillna(0) * fps\n            X[f'approach_speed_{w}'] = approach_speed.rolling(ws, **roll).mean()\n            \n            # 速度同步性 (Speed Correlation): 区分追逐(高相关)和路过(低相关)\n            vA = np.sqrt(mouse_pair['A']['body_center'].diff()['x']**2 + mouse_pair['A']['body_center'].diff()['y']**2)\n            vB = np.sqrt(mouse_pair['B']['body_center'].diff()['x']**2 + mouse_pair['B']['body_center'].diff()['y']**2)\n            X[f'speed_corr_{w}'] = vA.rolling(ws, **roll).corr(vB)\n\n    elite_keywords = [\n        'speed',        # 速度绝对是核心\n        'jitter',       # 抓挠/理毛的关键\n        'facing',       # 社交朝向变化\n        'spine_angle',  # 身体弯曲（转身/蜷缩）\n        'head_area'     # 头部伸缩（探测）\n    ]\n    \n    # 只保留同时满足：1. 在当前X中存在 2. 属于精英特征 的列\n    target_cols = [c for c in X.columns if any(k in c for k in elite_keywords)]\n    \n    # 2. 减少时间步长：如果内存还不够，只保留一个步长（比如只用 10）\n    # 去掉 [5, 15]，改成 [10] 可以减少一半的新增列\n    steps = [6] # 取个中间值，约 0.2秒\n    \n    for step in steps:\n        X[[f'{c}_lag{step}' for c in target_cols]] = X[target_cols].shift(step).fillna(method='bfill')\n        X[[f'{c}_lead{step}' for c in target_cols]] = X[target_cols].shift(-step).fillna(method='ffill')\n    \n    # 3. 差分只保留给 \"speed\" (加速度)\n    # 仅计算速度的变化，其他特征的差分暂时舍弃以省内存\n    speed_cols = [c for c in target_cols if 'speed' in c]\n    for c in speed_cols:\n        X[f'{c}_diff_6'] = X[c] - X[c].shift(6).fillna(0)\n        \n    return X.astype(np.float16)\n\n# ================= 数据生成器 =================\n\ndef generate_mouse_data(dataset, traintest, traintest_directory=None, generate_single=True, generate_pair=True):\n    \"\"\"\n    读取 Parquet 文件，生成 (特征, 元数据, 标签) 的元组流。\n    \n    [修复内容]: \n    1. 增加了对 behaviors_labeled 字段类型的检查，跳过非字符串 (NaN) 数据。\n    2. 增加了 JSON 解析的异常捕获。\n    \"\"\"\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n        \n    for _, row in dataset.iterrows():\n        # 【修复关键点 1】: 检查标签是否为字符串。如果是 float (NaN)，直接跳过该视频。\n        if not isinstance(row.behaviors_labeled, str):\n            continue\n            \n        lab_id = row.lab_id\n        video_id = row.video_id\n        \n        # 1. 读取追踪数据 (Tracking Data)\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        if not os.path.exists(path): \n            continue\n            \n        try:\n            vid = pd.read_parquet(path)\n        except Exception as e:\n            print(f\"Error reading tracking file {path}: {e}\")\n            continue\n\n        # 过滤冗余身体部位 (如果存在)\n        # 注意: drop_body_parts 需在外部定义，或者在这里硬编码\n        if 'drop_body_parts' in globals() and len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n            \n        # 转换数据格式: (Frame, Mouse, BodyPart) -> (Frame, Mouse_BodyPart_Coord)\n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        del vid; gc.collect()\n        \n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n        \n        # 空间归一化 (像素 -> 厘米)\n        if hasattr(row, 'pix_per_cm_approx') and row.pix_per_cm_approx > 0:\n            pvid /= row.pix_per_cm_approx\n\n        # 2. 解析行为标签 (JSON -> DataFrame)\n        try:\n            vid_behaviors = json.loads(row.behaviors_labeled)\n            # 清洗数据格式\n            vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in vid_behaviors}))\n            vid_behaviors = [b.split(',') for b in vid_behaviors]\n            # 转换为 DataFrame: [agent, target, action]\n            vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action'])\n        except Exception as e:\n            # 【修复关键点 2】: 如果 JSON 解析依然报错，跳过该视频\n            print(f\"Skipping video {video_id}: JSON parse error - {e}\")\n            continue\n        \n        # 3. 读取标注真值 (仅训练模式)\n        annot = None\n        if traintest == 'train':\n            annot_path = path.replace(f'{traintest}_tracking', 'train_annotation')\n            if os.path.exists(annot_path):\n                try:\n                    annot = pd.read_parquet(annot_path)\n                except:\n                    continue\n            else:\n                continue\n\n        # 4. 生成 Single 模式数据 (Target == 'self')\n        if generate_single:\n            subset = vid_behaviors.query(\"target == 'self'\")\n            unique_agents = np.unique(subset.agent)\n            \n            for agent_str in unique_agents:\n                try:\n                    # 解析 Mouse ID (兼容 'mouse1' 和 '1' 两种格式)\n                    if 'mouse' in agent_str:\n                        mouse_id = int(agent_str.replace('mouse', ''))\n                    else:\n                        mouse_id = int(agent_str)\n                        \n                    # 提取单只老鼠的坐标数据\n                    if mouse_id in pvid.columns.get_level_values(0):\n                        single_mouse = pvid.loc[:, mouse_id]\n                        \n                        meta = pd.DataFrame({\n                            'video_id': video_id,\n                            'agent_id': agent_str,\n                            'target_id': 'self',\n                            'video_frame': single_mouse.index,\n                            'frames_per_second': row.frames_per_second,\n                            'arena_width': row.get('arena_width_cm', np.nan),\n                            'arena_height': row.get('arena_height_cm', np.nan)\n                        })\n                        \n                        if traintest == 'train':\n                            possible_actions = np.unique(subset.query(\"agent == @agent_str\").action)\n                            # 初始化全0标签矩阵\n                            y = pd.DataFrame(0, index=single_mouse.index, columns=possible_actions)\n                            \n                            # 填入标注数据\n                            if annot is not None:\n                                agent_annot = annot.query(\"agent_id == @mouse_id and target_id == @mouse_id\")\n                                for _, r in agent_annot.iterrows():\n                                    if r['action'] in y.columns:\n                                        y.loc[r['start_frame']:r['stop_frame'], r['action']] = 1\n                            yield 'single', single_mouse, meta, y\n                        else:\n                            # 测试模式返回需要预测的动作列表\n                            possible_actions = np.unique(subset.query(\"agent == @agent_str\").action)\n                            yield 'single', single_mouse, meta, possible_actions\n                except Exception as e:\n                    # print(f\"Error yielding single data: {e}\")\n                    pass\n\n        # 5. 生成 Pair 模式数据 (Target != 'self')\n        if generate_pair:\n            subset = vid_behaviors.query(\"target != 'self'\")\n            if len(subset) > 0:\n                all_mice = np.unique(pvid.columns.get_level_values(0))\n                # 遍历所有两两组合\n                for m1, m2 in itertools.permutations(all_mice, 2):\n                    agent_str = f\"mouse{m1}\"\n                    target_str = f\"mouse{m2}\"\n                    \n                    # 检查这对老鼠是否有互动任务\n                    pair_actions = subset.query(\"agent == @agent_str and target == @target_str\")\n                    if len(pair_actions) == 0: \n                        continue\n                        \n                    # 拼接两只老鼠的数据\n                    mouse_pair = pd.concat([pvid[m1], pvid[m2]], axis=1, keys=['A', 'B'])\n                    \n                    meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': agent_str,\n                        'target_id': target_str,\n                        'video_frame': mouse_pair.index,\n                        'frames_per_second': row.frames_per_second\n                    })\n                    \n                    possible_actions = np.unique(pair_actions.action)\n                    \n                    if traintest == 'train':\n                        y = pd.DataFrame(0, index=mouse_pair.index, columns=possible_actions)\n                        if annot is not None:\n                            # 注意 annot 中的 target_id 也是 int\n                            pair_annot = annot.query(\"agent_id == @m1 and target_id == @m2\")\n                            for _, r in pair_annot.iterrows():\n                                if r['action'] in y.columns:\n                                    y.loc[r['start_frame']:r['stop_frame'], r['action']] = 1\n                        yield 'pair', mouse_pair, meta, y\n                    else:\n                        yield 'pair', mouse_pair, meta, possible_actions\n\nprint(\"Step 4: 特征提取与生成器定义完成。\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T05:53:29.306556Z","iopub.execute_input":"2025-12-25T05:53:29.306789Z","iopub.status.idle":"2025-12-25T05:53:29.349157Z","shell.execute_reply.started":"2025-12-25T05:53:29.306770Z","shell.execute_reply":"2025-12-25T05:53:29.348521Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 运行测试脚本，查看特征矩阵","metadata":{}},{"cell_type":"code","source":"# ================= Step 5: 特征工程测试脚本 =================\n\nprint(\"Step 5: 开始测试特征提取...\")\n\n# 1. 随机取一个样本视频 (这里取第一行)\nsample_row = train.iloc[0:1].copy()\nvideo_id = sample_row.video_id.iloc[0]\nlab_id = sample_row.lab_id.iloc[0]\nprint(f\"测试视频 ID: {video_id} (Lab: {lab_id})\")\n\n# 2. 获取该视频的 FPS 和 身体部位配置\n# 注意：这里需要重新获取一下 body_parts_tracked，因为 generate_mouse_data 内部不处理这个\nbpt_str = sample_row.body_parts_tracked.iloc[0]\nbpt = json.loads(bpt_str)\n# 过滤冗余部位\nif len(bpt) > 5:\n    bpt = [b for b in bpt if b not in drop_body_parts]\nprint(f\"追踪的身体部位 ({len(bpt)}个): {bpt}\")\n\n# 3. 初始化生成器\n# 注意：这里只生成 Single 和 Pair 各自的第一个片段来做演示\ngen = generate_mouse_data(\n    sample_row, \n    'train', \n    traintest_directory=CFG.train_tracking_path,\n    generate_single=True, \n    generate_pair=True\n)\n\nfound_single = False\nfound_pair = False\n\nprint(\"\\n\" + \"=\"*40)\n\n# 4. 遍历生成器结果\nfor mode, data, meta, label in gen:\n    \n    # --- 测试 Single 模式 ---\n    if mode == 'single' and not found_single:\n        print(f\"【模式: Single】 (Agent: {meta.agent_id.iloc[0]})\")\n        \n        # 获取 FPS\n        fps = _fps_from_meta(meta, {}, default_fps=30.0)\n        print(f\"视频 FPS: {fps}\")\n        \n        # === 核心：调用 transform_single ===\n        # 假设这里没有场地大小信息，传入 None\n        X_single = transform_single(data, bpt, fps, arena_dims=None)\n        \n        print(f\"特征矩阵形状: {X_single.shape} (行=帧数, 列=特征数)\")\n        \n        # 检查我们新加的特征是否存在\n        check_feats = ['head_area', 'spine_angle', 'jitter_5', 'rest_ratio_30']\n        print(f\"检查新特征是否存在:\")\n        for f in check_feats:\n            status = \"✅ 存在\" if f in X_single.columns else \"❌ 缺失 (可能缺少对应身体部位)\"\n            print(f\"  - {f}: {status}\")\n            \n        # 打印前5行数据预览\n        print(\"\\n数据预览 (前5行, 随机5列):\")\n        # 随机选几列展示，包含新特征\n        cols_to_show = [c for c in check_feats if c in X_single.columns]\n        if len(cols_to_show) < 5:\n            cols_to_show += list(X_single.columns[:5])\n        display(X_single[cols_to_show].head())\n        \n        found_single = True\n        print(\"-\" * 40)\n\n    # --- 测试 Pair 模式 ---\n    elif mode == 'pair' and not found_pair:\n        print(f\"【模式: Pair】 (Agent: {meta.agent_id.iloc[0]}, Target: {meta.target_id.iloc[0]})\")\n        \n        fps = _fps_from_meta(meta, {}, default_fps=30.0)\n        \n        # === 核心：调用 transform_pair ===\n        X_pair = transform_pair(data, bpt, fps)\n        \n        print(f\"特征矩阵形状: {X_pair.shape}\")\n        \n        # 检查新特征\n        check_feats = ['A_facing_B', 'relative_orientation', 'approach_speed_30', 'speed_corr_30']\n        print(f\"检查社交新特征:\")\n        for f in check_feats:\n            status = \"✅ 存在\" if f in X_pair.columns else \"❌ 缺失\"\n            print(f\"  - {f}: {status}\")\n            \n        print(\"\\n数据预览 (前5行, 随机5列):\")\n        cols_to_show = [c for c in check_feats if c in X_pair.columns]\n        if len(cols_to_show) < 5:\n            cols_to_show += list(X_pair.columns[:5])\n        display(X_pair[cols_to_show].head())\n        \n        found_pair = True\n        print(\"-\" * 40)\n        \n    # 如果两种模式都测过了，就退出循环\n    if found_single and found_pair:\n        break\n\nif not found_single and not found_pair:\n    print(\"警告: 该视频没有生成任何数据 (可能没有标注或文件缺失)。\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T05:53:29.350312Z","iopub.execute_input":"2025-12-25T05:53:29.350821Z","iopub.status.idle":"2025-12-25T05:53:31.665342Z","shell.execute_reply.started":"2025-12-25T05:53:29.350797Z","shell.execute_reply":"2025-12-25T05:53:31.664585Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 定义模型训练函数","metadata":{}},{"cell_type":"code","source":"# ================= Step 6: 定义两阶段训练函数 =================\n\nfrom lightgbm import LGBMClassifier\nfrom sklearn.model_selection import GroupKFold\nfrom sklearn.metrics import f1_score\nfrom tqdm.auto import tqdm  # 引入进度条工具\n\ndef tune_threshold(y_prob, y_true):\n    \"\"\"简单的阈值搜索，寻找最佳 F1 分数对应的阈值\"\"\"\n    best_f1 = 0\n    best_th = 0.3\n    # 在 0.1 到 0.9 之间搜索\n    for th in np.arange(0.1, 0.9, 0.05):\n        pred = (y_prob >= th).astype(int)\n        # 如果预测全是0，f1 undefined，这里处理一下\n        if pred.sum() == 0 and y_true.sum() == 0:\n            f1 = 1.0\n        else:\n            f1 = f1_score(y_true, pred, zero_division=0)\n            \n        if f1 > best_f1:\n            best_f1 = f1\n            best_th = th\n    return best_th\n\n\ndef train_stage1_binary(X, label, meta, section_name):\n    \"\"\"\n    Stage 1: 训练二分类 Proposal 模型 (GPU 加速)\n    目标: 快速筛选出包含动作的时间片段。\n    \"\"\"\n    print(f\"  [Stage 1] Training Binary Proposal Model (GPU)...\")\n    \n    y_binary = label.max(axis=1).values.astype(int)\n    groups = meta.video_id\n    \n    # GPU 参数\n    model_params = {\n        'objective': 'binary',\n        'metric': 'binary_logloss',\n        'n_estimators': 600,\n        'learning_rate': 0.03,\n        'num_leaves': 63,\n        'max_depth': 10,\n        'subsample': 0.8,\n        'colsample_bytree': 0.8,\n        'scale_pos_weight': 2.0,\n        'n_jobs': -1,\n        'random_state': 42,\n        'verbosity': -1,\n        'device': 'gpu',            # 启用 GPU\n        'gpu_platform_id': 0,\n        'gpu_device_id': 0\n    }\n    \n    cv = GroupKFold(n_splits=3)\n    oof_binary = np.zeros(len(y_binary))\n    models = []\n    \n    # 手动打印 Fold 进度\n    for fold, (train_idx, valid_idx) in enumerate(cv.split(X, y_binary, groups)):\n        try:\n            model = LGBMClassifier(**model_params)\n            model.fit(X.iloc[train_idx], y_binary[train_idx])\n        except Exception as e:\n            print(f\"    Fold {fold+1}: GPU failed, fallback to CPU.\")\n            model_params['device'] = 'cpu'\n            model = LGBMClassifier(**model_params)\n            model.fit(X.iloc[train_idx], y_binary[train_idx])\n            \n        oof_binary[valid_idx] = model.predict_proba(X.iloc[valid_idx])[:, 1]\n        models.append(model)\n        \n    best_th = tune_threshold(oof_binary, y_binary)\n    score = f1_score(y_binary, (oof_binary >= best_th).astype(int), zero_division=0)\n    print(f\"  [Stage 1] Binary F1: {score:.4f} (Threshold: {best_th:.2f})\")\n    \n    # 保存\n    save_dir = f\"{CFG.model_name}/{section_name}/stage1\"\n    os.makedirs(save_dir, exist_ok=True)\n    joblib.dump(models, f\"{save_dir}/binary_models.pkl\")\n    joblib.dump(best_th, f\"{save_dir}/binary_threshold.pkl\")\n    \n    return oof_binary, best_th\n\ndef train_stage2_specific(X, label, meta, section_name):\n    \"\"\"\n    Stage 2: 训练具体动作分类器 (带进度条 + 智能 GPU 切换)\n    \"\"\"\n    print(f\"  [Stage 2] Training Specific Action Classifiers...\")\n    \n    thresholds = {}\n    save_dir_base = f\"{CFG.model_name}/{section_name}/stage2\"\n    action_columns = label.columns\n    \n    # === 进度条包裹循环 ===\n    for action in tqdm(action_columns, desc=\"  Training Actions\", leave=False):\n        \n        # 1. 准备数据\n        action_mask = ~label[action].isna().values\n        if action_mask.sum() == 0: continue\n            \n        y_action = label[action][action_mask].values.astype(int)\n        \n        # 样本过少跳过\n        if y_action.sum() < 5:\n            thresholds[action] = 1.0 \n            continue\n            \n        X_action = X[action_mask]\n        groups_action = meta.video_id[action_mask]\n        \n        # 2. 智能设备选择: 样本 > 1万 用 GPU，否则用 CPU (避免传输开销)\n        use_gpu = len(X_action) > 10000\n        \n        model_params = {\n            'objective': 'binary',\n            'metric': 'auc',\n            'n_estimators': 300,\n            'learning_rate': 0.05,\n            'num_leaves': 31,\n            'n_jobs': -1,\n            'random_state': 42,\n            'verbosity': -1,\n            'device': 'gpu' if use_gpu else 'cpu'\n        }\n        \n        # 3. 交叉验证训练\n        cv = GroupKFold(n_splits=3)\n        oof_action = np.zeros(len(y_action))\n        models = []\n        \n        for train_idx, valid_idx in cv.split(X_action, y_action, groups_action):\n            try:\n                model = LGBMClassifier(**model_params)\n                model.fit(X_action.iloc[train_idx], y_action[train_idx])\n            except:\n                # 兜底回退 CPU\n                model_params['device'] = 'cpu'\n                model = LGBMClassifier(**model_params)\n                model.fit(X_action.iloc[train_idx], y_action[train_idx])\n                \n            oof_action[valid_idx] = model.predict_proba(X_action.iloc[valid_idx])[:, 1]\n            models.append(model)\n            \n        # 4. 保存结果\n        best_th = tune_threshold(oof_action, y_action)\n        thresholds[action] = best_th\n        \n        action_dir = f\"{save_dir_base}/{action}\"\n        os.makedirs(action_dir, exist_ok=True)\n        joblib.dump(models, f\"{action_dir}/models.pkl\")\n        \n    joblib.dump(thresholds, f\"{save_dir_base}/thresholds.pkl\")\n    print(f\"  [Stage 2] Finished. Thresholds saved.\")\n    \n    return thresholds\n\nprint(\"Step 6: 训练函数定义完成。\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T05:53:31.666406Z","iopub.execute_input":"2025-12-25T05:53:31.666641Z","iopub.status.idle":"2025-12-25T05:53:40.537088Z","shell.execute_reply.started":"2025-12-25T05:53:31.666620Z","shell.execute_reply":"2025-12-25T05:53:40.536362Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 执行训练主循环 (Main Loop)","metadata":{}},{"cell_type":"code","source":"# ================= Step 7: 执行训练主循环  =================\nfrom tqdm.auto import tqdm  \n\nprint(\"Step 7: 开始训练流程...\")\n\n# 获取总的 Section 数量\ntotal_sections = len(body_parts_tracked_list)\n\n# 使用 tqdm 包裹最外层循环\n# desc: 进度条左边的描述文字\n# position=0, leave=True: 保证进度条在 Notebook 中显示正常不闪烁\nstart_section = 0 \nfor section in tqdm(range(start_section, total_sections), desc=\"Overall Progress\", position=0, leave=True):\n    bpt_str = body_parts_tracked_list[section]\n    \n    # 解析配置\n    try:\n        bpt = json.loads(bpt_str)\n        if len(bpt) > 5:\n            bpt = [b for b in bpt if b not in drop_body_parts]\n    except:\n        continue\n        \n    # 打印当前正在处理的 Section 信息\n    # tqdm.write 是专门配合进度条的打印函数，不会打断进度条的显示\n    tqdm.write(f\"\\n{'='*10} Section {section+1}/{total_sections} {'='*10}\")\n    tqdm.write(f\"Body Parts ({len(bpt)}): {bpt}\")\n    \n    # 1. 筛选出属于该配置的训练数据\n    train_subset = train[train.body_parts_tracked == bpt_str]\n    if len(train_subset) == 0: continue\n\n    # 构建 FPS 查找表\n    _fps_lookup = (train_subset[['video_id', 'frames_per_second']]\n                   .drop_duplicates('video_id')\n                   .set_index('video_id')['frames_per_second'].to_dict())\n    \n    # 2. 分别处理 'single' 和 'pair' 模式\n    # 这里也可以加一个小的进度条，或者直接打印\n    for mode in ['single', 'pair']:\n        tqdm.write(f\"  >>> Mode: {mode} (Generating Data...)\")\n        \n        # --- 数据生成 ---\n        data_list, label_list, meta_list = [], [], []\n        \n        gen = generate_mouse_data(\n            train_subset, 'train', \n            traintest_directory=CFG.train_tracking_path,\n            generate_single=(mode=='single'), \n            generate_pair=(mode=='pair')\n        )\n        \n        # 收集数据\n        count = 0\n        # 这里的生成器因为不知道总长度，很难加精确进度条，我们每生成100个打印一个点\n        for switch, data, meta, label in gen:\n            if switch != mode: continue\n            data_list.append(data)\n            meta_list.append(meta)\n            label_list.append(label)\n            count += 1\n            if count % 100 == 0:\n                print(\".\", end=\"\") # 简单的可视化\n            \n        print(\"\") # 换行\n        tqdm.write(f\"  Generated {count} clips.\")\n        if count == 0: continue\n            \n        # --- 特征转换 ---\n        tqdm.write(f\"  Extracting features...\")\n        X_parts = []\n        # 使用 tqdm 显示特征提取进度\n        for data_i, meta_i in tqdm(zip(data_list, meta_list), total=len(data_list), desc=f\"  Feat Eng ({mode})\", leave=False):\n            fps = _fps_from_meta(meta_i, _fps_lookup)\n            # 根据模式调用不同的特征函数\n            if mode == 'single':\n                Xi = transform_single(data_i, bpt, fps)\n            else:\n                Xi = transform_pair(data_i, bpt, fps)\n            X_parts.append(Xi)\n            \n        # 拼接大矩阵\n        X_all = pd.concat(X_parts, axis=0, ignore_index=True)\n        label_all = pd.concat(label_list, axis=0, ignore_index=True)\n        meta_all = pd.concat(meta_list, axis=0, ignore_index=True)\n        \n        del data_list, meta_list, label_list, X_parts\n        gc.collect()\n        \n        tqdm.write(f\"  Data shape: {X_all.shape}\")\n        \n        # --- 开始训练 ---\n        section_name = f\"{section}_{mode}\"\n        \n        # Stage 1: 二分类\n        tqdm.write(f\"  Training Stage 1...\")\n        train_stage1_binary(X_all, label_all, meta_all, section_name)\n        \n        # Stage 2: 多分类\n        tqdm.write(f\"  Training Stage 2...\")\n        train_stage2_specific(X_all, label_all, meta_all, section_name)\n        \n        del X_all, label_all, meta_all\n        gc.collect()\n\n    tqdm.write(f\"Section {section} done.\")\n    \n    # [重要] 如果你想跑完全部数据，请保留下面这行注释\n    #break \n\nprint(\"\\nStep 7: 训练流程结束。\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-25T05:53:40.538098Z","iopub.execute_input":"2025-12-25T05:53:40.538803Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 推理与生成提交文件","metadata":{}},{"cell_type":"code","source":"# ================= Step 8: 两阶段推理 (Stage 1 + Stage 2) =================\n# 包含：流式处理、模型预加载、内存保护\n\nprint(\"Step 8: 开始两阶段推理 (Stage 1 + Stage 2)...\")\n\ndef get_proposals(binary_prob, threshold=0.25, min_duration=5, smoothing_window=15):\n    \"\"\"\n    改进后的后处理函数：\n    1. 更大的平滑窗口\n    2. 二值化\n    3. 闭运算 (Filling Gaps)\n    \"\"\"\n    # 1. 概率平滑 (加大窗口，比如 15)\n    probs_smooth = pd.Series(binary_prob).rolling(window=smoothing_window, center=True, min_periods=1).mean().fillna(0).values\n    \n    # 2. 初始 Mask\n    binary_mask = (probs_smooth > threshold).astype(int)\n    \n    # 3. 填补微小空洞 (Gap Filling)\n    # 逻辑：如果两个动作之间只有 < 10 帧的间隔，认为它是连续的\n    # 这里用简单的循环实现，也可以用 scipy.ndimage.binary_closing\n    max_gap = 10 \n    zeros = np.where(binary_mask == 0)[0]\n    # 这一步比较耗时，这里提供一个简化的 Pandas 实现思路：\n    # 将 mask 转换为 Series，计算 cumsum 分组，合并短的 0 组\n    \n    # 简单实现：再做一次 Rolling Max 膨胀，然后腐蚀 (类似形态学闭运算)\n    # 先让动作“变胖”填满缝隙，再“变瘦”回原大小\n    import scipy.ndimage as ndimage\n    # 结构元素大小控制填补能力\n    struct = np.ones(max_gap) \n    binary_mask = ndimage.binary_closing(binary_mask, structure=struct).astype(int)\n\n    # 4. 提取片段\n    diffs = np.diff(np.concatenate(([0], binary_mask, [0])))\n    starts = np.where(diffs == 1)[0]\n    stops = np.where(diffs == -1)[0]\n    \n    proposals = []\n    for s, e in zip(starts, stops):\n        # 过滤过短的动作 (可能是噪声)\n        if (e - s) >= min_duration:\n            proposals.append((s, e))\n            \n    return proposals\n\n# --- 1. 预加载所有模型 (Pre-load Models) ---\n# 这样不用每次处理视频都读硬盘，速度快很多\nprint(\"Pre-loading models into memory...\")\nmodels_cache = {}\n\n# 遍历所有可能的 Section\nfor section in range(len(body_parts_tracked_list)):\n    models_cache[section] = {}\n    for mode in ['single', 'pair']:\n        section_name = f\"{section}_{mode}\"\n        model_dir = f\"{CFG.model_name}/{section_name}\"\n        \n        # 检查该 Section 是否训练过\n        if not os.path.exists(f\"{model_dir}/stage1/binary_models.pkl\"):\n            continue\n            \n        models_cache[section][mode] = {}\n        \n        # 加载 Stage 1\n        try:\n            models_cache[section][mode]['s1_models'] = joblib.load(f\"{model_dir}/stage1/binary_models.pkl\")\n            # 尝试加载阈值，如果没有就用默认 0.25\n            if os.path.exists(f\"{model_dir}/stage1/binary_threshold.pkl\"):\n                models_cache[section][mode]['s1_th'] = joblib.load(f\"{model_dir}/stage1/binary_threshold.pkl\")\n            else:\n                models_cache[section][mode]['s1_th'] = 0.25\n        except Exception as e:\n            print(f\"Warning: Failed to load Stage 1 for {section_name}: {e}\")\n            continue\n\n        # 加载 Stage 2\n        try:\n            if os.path.exists(f\"{model_dir}/stage2/thresholds.pkl\"):\n                thresholds = joblib.load(f\"{model_dir}/stage2/thresholds.pkl\")\n                models_cache[section][mode]['s2_ths'] = thresholds\n                models_cache[section][mode]['s2_models'] = {}\n                \n                for action in thresholds.keys():\n                    m_path = f\"{model_dir}/stage2/{action}/models.pkl\"\n                    if os.path.exists(m_path):\n                        models_cache[section][mode]['s2_models'][action] = joblib.load(m_path)\n        except Exception as e:\n            print(f\"Warning: Failed to load Stage 2 for {section_name}: {e}\")\n\nprint(f\"Models loaded for sections: {[k for k in models_cache.keys() if models_cache[k]]}\")\n\n\n# --- 2. 主推理循环 (Stream Inference) ---\nsubmission_list = []\n\nfor section in range(len(body_parts_tracked_list)):\n    bpt_str = body_parts_tracked_list[section]\n    \n    # [关键] 筛选当前 Section 的测试视频\n    # 如果你在本地只训练了 Section 9，而测试视频是 Section 0 的，这里就会过滤掉，这是正常的\n    test_subset = test[test.body_parts_tracked == bpt_str]\n    if len(test_subset) == 0: \n        continue\n    \n    # 检查模型是否加载了\n    if section not in models_cache:\n        # print(f\"Skipping Section {section} (Videos exist but no model trained)\")\n        continue\n\n    print(f\"Processing Section {section} ({len(test_subset)} videos)...\")\n    \n    try:\n        bpt = json.loads(bpt_str)\n        if len(bpt) > 5:\n            bpt = [b for b in bpt if b not in drop_body_parts]\n    except:\n        continue\n\n    # FPS 查找表\n    _fps_lookup = (test_subset[['video_id', 'frames_per_second']]\n                   .drop_duplicates('video_id')\n                   .set_index('video_id')['frames_per_second'].to_dict())\n\n    for mode in ['single', 'pair']:\n        # 检查该模式是否有模型\n        if mode not in models_cache[section] or 's1_models' not in models_cache[section][mode]:\n            continue\n            \n        # 拿缓存的模型\n        cache = models_cache[section][mode]\n        s1_models = cache['s1_models']\n        s1_th = cache['s1_th']\n        s2_models_dict = cache.get('s2_models', {})\n        s2_ths = cache.get('s2_ths', {})\n\n        # 生成器：一次只拿一个视频的数据\n        gen = generate_mouse_data(\n            test_subset, 'test', \n            traintest_directory=CFG.test_tracking_path, \n            generate_single=(mode=='single'), \n            generate_pair=(mode=='pair')\n        )\n        \n        for switch, data, meta, _ in gen:\n            if switch != mode: continue\n            \n            try:\n                # 1. 特征提取\n                fps = _fps_from_meta(meta, _fps_lookup)\n                if mode == 'single':\n                    X = transform_single(data, bpt, fps)\n                else:\n                    X = transform_pair(data, bpt, fps)\n                \n                # 2. === Stage 1: 找候选片段 (Proposals) ===\n                # Ensemble 平均预测\n                s1_probs = np.mean([clf.predict_proba(X)[:, 1] for clf in s1_models], axis=0)\n                \n                # 获取时间段\n                proposals = get_proposals(s1_probs, threshold=s1_th)\n                \n                if not proposals:\n                    del X, data, meta\n                    continue\n                    \n                # 3. === Stage 2: 具体分类 ===\n                # 如果没有 Stage 2 模型 (比如全被过滤了)，就跳过\n                if not s2_models_dict:\n                    del X, data, meta\n                    continue\n\n                for start_idx, stop_idx in proposals:\n                    # 只看这一个小片段\n                    X_clip = X.iloc[start_idx:stop_idx]\n                    \n                    best_action = None\n                    best_score = -1.0\n                    \n                    # 遍历所有动作分类器\n                    for action, models in s2_models_dict.items():\n                        # 预测概率\n                        probs = np.mean([m.predict_proba(X_clip)[:, 1] for m in models], axis=0)\n                        \n                        # 取 90% 分位数作为该片段的得分\n                        raw_score = np.percentile(probs, 90)\n                        \n                        # 归一化比较\n                        th = s2_ths.get(action, 0.5)\n                        norm_score = raw_score / (th + 1e-6)\n                        \n                        if raw_score > th and norm_score > best_score:\n                            best_score = norm_score\n                            best_action = action\n                    \n                    # 只有选出了最佳动作才保存\n                    if best_action:\n                        real_start = meta.video_frame.iloc[start_idx]\n                        real_stop = meta.video_frame.iloc[stop_idx-1] + 1\n                        \n                        submission_list.append({\n                            'video_id': meta.video_id.iloc[0],\n                            'agent_id': meta.agent_id.iloc[0],\n                            'target_id': meta.target_id.iloc[0],\n                            'action': best_action,\n                            'start_frame': real_start,\n                            'stop_frame': real_stop\n                        })\n\n                # 及时清理内存\n                del X, data, meta, s1_probs\n                # gc.collect() # 不必每次都调，影响速度\n                \n            except Exception as e:\n                print(f\"Error processing video: {e}\")\n                continue\n    \n    # 跑完一个 Section 清理一次\n    gc.collect()\n\n# --- 生成最终 CSV ---\nprint(\"Inference done. Generating submission...\")\n\nif len(submission_list) > 0:\n    submission_df = pd.DataFrame(submission_list)\nelse:\n    # 兜底：如果本地测试因为没有对应的模型导致没结果，这是正常的\n    # 但为了保证代码不报错，我们生成一行假数据\n    print(\"Warning: No predictions generated (likely due to Section mismatch in local test).\")\n    submission_df = pd.DataFrame({\n        'video_id': [test.video_id.iloc[0]],\n        'agent_id': ['mouse1'],\n        'target_id': ['self'],\n        'action': ['other'],\n        'start_frame': [0],\n        'stop_frame': [1]\n    })\n\n# 最后的清洗和填充\nfinal_submission = robustify(submission_df, test, 'test')\nfinal_submission.to_csv('submission.csv', index=True, index_label='row_id')\n\nprint(f\"SUCCESS: submission.csv generated with {len(final_submission)} rows.\")\nprint(final_submission.head())","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}