{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport json\nimport pandas as pd\nfrom pathlib import Path","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-13T16:19:29.889041Z","iopub.execute_input":"2025-10-13T16:19:29.890404Z","iopub.status.idle":"2025-10-13T16:19:29.894775Z","shell.execute_reply.started":"2025-10-13T16:19:29.89036Z","shell.execute_reply":"2025-10-13T16:19:29.893936Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 自动查找 MABe 数据集路径\n# =====================================================\ndef find_mabe_dataset():\n    input_root = Path(\"/kaggle/input\")\n    candidates = [p for p in input_root.iterdir() if \"mabe\" in p.name.lower()]\n    if not candidates:\n        raise FileNotFoundError(\"没找到包含 'mabe' 的数据集目录，请检查数据集挂载。\")\n    dataset_dir = candidates[0]\n    print(f\"自动检测到数据集路径：{dataset_dir}\")\n    return dataset_dir","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-13T16:19:29.896165Z","iopub.execute_input":"2025-10-13T16:19:29.896416Z","iopub.status.idle":"2025-10-13T16:19:29.913826Z","shell.execute_reply.started":"2025-10-13T16:19:29.896393Z","shell.execute_reply":"2025-10-13T16:19:29.912782Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =====================================================\n# 基本路径设定\n# =====================================================\nDATASET_DIR = find_mabe_dataset()\nTRAIN_CSV = DATASET_DIR / \"train.csv\"\nANNOTATION_DIR = DATASET_DIR / \"train_annotation\"\nTRACKING_DIR = DATASET_DIR / \"train_tracking\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-13T16:19:29.914941Z","iopub.execute_input":"2025-10-13T16:19:29.915536Z","iopub.status.idle":"2025-10-13T16:19:29.930827Z","shell.execute_reply.started":"2025-10-13T16:19:29.915513Z","shell.execute_reply":"2025-10-13T16:19:29.929823Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Step 1 清洗 train.csv，解析 behaviors_labeled\n# =====================================================\ndef clean_train_csv():\n    df = pd.read_csv(TRAIN_CSV)\n    print(f\"读取 train.csv 成功，共 {len(df)} 行\")\n\n    # 删除 body_parts_tracked 为空的行\n    df = df[df[\"body_parts_tracked\"].notna()].reset_index(drop=True)\n    print(f\"清洗后剩余 {len(df)} 行（已删除 body_parts_tracked 为空的数据）\")\n\n    # 解析 behaviors_labeled 字符串为列表\n    behavior_records = []\n    for _, row in df.iterrows():\n        try:\n            behaviors = json.loads(row[\"behaviors_labeled\"])\n            for b in behaviors:\n                agent, target, action = b.split(\",\")\n                behavior_records.append({\n                    \"lab_id\": row[\"lab_id\"],\n                    \"video_id\": row[\"video_id\"],\n                    \"agent_id\": agent,\n                    \"target_id\": target,\n                    \"action\": action\n                })\n        except Exception as e:\n            print(f\"解析 behaviors_labeled 出错: {e}\")\n\n    behaviors_df = pd.DataFrame(behavior_records)\n    print(f\"解析 behaviors_labeled 完成，共 {len(behaviors_df)} 条行为记录\")\n    \n    df.to_csv(\"/kaggle/working/train_without_MABe.csv\", index=False)\n    behaviors_df.to_csv(\"/kaggle/working/behaviors_labeled.csv\", index=False)\n\n    return df, behaviors_df\n\n# =====================================================\n# Step 2 整合 annotation → solution\n# =====================================================\ndef build_solution(train_df):\n    records = []\n    for _, row in train_df.iterrows():\n        lab_id = row[\"lab_id\"]\n        video_id = row[\"video_id\"]\n        ann_path = ANNOTATION_DIR / str(lab_id) / f\"{video_id}.parquet\"\n        if not ann_path.exists():\n            continue\n        ann_df = pd.read_parquet(ann_path)\n        ann_df[\"lab_id\"] = lab_id\n        ann_df[\"video_id\"] = video_id\n        ann_df[\"behaviors_labeled\"] = row[\"behaviors_labeled\"]\n        records.append(ann_df)\n    solution_df = pd.concat(records, ignore_index=True)\n    solution_df.to_csv(\"/kaggle/working/solution.csv\", index=False)\n    print(f\"✅ solution 构建完成，共 {len(solution_df)} 行\")\n    return solution_df\n\n# =====================================================\n# Step 3 整合 tracking → pvid（像素转厘米）\n# =====================================================\ndef build_pvid(train_df):\n    pvid_dict = {}\n    for _, row in train_df.iterrows():\n        lab_id = row[\"lab_id\"]\n        video_id = row[\"video_id\"]\n        pix_per_cm = row[\"pix_per_cm_approx\"]\n        track_path = TRACKING_DIR / str(lab_id) / f\"{video_id}.parquet\"\n        if not track_path.exists():\n            continue\n        df = pd.read_parquet(track_path)\n        df[\"x\"] = df[\"x\"] / pix_per_cm\n        df[\"y\"] = df[\"y\"] / pix_per_cm\n        pvid = df.pivot_table(\n            index=\"video_frame\",\n            columns=[\"mouse_id\", \"bodypart\"],\n            values=[\"x\", \"y\"]\n        )\n        pvid_dict[video_id] = pvid\n        # 可选：保存每个视频的 pvid\n        pvid.to_parquet(f\"/kaggle/working/pvid_{video_id}.parquet\")\n    print(f\"✅ pvid 构建完成，共 {len(pvid_dict)} 个视频\")\n    return pvid_dict\n\n# =====================================================\n# Step 4 整合 solution + pvid + behaviors → single/pair\n# =====================================================\ndef build_single_and_pair(solution_df, pvid_dict):\n    single_records = []\n    pair_records = []\n\n    for _, row in solution_df.iterrows():\n        vid = row[\"video_id\"]\n        if vid not in pvid_dict:\n            continue\n        pvid = pvid_dict[vid]\n        sub = pvid.loc[row[\"start_frame\"]:row[\"stop_frame\"]].copy()\n        sub[\"video_id\"] = vid\n        sub[\"lab_id\"] = row[\"lab_id\"]\n        sub[\"agent_id\"] = row[\"agent_id\"]\n        sub[\"target_id\"] = row[\"target_id\"]\n        sub[\"action\"] = row[\"action\"]\n        sub[\"start_frame\"] = row[\"start_frame\"]\n        sub[\"stop_frame\"] = row[\"stop_frame\"]\n\n        if row[\"agent_id\"] == \"self\" or row[\"agent_id\"] == row[\"target_id\"]:\n            single_records.append(sub)\n        else:\n            pair_records.append(sub)\n\n    if single_records:\n        single_df = pd.concat(single_records)\n        single_df.to_csv(\"/kaggle/working/single_mouse.csv\")\n        print(f\"✅ single_mouse.csv 输出完成，共 {len(single_df)} 行\")\n    else:\n        single_df = pd.DataFrame()\n\n    if pair_records:\n        pair_df = pd.concat(pair_records)\n        pair_df.to_csv(\"/kaggle/working/mouse_pair.csv\")\n        print(f\"✅ mouse_pair.csv 输出完成，共 {len(pair_df)} 行\")\n    else:\n        pair_df = pd.DataFrame()\n\n    return single_df, pair_df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-13T16:19:29.932427Z","iopub.execute_input":"2025-10-13T16:19:29.932667Z","iopub.status.idle":"2025-10-13T16:19:29.949973Z","shell.execute_reply.started":"2025-10-13T16:19:29.932647Z","shell.execute_reply":"2025-10-13T16:19:29.949146Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =====================================================\n# 主流程\n# =====================================================\nif __name__ == \"__main__\":\n    train_df, behaviors_df = clean_train_csv()\n    solution_df = build_solution(train_df)\n    pvid_dict = build_pvid(train_df)\n    single_df, pair_df = build_single_and_pair(solution_df, pvid_dict)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-13T16:19:29.951107Z","iopub.execute_input":"2025-10-13T16:19:29.951859Z","execution_failed":"2025-10-13T16:40:49.687Z"}},"outputs":[],"execution_count":null}]}