{"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":"gpu","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"},{"sourceId":13497318,"sourceType":"datasetVersion","datasetId":8569736}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ================================================================\n# MABe Challenge - Reworked Baseline (Single-file, unified main)\n# - Unified Config flags (VIS / TRAIN / SUB)\n# - FPS-aware feature engineering preserved from baseline\n# - 5-fold CV (GroupKFold by video_id) for XGB (+ full XGB)\n# - CatBoost full, and 3x LightGBM full (different params)\n# - Train-time light Gaussian noise (optional)\n# - Test-time augmentation (TTA) with Gaussian noise\n# - Adaptive thresholds + temporal smoothing -> [start, stop) segments\n# - robustify() for non-overlap & full coverage\n# - All models saved per (section/mode/action) under CFG.MODEL_DIR\n# ================================================================\n\nimport os, sys, gc, json, math, random, hashlib, warnings\nfrom pathlib import Path\nwarnings.filterwarnings(\"ignore\")\n\nimport numpy as np\nimport pandas as pd\nimport itertools\nfrom tqdm.auto import tqdm\n\n# tree models\nimport lightgbm as lgb\nXGBOOST_AVAILABLE = False\nCATBOOST_AVAILABLE = False\ntry:\n    from xgboost import XGBClassifier\n    XGBOOST_AVAILABLE = True\nexcept Exception:\n    pass\n\ntry:\n    from catboost import CatBoostClassifier\n    CATBOOST_AVAILABLE = True\nexcept Exception:\n    pass\n\nfrom sklearn.model_selection import GroupKFold\nfrom sklearn.base import clone\nimport joblib\n\n# Optional GPU probe\ntry:\n    import torch\n    TORCH_OK = True\nexcept Exception:\n    TORCH_OK = False\n\n# ---------------------------\n# 1) CONFIG\n# ---------------------------\nclass CFG:\n    # Flags\n    VIS   = False    # run EDA only\n    TRAIN = False    # train & save weights\n    SUB   = True    # load weights & predict & submit\n\n    # Repro\n    SEED = 42\n\n    # Paths (Kaggle)\n    INPUT_DIR = Path(\"/kaggle/input/MABe-mouse-behavior-detection\")\n    WORK_DIR  = Path(\"/kaggle/working\")\n\n    if TRAIN:\n        MODEL_DIR = WORK_DIR / \"models\"\n        META_DIR  = WORK_DIR / \"model_meta\"\n        SUB_PATH  = WORK_DIR / \"submission.csv\"\n    elif SUB:\n        MODEL_DIR = Path(\"/kaggle/input/mice2025exp-notebook-test/models\")\n        META_DIR  = Path(\"/kaggle/input/mice2025exp-notebook-test/model_meta\")\n        SUB_PATH  = WORK_DIR / \"submission.csv\"\n\n    # CV\n    N_FOLDS_XGB = 5\n\n    # Train-time augmentation\n    TRAIN_NOISE_STD   = 0.0   # 0 to disable (e.g. 0.003~0.01 for light noise)\n    TRAIN_NOISE_TIMES = 0     # 0 to disable\n\n    # Test-time augmentation (TTA)\n    TTA_NOISE_STD = 0.005\n    TTA_N         = 5\n\n    # GPU preference\n    USE_GPU = True\n\n    # LightGBM (3 variants)\n    LGBM_PARAMSETS = [\n        dict(n_estimators=225, learning_rate=0.07, min_child_samples=40,\n             num_leaves=31, subsample=0.8, colsample_bytree=0.8, verbosity=-1),\n        dict(n_estimators=150, learning_rate=0.10, min_child_samples=20,\n             num_leaves=63, max_depth=8, subsample=0.7, colsample_bytree=0.9,\n             reg_alpha=0.1, reg_lambda=0.1, verbosity=-1),\n        dict(n_estimators=100, learning_rate=0.05, min_child_samples=30,\n             num_leaves=127, max_depth=10, subsample=0.75, verbosity=-1),\n    ]\n\n    # XGB (fold & full)\n    XGB_BASE = dict(\n        n_estimators=180, learning_rate=0.08, max_depth=6,\n        min_child_weight=5, subsample=0.8, colsample_bytree=0.8,\n        eval_metric=\"logloss\", n_jobs=-1, random_state=SEED\n    )\n\n    # CatBoost (full)\n    CAT_BASE = dict(\n        iterations=120, learning_rate=0.1, depth=6,\n        verbose=False, allow_writing_files=False, random_seed=SEED\n    )\n\n    # ---------- EDA knobs ----------\n    EDA_SAMPLE_TRAIN_IDXS = [5772, 484, 397, 428, 8669, 306]   # 會逐個載入並畫幾個關鍵幀\n    EDA_FRAMES_PER_VIDEO  = 3                                  # 每支影片取幾個位置等距抽樣畫圖\n    EDA_DO_ANIM           = True                               # 是否做行為片段動態骨架動畫\n    EDA_ANIM_PADDING_FR   = 20                                 # 動畫上下游額外幀數\n    EDA_MAX_ANNOT_VIDEOS  = -1                                 # 建全量標註表時最多掃多少支影片；-1=不限制（全掃）\n    EDA_SAVE_ANIM_HTML    = False                              # 是否把動畫存成 HTML（kaggle/working 下）\n    EDA_RANDOM_SEED       = 42                                 # EDA 隨機選樣的種子\n\n# Ensure dirs\nif CFG.TRAIN:\n    CFG.MODEL_DIR.mkdir(parents=True, exist_ok=True)\n    CFG.META_DIR.mkdir(parents=True, exist_ok=True)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T12:44:24.100951Z","iopub.execute_input":"2025-10-25T12:44:24.101696Z","iopub.status.idle":"2025-10-25T12:44:24.112309Z","shell.execute_reply.started":"2025-10-25T12:44:24.101649Z","shell.execute_reply":"2025-10-25T12:44:24.111620Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# 2) UTILS\n# ---------------------------\ndef seed_everything(seed: int):\n    random.seed(seed)\n    np.random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    try:\n        import torch\n        torch.manual_seed(seed)\n        torch.cuda.manual_seed_all(seed)\n    except Exception:\n        pass\n\ndef has_gpu() -> bool:\n    if not CFG.USE_GPU:\n        return False\n    if TORCH_OK and torch.cuda.is_available():\n        return True\n    return os.environ.get(\"CUDA_VISIBLE_DEVICES\", \"\") not in (\"\", \"-1\")\n\nGPU_AVAILABLE = has_gpu()\nseed_everything(CFG.SEED)\nprint(f\"[ENV] GPU_AVAILABLE={GPU_AVAILABLE}, XGB={XGBOOST_AVAILABLE}, CAT={CATBOOST_AVAILABLE}\")\n\ndef bpt_slug(s: str) -> str:\n    return hashlib.md5(s.encode(\"utf-8\")).hexdigest()[:10]\n\ndef save_columns(cols, path: Path):\n    path.parent.mkdir(parents=True, exist_ok=True)\n    with open(path, \"w\") as f:\n        for c in cols:\n            f.write(str(c) + \"\\n\")\n\ndef save_model_any(model, path: Path, model_type: str):\n    path.parent.mkdir(parents=True, exist_ok=True)\n    if model_type == \"xgb\":\n        try:\n            model.save_model(str(path))\n        except Exception:\n            joblib.dump(model, path)\n    elif model_type == \"cat\":\n        try:\n            model.save_model(str(path))\n        except Exception:\n            joblib.dump(model, path)\n    elif model_type == \"lgbm\":\n        try:\n            model.booster_.save_model(str(path))\n        except Exception:\n            joblib.dump(model, path)\n    else:\n        joblib.dump(model, path)\n\n# ---------------------------\n# 3) DATA LOADING\n# ---------------------------\ntrain = pd.read_csv(CFG.INPUT_DIR / \"train.csv\")\ntest  = pd.read_csv(CFG.INPUT_DIR / \"test.csv\")\ntrain[\"n_mice\"] = 4 - train[[\"mouse1_strain\",\"mouse2_strain\",\"mouse3_strain\",\"mouse4_strain\"]].isna().sum(axis=1)\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\ndrop_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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T12:44:24.113267Z","iopub.execute_input":"2025-10-25T12:44:24.113455Z","iopub.status.idle":"2025-10-25T12:44:24.197301Z","shell.execute_reply.started":"2025-10-25T12:44:24.113441Z","shell.execute_reply":"2025-10-25T12:44:24.196729Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# 4) DATA PIPELINE (baseline logic)\n# ---------------------------\ndef generate_mouse_data(dataset, traintest, traintest_directory=None, generate_single=True, generate_pair=True, verbose=False):\n    assert traintest in [\"train\", \"test\"]\n    if traintest_directory is None:\n        traintest_directory = f\"{CFG.INPUT_DIR}/{traintest}_tracking\"\n\n    for _, row in dataset.iterrows():\n        lab_id  = row.lab_id\n        video_id = row.video_id\n\n        if type(row.behaviors_labeled) != str:\n            if verbose: print(\"No labeled behaviors:\", lab_id, video_id)\n            continue\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n        if len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n        pvid = vid.pivot(columns=[\"mouse_id\",\"bodypart\"], index=\"video_frame\", values=[\"x\",\"y\"])\n        del vid\n        pvid = pvid.reorder_levels([1,2,0], axis=1).T.sort_index().T\n        pvid /= row.pix_per_cm_approx  # convert to cm\n\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]\n        vid_behaviors = pd.DataFrame(vid_behaviors, columns=[\"agent\",\"target\",\"action\"])\n\n        if traintest == \"train\":\n            try:\n                annot = pd.read_parquet(path.replace(\"train_tracking\",\"train_annotation\"))\n            except FileNotFoundError:\n                continue\n\n        if generate_single:\n            vb_single = vid_behaviors.query(\"target == 'self'\")\n            for mouse_id_str in np.unique(vb_single.agent):\n                try:\n                    mouse_id = int(mouse_id_str[-1])\n                    actions = np.unique(vb_single.query(\"agent == @mouse_id_str\").action)\n                    single_mouse = pvid.loc[:, mouse_id]\n                    single_meta = pd.DataFrame({\n                        \"video_id\": video_id, \"agent_id\": mouse_id_str, \"target_id\": \"self\",\n                        \"video_frame\": single_mouse.index, \"frames_per_second\": row.frames_per_second\n                    })\n                    if traintest == \"train\":\n                        single_label = pd.DataFrame(0.0, columns=actions, index=single_mouse.index)\n                        annot_subset = annot.query(\"(agent_id == @mouse_id) & (target_id == @mouse_id)\")\n                        for i in range(len(annot_subset)):\n                            ar = annot_subset.iloc[i]\n                            single_label.loc[ar[\"start_frame\"]:ar[\"stop_frame\"], ar.action] = 1.0\n                        yield \"single\", single_mouse, single_meta, single_label\n                    else:\n                        yield \"single\", single_mouse, single_meta, actions\n                except KeyError:\n                    pass\n\n        if generate_pair:\n            vb_pair = vid_behaviors.query(\"target != 'self'\")\n            if len(vb_pair) > 0:\n                for agent, target in itertools.permutations(np.unique(pvid.columns.get_level_values(\"mouse_id\")), 2):\n                    agent_str = f\"mouse{agent}\"\n                    target_str = f\"mouse{target}\"\n                    actions = np.unique(vb_pair.query(\"(agent == @agent_str) & (target == @target_str)\").action)\n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=[\"A\",\"B\"])\n                    pair_meta = pd.DataFrame({\n                        \"video_id\": video_id, \"agent_id\": agent_str, \"target_id\": target_str,\n                        \"video_frame\": mouse_pair.index, \"frames_per_second\": row.frames_per_second\n                    })\n                    if traintest == \"train\":\n                        pair_label = pd.DataFrame(0.0, columns=actions, index=mouse_pair.index)\n                        annot_subset = annot.query(\"(agent_id == @agent) & (target_id == @target)\")\n                        for i in range(len(annot_subset)):\n                            ar = annot_subset.iloc[i]\n                            pair_label.loc[ar[\"start_frame\"]:ar[\"stop_frame\"], ar.action] = 1.0\n                        yield \"pair\", mouse_pair, pair_meta, pair_label\n                    else:\n                        yield \"pair\", mouse_pair, pair_meta, actions\n\n# ---- Feature helpers (fps-aware)\ndef _scale(n_frames_at_30fps, fps, ref=30.0):\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    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    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\ndef add_curvature_features(X, cx, cy, fps):\n    vx, vy = cx.diff(), cy.diff()\n    ax, ay = vx.diff(), vy.diff()\n    cross = vx * ay - vy * ax\n    vmag = np.sqrt(vx**2 + vy**2)\n    curv = np.abs(cross) / (vmag**3 + 1e-6)\n    for w in [30, 60]:\n        ws = _scale(w, fps); X[f\"curv_mean_{w}\"] = curv.rolling(ws, min_periods=max(1, ws//6)).mean()\n    ang = np.arctan2(vy, vx); ang_ch = np.abs(ang.diff())\n    ws = _scale(30, fps); X[f\"turn_rate_{30}\"] = ang_ch.rolling(ws, min_periods=max(1, ws//6)).sum()\n    return X\n\ndef add_multiscale_features(X, cx, cy, fps):\n    speed = np.sqrt(cx.diff()**2 + cy.diff()**2) * float(fps)\n    for s in [10, 40, 160]:\n        ws = _scale(s, fps)\n        if len(speed) >= ws:\n            X[f\"sp_m{s}\"] = speed.rolling(ws, min_periods=max(1, ws//4)).mean()\n            X[f\"sp_s{s}\"] = speed.rolling(ws, min_periods=max(1, ws//4)).std()\n    if \"sp_m10\" in X.columns and \"sp_m160\" in X.columns:\n        X[\"sp_ratio\"] = X[\"sp_m10\"] / (X[\"sp_m160\"] + 1e-6)\n    return X\n\ndef add_state_features(X, cx, cy, fps):\n    speed = np.sqrt(cx.diff()**2 + cy.diff()**2) * float(fps)\n    ws_ma = _scale(15, fps); sp_ma = speed.rolling(ws_ma, min_periods=max(1, ws_ma//3)).mean()\n    try:\n        bins = [-np.inf, 0.5*fps, 2.0*fps, 5.0*fps, np.inf]\n        states = pd.cut(sp_ma, bins=bins, labels=[0,1,2,3]).astype(float)\n        for w in [60, 120]:\n            ws = _scale(w, fps)\n            if len(states) >= ws:\n                for st in [0,1,2,3]:\n                    X[f\"s{st}_{w}\"] = ((states == st).astype(float)).rolling(ws, min_periods=max(1, ws//6)).mean()\n                X[f\"trans_{w}\"] = (states != states.shift(1)).astype(float).rolling(ws, min_periods=max(1, ws//6)).sum()\n    except Exception:\n        pass\n    return X\n\ndef add_longrange_features(X, cx, cy, fps):\n    for w in [120, 240]:\n        ws = _scale(w, fps)\n        if len(cx) >= ws:\n            X[f\"x_ml{w}\"] = cx.rolling(ws, min_periods=max(5, ws//6)).mean()\n            X[f\"y_ml{w}\"] = cy.rolling(ws, min_periods=max(5, ws//6)).mean()\n    for span in [60, 120]:\n        s = _scale(span, fps)\n        X[f\"x_e{span}\"] = cx.ewm(span=s, min_periods=1).mean()\n        X[f\"y_e{span}\"] = cy.ewm(span=s, min_periods=1).mean()\n    speed = np.sqrt(cx.diff()**2 + cy.diff()**2) * float(fps)\n    for w in [60, 120]:\n        ws = _scale(w, fps)\n        if len(speed) >= ws:\n            X[f\"sp_pct{w}\"] = speed.rolling(ws, min_periods=max(5, ws//6)).rank(pct=True)\n    return X\n\ndef add_interaction_features(X, pair, avail_A, avail_B, fps):\n    if \"body_center\" not in avail_A or \"body_center\" not in avail_B: return X\n    rel_x = pair[\"A\"][\"body_center\"][\"x\"] - pair[\"B\"][\"body_center\"][\"x\"]\n    rel_y = pair[\"A\"][\"body_center\"][\"y\"] - pair[\"B\"][\"body_center\"][\"y\"]\n    rel_d = np.sqrt(rel_x**2 + rel_y**2)\n    Avx = pair[\"A\"][\"body_center\"][\"x\"].diff(); Avy = pair[\"A\"][\"body_center\"][\"y\"].diff()\n    Bvx = pair[\"B\"][\"body_center\"][\"x\"].diff(); Bvy = pair[\"B\"][\"body_center\"][\"y\"].diff()\n    A_lead = (Avx*rel_x + Avy*rel_y) / (np.sqrt(Avx**2 + Avy**2)*rel_d + 1e-6)\n    B_lead = (Bvx*(-rel_x) + Bvy*(-rel_y)) / (np.sqrt(Bvx**2 + Bvy**2)*rel_d + 1e-6)\n    for w in [30, 60]:\n        ws = _scale(w, fps)\n        X[f\"A_ld{w}\"] = A_lead.rolling(ws, min_periods=max(1, ws//6)).mean()\n        X[f\"B_ld{w}\"] = B_lead.rolling(ws, min_periods=max(1, ws//6)).mean()\n    approach = -rel_d.diff()\n    chase = approach * B_lead\n    ws = _scale(30, fps); X[f\"chase_{30}\"] = chase.rolling(ws, min_periods=max(1, ws//6)).mean()\n    for w in [60,120]:\n        ws = _scale(w, fps)\n        A_sp = np.sqrt(Avx**2 + Avy**2); B_sp = np.sqrt(Bvx**2 + Bvy**2)\n        X[f\"sp_cor{w}\"] = A_sp.rolling(ws, min_periods=max(1, ws//6)).corr(B_sp)\n    return X\n\ndef transform_single(single_mouse, body_parts_tracked, fps):\n    avail = single_mouse.columns.get_level_values(0)\n    X = pd.DataFrame({\n        f\"{p1}+{p2}\": np.square(single_mouse[p1] - single_mouse[p2]).sum(axis=1, skipna=False)\n        for p1, p2 in itertools.combinations(body_parts_tracked, 2)\n        if p1 in avail and p2 in avail\n    })\n    X = X.reindex(columns=[f\"{p1}+{p2}\" for p1, p2 in itertools.combinations(body_parts_tracked, 2)], copy=False)\n    if all(p in single_mouse.columns for p in [\"ear_left\",\"ear_right\",\"tail_base\"]):\n        lag = _scale(10, fps)\n        sh = single_mouse[[\"ear_left\",\"ear_right\",\"tail_base\"]].shift(lag)\n        speeds = pd.DataFrame({\n            \"sp_lf\":  np.square(single_mouse[\"ear_left\"]  - sh[\"ear_left\"]).sum(axis=1, skipna=False),\n            \"sp_rt\":  np.square(single_mouse[\"ear_right\"] - sh[\"ear_right\"]).sum(axis=1, skipna=False),\n            \"sp_lf2\": np.square(single_mouse[\"ear_left\"]  - sh[\"tail_base\"]).sum(axis=1, skipna=False),\n            \"sp_rt2\": np.square(single_mouse[\"ear_right\"] - sh[\"tail_base\"]).sum(axis=1, skipna=False),\n        })\n        X = pd.concat([X, speeds], axis=1)\n    if \"nose+tail_base\" in X.columns and \"ear_left+ear_right\" in X.columns:\n        X[\"elong\"] = X[\"nose+tail_base\"] / (X[\"ear_left+ear_right\"] + 1e-6)\n    if all(p in avail 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[\"body_ang\"] = (v1[\"x\"]*v2[\"x\"] + v1[\"y\"]*v2[\"y\"]) / (np.sqrt(v1[\"x\"]**2+v1[\"y\"]**2) * np.sqrt(v2[\"x\"]**2+v2[\"y\"]**2) + 1e-6)\n    if \"body_center\" in avail:\n        cx = single_mouse[\"body_center\"][\"x\"]; cy = single_mouse[\"body_center\"][\"y\"]\n        for w in [5,15,30,60]:\n            ws = _scale(w, fps); roll = dict(min_periods=1, center=True)\n            X[f\"cx_m{w}\"] = cx.rolling(ws, **roll).mean()\n            X[f\"cy_m{w}\"] = cy.rolling(ws, **roll).mean()\n            X[f\"cx_s{w}\"] = cx.rolling(ws, **roll).std()\n            X[f\"cy_s{w}\"] = cy.rolling(ws, **roll).std()\n            X[f\"x_rng{w}\"] = cx.rolling(ws, **roll).max() - cx.rolling(ws, **roll).min()\n            X[f\"y_rng{w}\"] = cy.rolling(ws, **roll).max() - cy.rolling(ws, **roll).min()\n            X[f\"disp{w}\"] = np.sqrt(cx.diff().rolling(ws, min_periods=1).sum()**2 +\n                                    cy.diff().rolling(ws, min_periods=1).sum()**2)\n            X[f\"act{w}\"]  = np.sqrt(cx.diff().rolling(ws, min_periods=1).var() +\n                                    cy.diff().rolling(ws, min_periods=1).var())\n        X = add_curvature_features(X, cx, cy, fps)\n        X = add_multiscale_features(X, cx, cy, fps)\n        X = add_state_features(X, cx, cy, fps)\n        X = add_longrange_features(X, cx, cy, fps)\n    if all(p in avail for p in [\"nose\",\"tail_base\"]):\n        nt = np.sqrt((single_mouse[\"nose\"][\"x\"] - single_mouse[\"tail_base\"][\"x\"])**2 +\n                     (single_mouse[\"nose\"][\"y\"] - single_mouse[\"tail_base\"][\"y\"])**2)\n        for lag in [10,20,40]:\n            l = _scale(lag, fps); X[f\"nt_lg{lag}\"] = nt.shift(l); X[f\"nt_df{lag}\"] = nt - nt.shift(l)\n    if all(p in avail for p in [\"ear_left\",\"ear_right\"]):\n        ed = np.sqrt((single_mouse[\"ear_left\"][\"x\"] - single_mouse[\"ear_right\"][\"x\"])**2 +\n                     (single_mouse[\"ear_left\"][\"y\"] - single_mouse[\"ear_right\"][\"y\"])**2)\n        for off in [-20,-10,10,20]:\n            o = _scale_signed(off, fps); X[f\"ear_o{off}\"] = ed.shift(-o)\n        w = _scale(30, fps)\n        X[\"ear_con\"] = ed.rolling(w, min_periods=1, center=True).std() / (ed.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n    return X.astype(np.float32, copy=False)\n\ndef transform_pair(pair, body_parts_tracked, fps):\n    avail_A = pair[\"A\"].columns.get_level_values(0)\n    avail_B = pair[\"B\"].columns.get_level_values(0)\n    X = pd.DataFrame({\n        f\"12+{p1}+{p2}\": np.square(pair[\"A\"][p1] - pair[\"B\"][p2]).sum(axis=1, skipna=False)\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    X = X.reindex(columns=[f\"12+{p1}+{p2}\" for p1, p2 in itertools.product(body_parts_tracked, repeat=2)], copy=False)\n    if (\"A\",\"ear_left\") in pair.columns and (\"B\",\"ear_left\") in pair.columns:\n        lag = _scale(10, fps)\n        shA = pair[\"A\"][\"ear_left\"].shift(lag); shB = pair[\"B\"][\"ear_left\"].shift(lag)\n        speeds = pd.DataFrame({\n            \"sp_A\":  np.square(pair[\"A\"][\"ear_left\"] - shA).sum(axis=1, skipna=False),\n            \"sp_AB\": np.square(pair[\"A\"][\"ear_left\"] - shB).sum(axis=1, skipna=False),\n            \"sp_B\":  np.square(pair[\"B\"][\"ear_left\"] - shB).sum(axis=1, skipna=False),\n        })\n        X = pd.concat([X, speeds], axis=1)\n    if \"nose+tail_base\" in X.columns and \"ear_left+ear_right\" in X.columns:\n        X[\"elong\"] = X[\"nose+tail_base\"] / (X[\"ear_left+ear_right\"] + 1e-6)\n    if all(p in avail_A for p in [\"nose\",\"tail_base\"]) and all(p in avail_B for p in [\"nose\",\"tail_base\"]):\n        dir_A = pair[\"A\"][\"nose\"] - pair[\"A\"][\"tail_base\"]\n        dir_B = pair[\"B\"][\"nose\"] - pair[\"B\"][\"tail_base\"]\n        X[\"rel_ori\"] = (dir_A[\"x\"]*dir_B[\"x\"] + dir_A[\"y\"]*dir_B[\"y\"]) / (np.sqrt(dir_A[\"x\"]**2+dir_A[\"y\"]**2) * np.sqrt(dir_B[\"x\"]**2+dir_B[\"y\"]**2) + 1e-6)\n    if all(p in avail_A for p in [\"nose\"]) and all(p in avail_B for p in [\"nose\"]):\n        cur = np.square(pair[\"A\"][\"nose\"] - pair[\"B\"][\"nose\"]).sum(axis=1, skipna=False)\n        lag = _scale(10, fps)\n        shA = pair[\"A\"][\"nose\"].shift(lag); shB = pair[\"B\"][\"nose\"].shift(lag)\n        past = np.square(shA - shB).sum(axis=1, skipna=False)\n        X[\"appr\"] = cur - past\n    if \"body_center\" in avail_A and \"body_center\" in avail_B:\n        cd = np.sqrt((pair[\"A\"][\"body_center\"][\"x\"] - pair[\"B\"][\"body_center\"][\"x\"])**2 +\n                     (pair[\"A\"][\"body_center\"][\"y\"] - pair[\"B\"][\"body_center\"][\"y\"])**2)\n        X[\"v_cls\"] = (cd < 5.0).astype(float)\n        X[\"cls\"]   = ((cd >= 5.0) & (cd < 15.0)).astype(float)\n        X[\"med\"]   = ((cd >= 15.0) & (cd < 30.0)).astype(float)\n        X[\"far\"]   = (cd >= 30.0).astype(float)\n        cd_full = np.square(pair[\"A\"][\"body_center\"] - pair[\"B\"][\"body_center\"]).sum(axis=1, skipna=False)\n        for w in [5,15,30,60]:\n            ws = _scale(w, fps); roll = dict(min_periods=1, center=True)\n            X[f\"d_m{w}\"]  = cd_full.rolling(ws, **roll).mean()\n            X[f\"d_s{w}\"]  = cd_full.rolling(ws, **roll).std()\n            X[f\"d_mn{w}\"] = cd_full.rolling(ws, **roll).min()\n            X[f\"d_mx{w}\"] = cd_full.rolling(ws, **roll).max()\n            d_var = cd_full.rolling(ws, **roll).var()\n            X[f\"int{w}\"] = 1 / (1 + d_var)\n            Axd = pair[\"A\"][\"body_center\"][\"x\"].diff()\n            Ayd = pair[\"A\"][\"body_center\"][\"y\"].diff()\n            Bxd = pair[\"B\"][\"body_center\"][\"x\"].diff()\n            Byd = pair[\"B\"][\"body_center\"][\"y\"].diff()\n            coord = Axd*Bxd + Ayd*Byd\n            X[f\"co_m{w}\"] = coord.rolling(ws, **roll).mean()\n            X[f\"co_s{w}\"] = coord.rolling(ws, **roll).std()\n    if \"nose\" in avail_A and \"nose\" in avail_B:\n        nn = np.sqrt((pair[\"A\"][\"nose\"][\"x\"] - pair[\"B\"][\"nose\"][\"x\"])**2 +\n                     (pair[\"A\"][\"nose\"][\"y\"] - pair[\"B\"][\"nose\"][\"y\"])**2)\n        for lag in [10,20,40]:\n            l = _scale(lag, fps)\n            X[f\"nn_lg{lag}\"]  = nn.shift(l)\n            X[f\"nn_ch{lag}\"]  = nn - nn.shift(l)\n            is_cl = (nn < 10.0).astype(float)\n            X[f\"cl_ps{lag}\"]  = is_cl.rolling(l, min_periods=1).mean()\n    if \"body_center\" in avail_A and \"body_center\" in avail_B:\n        Avx = pair[\"A\"][\"body_center\"][\"x\"].diff(); Avy = pair[\"A\"][\"body_center\"][\"y\"].diff()\n        Bvx = pair[\"B\"][\"body_center\"][\"x\"].diff(); Bvy = pair[\"B\"][\"body_center\"][\"y\"].diff()\n        val = (Avx*Bvx + Avy*Bvy) / (np.sqrt(Avx**2+Avy**2) * np.sqrt(Bvx**2+Bvy**2) + 1e-6)\n        for off in [-20,-10,0,10,20]:\n            o = _scale_signed(off, fps); X[f\"va_{off}\"] = val.shift(-o)\n        w = _scale(30, fps)\n        cd_full = np.square(pair[\"A\"][\"body_center\"] - pair[\"B\"][\"body_center\"]).sum(axis=1, skipna=False)\n        X[\"int_con\"] = cd_full.rolling(w, min_periods=1, center=True).std() / (cd_full.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n        X = add_interaction_features(X, pair, avail_A, avail_B, fps)\n    return X.astype(np.float32, copy=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T12:44:24.222594Z","iopub.execute_input":"2025-10-25T12:44:24.223002Z","iopub.status.idle":"2025-10-25T12:44:24.274054Z","shell.execute_reply.started":"2025-10-25T12:44:24.222982Z","shell.execute_reply":"2025-10-25T12:44:24.273330Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# 5) TRAIN PREP\n# ---------------------------\ndef prepare_train_mats(train_subset: pd.DataFrame, body_parts_tracked, mode: str):\n    feats_parts, meta_list, label_list = [], [], []\n    fps_lookup = (\n        train_subset[[\"video_id\",\"frames_per_second\"]]\n        .drop_duplicates(\"video_id\").set_index(\"video_id\")[\"frames_per_second\"].to_dict()\n    )\n    for switch, data_i, meta_i, label_i in generate_mouse_data(\n        train_subset, \"train\",\n        generate_single=(mode==\"single\"),\n        generate_pair=(mode==\"pair\"),\n        verbose=False\n    ):\n        if switch != mode: \n            continue\n        fps_i = _fps_from_meta(meta_i, fps_lookup, default_fps=30.0)\n        if mode == \"single\":\n            Xi = transform_single(data_i, body_parts_tracked, fps_i).astype(np.float32)\n        else:\n            Xi = transform_pair(data_i, body_parts_tracked, fps_i).astype(np.float32)\n        feats_parts.append(Xi); meta_list.append(meta_i); label_list.append(label_i)\n\n    if len(feats_parts) == 0:\n        return None, None, None, fps_lookup\n\n    X_tr  = pd.concat(feats_parts, axis=0, ignore_index=True)\n    meta  = pd.concat(meta_list,  axis=0, ignore_index=True)\n    label = pd.concat(label_list, axis=0, ignore_index=True)\n    return X_tr, label, meta, fps_lookup\n\n\n# ---------------------------\n# 6) TRAINERS\n# ---------------------------\ndef _maybe_augment(X: np.ndarray, y: np.ndarray, noise_std: float, times: int, seed: int):\n    if noise_std <= 0 or times <= 0:\n        return X, y\n    rng = np.random.default_rng(seed)\n    Xs = [X]; ys = [y]\n    for _ in range(times):\n        Xs.append(X + rng.normal(0.0, noise_std, size=X.shape).astype(np.float32))\n        ys.append(y.copy())\n    return np.concatenate(Xs, 0), np.concatenate(ys, 0)\n\ndef _xgb_params():\n    p = CFG.XGB_BASE.copy()\n    p[\"tree_method\"] = \"gpu_hist\" if (GPU_AVAILABLE and CFG.USE_GPU) else \"hist\"\n    return p\n\ndef _cat_params():\n    p = CFG.CAT_BASE.copy()\n    if GPU_AVAILABLE and CFG.USE_GPU:\n        p[\"task_type\"] = \"GPU\"\n    return p\n\ndef _lgbm_wrap_params(base_params: dict):\n    p = base_params.copy()\n    # 固定必要参数\n    p.update(dict(objective=\"binary\", random_state=CFG.SEED, n_jobs=-1))\n    # 正确的 GPU 開關鍵名是 device_type\n    if GPU_AVAILABLE and CFG.USE_GPU:\n        p[\"device_type\"] = \"gpu\"\n        # 用 double precision，GPU 分裂更稳定一点\n        p[\"gpu_use_dp\"] = True\n    else:\n        p[\"device_type\"] = \"cpu\"\n\n    # 提高稳定性的一些保守设置\n    p.setdefault(\"min_data_in_leaf\", max(20, p.get(\"min_child_samples\", 20)))\n    p.setdefault(\"feature_pre_filter\", False)\n    p.setdefault(\"zero_as_missing\", True)\n    p.setdefault(\"max_bin\", 255)\n    return p\n\n\ndef train_models_for_action(X_df: pd.DataFrame, y_series: pd.Series, meta_df: pd.DataFrame,\n                            save_root: Path, action: str):\n    save_root.mkdir(parents=True, exist_ok=True)\n\n    # 先做 y 掩码（只在有标签的行上训练/筛特征）\n    y_raw = y_series.to_numpy()\n    mask  = ~pd.isna(y_raw)\n    if mask.sum() == 0:\n        print(f\"    [SKIP] {action}: no labels\")\n        return\n\n    # 丢掉常数/全NaN特征（仅在有标签的样本上统计）\n    const_mask = (X_df.loc[mask].nunique(dropna=False) <= 1)\n    if const_mask.any():\n        tqdm.write(f\"    [Info] drop {int(const_mask.sum())} constant features for {action}\")\n        X_df = X_df.loc[:, ~const_mask].copy()\n\n    if X_df.shape[1] == 0:\n        print(f\"    [SKIP] {action}: no usable features after constant-drop\")\n        return\n\n    # 保存最终特征列（供推理对齐）\n    save_columns(list(X_df.columns), save_root / \"feature_columns.txt\")\n\n    # 准备 X / y / groups（仅取有标签的行）\n    idx = np.flatnonzero(mask)\n    y   = y_raw[mask].astype(int)\n    X   = X_df.to_numpy(np.float32, copy=False)[idx]\n    groups = meta_df.loc[mask, \"video_id\"].to_numpy()\n\n    # 少正样/单类直接跳过\n    if (y.sum() < 5) or (np.unique(y).size == 1):\n        print(f\"    [SKIP] {action}: insufficient positives ({y.sum()})\")\n        return\n\n    # ---------------- XGB: CV + full ----------------\n    if XGBOOST_AVAILABLE:\n        params = _xgb_params()\n\n        # 动态降折：折数 <= 真实 group 数；不足 2 折就跳过 CV\n        n_groups = int(pd.Series(groups).nunique())\n        n_folds = min(CFG.N_FOLDS_XGB, n_groups)\n\n        if n_folds < 2:\n            tqdm.write(f\"    [XGB] skip CV for {action}: only {n_groups} group(s)\")\n        else:\n            gkf = GroupKFold(n_splits=n_folds)\n            for k, (tr_idx, va_idx) in tqdm(\n                enumerate(gkf.split(X, y, groups)),\n                total=n_folds,\n                desc=f\"[{action}] XGB {n_folds}-fold\",\n                leave=True\n            ):\n                X_tr, y_tr = X[tr_idx], y[tr_idx]\n                X_va, y_va = X[va_idx], y[va_idx]\n                X_tr_aug, y_tr_aug = _maybe_augment(\n                    X_tr, y_tr, CFG.TRAIN_NOISE_STD, CFG.TRAIN_NOISE_TIMES, CFG.SEED + k\n                )\n                clf = XGBClassifier(**params)\n                clf.fit(X_tr_aug, y_tr_aug, eval_set=[(X_va, y_va)], verbose=False)\n                save_model_any(clf, save_root / f\"xgb_fold{k}.json\", \"xgb\")\n                del clf; gc.collect()\n\n        # full 模型（无论是否做了 CV 都训）\n        tqdm.write(f\"[{action}] XGB full\")\n        clf_full = XGBClassifier(**params)\n        X_aug, y_aug = _maybe_augment(X, y, CFG.TRAIN_NOISE_STD, CFG.TRAIN_NOISE_TIMES, CFG.SEED + 777)\n        clf_full.fit(X_aug, y_aug, verbose=False)\n        save_model_any(clf_full, save_root / \"xgb_full.json\", \"xgb\")\n        del clf_full; gc.collect()\n    else:\n        print(\"    [WARN] XGBoost unavailable, skip XGB models.\")\n\n    # ---------------- CatBoost: full ----------------\n    if CATBOOST_AVAILABLE:\n        cat = CatBoostClassifier(**_cat_params())\n        X_aug, y_aug = _maybe_augment(X, y, CFG.TRAIN_NOISE_STD, CFG.TRAIN_NOISE_TIMES, CFG.SEED + 888)\n        cat.fit(X_aug, y_aug, verbose=False)\n        save_model_any(cat, save_root / \"cat_full.cbm\", \"cat\")\n        del cat; gc.collect()\n    else:\n        print(\"    [WARN] CatBoost unavailable, skip CatBoost.\")\n\n    # ---------------- LightGBM: 3 组全量 ----------------\n    for i, base in tqdm(\n        list(enumerate(CFG.LGBM_PARAMSETS, 1)),\n        total=len(CFG.LGBM_PARAMSETS),\n        desc=f\"[{action}] LGBM variants\",\n        leave=False\n    ):\n        params = _lgbm_wrap_params(base)\n        X_aug, y_aug = _maybe_augment(X, y, CFG.TRAIN_NOISE_STD, CFG.TRAIN_NOISE_TIMES, CFG.SEED + 999 + i)\n\n        def _fit_one(params_):\n            model_ = lgb.LGBMClassifier(**params_)\n            model_.fit(X_aug, y_aug)  # 不要传 verbose\n            return model_\n\n        try:\n            lgbm = _fit_one(params_) if False else _fit_one(params)  # 占位，便于阅读\n        except Exception as e:\n            tqdm.write(f\"    [LGBM-{i}] GPU failed: {str(e).splitlines()[0]} -> fallback CPU/row-wise\")\n            params[\"device_type\"] = \"cpu\"\n            params[\"force_row_wise\"] = True\n            lgbm = _fit_one(params)\n\n        save_model_any(lgbm, save_root / f\"lgbm_full_{i}.txt\", \"lgbm\")\n        del lgbm; gc.collect()\n\ndef train_for_section(section_idx: int, body_parts_tracked_str: str):\n    bslug = bpt_slug(body_parts_tracked_str)\n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        if len(body_parts_tracked) > 5:\n            body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n    except Exception:\n        print(f\"[WARN] invalid body_parts json at section {section_idx}\")\n        return\n\n    train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n    for mode in (\"single\",\"pair\"):\n        print(f\"[Section {section_idx}] mode={mode} | #videos={train_subset.video_id.nunique()}\")\n        X_tr, label_df, meta_df, _ = prepare_train_mats(train_subset, body_parts_tracked, mode)\n        if X_tr is None:\n            print(\"  -> No data, skip.\")\n            continue\n        actions = list(label_df.columns)\n        print(f\"  -> X shape={X_tr.shape}, actions={len(actions)}\")\n        # save meta for inference\n        with open(CFG.META_DIR / f\"{bslug}_{mode}.json\", \"w\") as f:\n            json.dump({\"body_parts_tracked\": body_parts_tracked, \"mode\": mode, \"actions\": actions}, f)\n        for action in tqdm(actions, desc=f\"[Section {section_idx}] {mode} actions\", leave=True):\n            save_root = CFG.MODEL_DIR / bslug / mode / action\n            tqdm.write(f\"  [Train] action={action}\")\n            train_models_for_action(X_tr, label_df[action], meta_df, save_root, action)\n            gc.collect()\n        del X_tr, label_df, meta_df; gc.collect()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T12:44:24.275188Z","iopub.execute_input":"2025-10-25T12:44:24.275432Z","iopub.status.idle":"2025-10-25T12:44:24.303230Z","shell.execute_reply.started":"2025-10-25T12:44:24.275417Z","shell.execute_reply":"2025-10-25T12:44:24.302507Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# 7) EDA (VIS)\n# ---------------------------\nimport matplotlib.pyplot as plt\n# ===================== EDA: Visualizer + Plots + Animations =====================\ndef _safe_display_df_head(df: pd.DataFrame, n: int = 5, name: str = \"\"):\n    print(f\"\\n--- {name} (head) ---\")\n    try:\n        from IPython.display import display\n        display(df.head(n))\n    except Exception:\n        print(df.head(n).to_string())\n\nclass Visualizer:\n    \"\"\"\n    Visualize a single frame of a mouse video (skeleton + keypoints), with optional action title.\n    Inspired by: https://www.kaggle.com/code/ambrosm/mabe-eda-which-makes-sense\n    \"\"\"\n    paws = ['forepaw_left', 'forepaw_right', 'hindpaw_left', 'hindpaw_right']\n    head = ['ear_left', 'ear_right', 'nose', 'ear_left']\n\n    def __init__(self, train_df: pd.DataFrame, input_dir: Path):\n        self.train = train_df\n        self.input_dir = input_dir\n\n    def load_video(self, train_idx: int):\n        self.train_idx = train_idx\n        row = self.train.iloc[train_idx]\n        lab_id, video_id = row.lab_id, row.video_id\n        tpath = self.input_dir / \"train_tracking\" / lab_id / f\"{video_id}.parquet\"\n        self.video_name = f\"{lab_id}/{video_id}\"\n        self.vid = pd.read_parquet(tpath)\n\n        try:\n            apath = self.input_dir / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n            self.annot = pd.read_parquet(apath)\n        except FileNotFoundError:\n            self.annot = None\n\n        self.pvid = self.vid.pivot(columns=['mouse_id','bodypart'],\n                                   index='video_frame', values=['x','y'])\n        # convenience\n        self.n_mouses = len(np.unique(self.pvid.columns.get_level_values('mouse_id')))\n\n    def __len__(self):\n        return len(self.pvid)\n\n    def plot_frame(self, frame_idx: int):\n        import matplotlib.pyplot as plt\n        video_frame = self.pvid.index[frame_idx]\n        # empty guard\n        if (self.pvid.loc[video_frame] == 0).all().all():\n            print(f\"{self.train_idx}.{frame_idx} is empty.\")\n            return\n\n        colors = ['g','b','orange','brown']\n        for mouse in range(self.n_mouses):\n            color = colors[mouse % len(colors)]\n            mouse_id = mouse + 1\n            mx = self.pvid.loc[video_frame, ('x', mouse_id)].copy()\n            my = self.pvid.loc[video_frame, ('y', mouse_id)].copy()\n\n            # head\n            if 'nose' in mx.index and mx['nose'] != 0:\n                plt.fill(mx[self.head], my[self.head], color=color, alpha=0.5)\n                plt.scatter([mx['nose']], [my['nose']], s=100, color=color)\n            else:\n                # fall back to ears line\n                have_ears = [p for p in ['ear_left','ear_right'] if p in mx.index]\n                if len(have_ears) == 2:\n                    plt.plot(mx[have_ears], my[have_ears], color=color)\n\n            # synthesize head center if missing\n            if 'head' not in mx.index:\n                if all(p in mx.index for p in ['ear_left','ear_right']):\n                    mx['head'] = mx[['ear_left','ear_right']].mean()\n                    my['head'] = my[['ear_left','ear_right']].mean()\n\n            # body + tail polyline\n            parts = ['head']\n            if 'neck' in mx.index and mx['neck'] != 0: parts.append('neck')\n            if 'body_center' in mx.index and mx['body_center'] != 0: parts.append('body_center')\n            if 'tail_base' in mx.index and mx['tail_base'] != 0: parts.append('tail_base')\n            if 'tail_tip' in mx.index and mx['tail_tip'] != 0: parts.append('tail_tip')\n            parts = [p for p in parts if p in mx.index]\n            if len(parts) >= 2:\n                plt.plot(mx[parts], my[parts], color=color)\n\n            # width (laterals)\n            if all(p in mx.index for p in ['lateral_right','lateral_left']):\n                plt.plot(mx[['lateral_right','lateral_left']], my[['lateral_right','lateral_left']], color=color)\n\n            # hips\n            if all(p in mx.index for p in ['hip_right','hip_left']):\n                plt.plot(mx[['hip_right','hip_left']], my[['hip_right','hip_left']], color=color)\n\n            # paws\n            if 'forepaw_left' in mx.index:\n                have_paws = [p for p in self.paws if p in mx.index]\n                if have_paws:\n                    plt.scatter(mx[have_paws], my[have_paws], color=color)\n\n        # title with active actions (if any)\n        actions = ''\n        if self.annot is not None:\n            cur = set(self.annot.action[(self.annot.start_frame <= video_frame) & (video_frame <= self.annot.stop_frame)])\n            actions = ', '.join(sorted(list(cur)))\n        plt.title(f'{self.video_name} | frame={video_frame} | actions=[{actions}]')\n        plt.gca().set_aspect('equal')\n        plt.gca().invert_yaxis()\n        plt.tight_layout()\n        plt.show()\n\n\ndef _pivot_wide_xy(df_tracking: pd.DataFrame) -> pd.DataFrame:\n    \"\"\"Long->wide: (frame) x (mouse,part,{x,y}) → flat columns mouse{m}_{part}_{x|y}\"\"\"\n    pivot_x = df_tracking.pivot(index='video_frame', columns=['mouse_id','bodypart'], values='x')\n    pivot_y = df_tracking.pivot(index='video_frame', columns=['mouse_id','bodypart'], values='y')\n    pivot_x.columns = [f\"mouse{m}_{bp}_x\" for m, bp in pivot_x.columns]\n    pivot_y.columns = [f\"mouse{m}_{bp}_y\" for m, bp in pivot_y.columns]\n    wide = pd.concat([pivot_x, pivot_y], axis=1).sort_index(axis=1)\n    return wide\n\n\ndef _plot_skeleton_frame(frame_row: pd.Series, frame_idx: int, anatomical_parts=None, connections=None):\n    import matplotlib.pyplot as plt\n    if anatomical_parts is None:\n        anatomical_parts = ['nose','ear_left','ear_right','neck','body_center','lateral_left','lateral_right','tail_base']\n    if connections is None:\n        connections = [\n            ('nose','ear_left'), ('nose','ear_right'), ('ear_left','ear_right'),\n            ('nose','neck'), ('neck','body_center'),\n            ('body_center','lateral_left'), ('body_center','lateral_right'),\n            ('body_center','tail_base')\n        ]\n    mouse_colors = {1:'blue',2:'orange',3:'green',4:'red'}\n    plt.figure(figsize=(7.5,7.5))\n    for mouse_id in range(1,5):\n        # require a basic key to decide if this mouse exists\n        kx = f'mouse{mouse_id}_nose_x'\n        if kx not in frame_row or pd.isna(frame_row[kx]):\n            continue\n        for part in anatomical_parts:\n            cx, cy = f'mouse{mouse_id}_{part}_x', f'mouse{mouse_id}_{part}_y'\n            if cx in frame_row and cy in frame_row and pd.notna(frame_row[cx]) and pd.notna(frame_row[cy]):\n                plt.scatter(frame_row[cx], frame_row[cy], color=mouse_colors[mouse_id])\n        for a,b in connections:\n            a_x, a_y = f'mouse{mouse_id}_{a}_x', f'mouse{mouse_id}_{a}_y'\n            b_x, b_y = f'mouse{mouse_id}_{b}_x', f'mouse{mouse_id}_{b}_y'\n            if all(k in frame_row for k in [a_x,a_y,b_x,b_y]) and pd.notna(frame_row[a_x]) and pd.notna(frame_row[b_x]):\n                plt.plot([frame_row[a_x], frame_row[b_x]], [frame_row[a_y], frame_row[b_y]], color=mouse_colors[mouse_id], alpha=0.75)\n    plt.title(f\"Skeleton @ frame {frame_idx}\")\n    plt.gca().invert_yaxis()\n    plt.gca().set_aspect('equal')\n    plt.tight_layout()\n    plt.show()\n\n\ndef _animate_behavior_segment(wide_df: pd.DataFrame, annot_row: pd.Series, padding: int = 20, save_html: bool = False, out_html_path: Path = None):\n    \"\"\"Create a simple skeleton animation around one annotated behavior.\"\"\"\n    import matplotlib.pyplot as plt\n    from matplotlib.animation import FuncAnimation\n    try:\n        from IPython.display import HTML, display\n    except Exception:\n        HTML = None; display = print\n\n    start, stop = int(annot_row['start_frame']), int(annot_row['stop_frame'])\n    act  = annot_row['action']; agent = int(annot_row['agent_id'])\n    a0 = max(0, start - padding); a1 = stop + padding\n    anim_df = wide_df.loc[a0:a1].copy()\n\n    fig, ax = plt.subplots(figsize=(7.5,7.5))\n    x_min, x_max = anim_df.filter(like='_x').min().min(), anim_df.filter(like='_x').max().max()\n    y_min, y_max = anim_df.filter(like='_y').min().min(), anim_df.filter(like='_y').max().max()\n    pad = 50\n    ax.set_xlim(x_min - pad, x_max + pad)\n    ax.set_ylim(y_min - pad, y_max + pad)\n\n    anatomical_parts = ['nose','ear_left','ear_right','neck','body_center','lateral_left','lateral_right','tail_base']\n    connections = [\n        ('nose','ear_left'), ('nose','ear_right'), ('ear_left','ear_right'),\n        ('nose','neck'), ('neck','body_center'),\n        ('body_center','lateral_left'), ('body_center','lateral_right'),\n        ('body_center','tail_base')\n    ]\n    mouse_colors = {1:'blue',2:'orange',3:'green',4:'red'}\n\n    def update(i):\n        ax.clear()\n        fr = anim_df.iloc[i]; real_fr = anim_df.index[i]\n        for mouse_id in range(1,5):\n            base = f'mouse{mouse_id}_nose_x'\n            if base not in fr or pd.isna(fr[base]): \n                continue\n            for p in anatomical_parts:\n                cx,cy = f'mouse{mouse_id}_{p}_x', f'mouse{mouse_id}_{p}_y'\n                if cx in fr and cy in fr and pd.notna(fr[cx]) and pd.notna(fr[cy]):\n                    ax.scatter(fr[cx], fr[cy], color=mouse_colors[mouse_id])\n            for a,b in connections:\n                ax1, ay1 = f'mouse{mouse_id}_{a}_x', f'mouse{mouse_id}_{a}_y'\n                bx1, by1 = f'mouse{mouse_id}_{b}_x', f'mouse{mouse_id}_{b}_y'\n                if all(k in fr for k in [ax1,ay1,bx1,by1]) and pd.notna(fr[ax1]) and pd.notna(fr[bx1]):\n                    ax.plot([fr[ax1], fr[bx1]], [fr[ay1], fr[by1]], color=mouse_colors[mouse_id], alpha=0.7)\n        ax.set_title(f\"Behavior '{act}' by mouse {agent} | frame={real_fr}\")\n        ax.set_xlim(x_min - pad, x_max + pad)\n        ax.set_ylim(y_min - pad, y_max + pad)\n        ax.invert_yaxis()\n        return ax,\n\n    ani = FuncAnimation(fig, update, frames=len(anim_df), interval=50, blit=False)\n    if save_html:\n        out_html_path = out_html_path or (CFG.WORK_DIR / \"behavior_anim.html\")\n        try:\n            html = ani.to_jshtml()\n            with open(out_html_path, \"w\", encoding=\"utf-8\") as f:\n                f.write(html)\n            print(f\"[EDA] animation saved to: {out_html_path}\")\n        except Exception as e:\n            print(f\"[EDA] failed to save animation: {e}\")\n\n    if 'display' in globals():\n        try:\n            from IPython.display import HTML, display\n            display(HTML(ani.to_jshtml()))\n        except Exception:\n            pass\n\n\ndef _build_full_annotations(train_meta: pd.DataFrame, input_dir: Path, max_videos: int = -1) -> pd.DataFrame:\n    \"\"\"Concatenate all train_annotation parquet files. If max_videos>0, truncate for speed.\"\"\"\n    from tqdm.auto import tqdm\n    rows = []\n    it = train_meta.iterrows()\n    if max_videos and max_videos > 0:\n        it = list(it)[:max_videos]\n    for _, row in tqdm(it, total=(len(train_meta) if (max_videos in [-1, None, 0]) else min(max_videos, len(train_meta)))):\n        lab_id = row['lab_id']; video_id = row['video_id']\n        apath = input_dir / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n        if apath.exists():\n            tmp = pd.read_parquet(apath)\n            tmp['video_id'] = video_id\n            rows.append(tmp)\n    if len(rows) == 0:\n        return pd.DataFrame(columns=['start_frame','stop_frame','action','agent_id','target_id','video_id'])\n    return pd.concat(rows, ignore_index=True)\n\n\ndef run_eda():\n    import seaborn as sns\n    import matplotlib.pyplot as plt\n    np.random.seed(CFG.EDA_RANDOM_SEED)\n\n    print(\"=== EDA (rich) ===\")\n\n    # ---------- Meta quick look ----------\n    print(\"\\n[Meta] train.csv & test.csv\")\n    print(f\"train shape: {train.shape} | test shape: {test.shape}\")\n    _safe_display_df_head(train, 5, \"train.csv\")\n    _safe_display_df_head(test, 5, \"test.csv\")\n\n    # FPS distribution\n    if 'frames_per_second' in train.columns:\n        fps_counts = train['frames_per_second'].value_counts().sort_index()\n        print(\"\\n[Meta] FPS distribution:\")\n        print(fps_counts)\n        plt.figure(figsize=(6,3.5)); fps_counts.plot(kind='bar'); plt.title(\"FPS distribution (train)\")\n        plt.xlabel(\"FPS\"); plt.ylabel(\"count\"); plt.tight_layout(); plt.show()\n\n    # body_parts_tracked combos\n    bp_counts = train['body_parts_tracked'].value_counts()\n    print(\"\\n[Meta] top body_parts_tracked combos:\")\n    print(bp_counts.head(15))\n    plt.figure(figsize=(8,3.5)); bp_counts.head(20).plot(kind='bar')\n    plt.title(\"Top body_parts_tracked (top20)\"); plt.xticks(rotation=90); plt.tight_layout(); plt.show()\n\n    # labeled presence\n    has_label = train['behaviors_labeled'].notna().astype(int)\n    print(f\"\\n[Meta] videos with labeled behaviors: {has_label.sum()} / {len(train)}\")\n    plt.figure(figsize=(4,3)); has_label.value_counts().sort_index().plot(kind='bar')\n    plt.title(\"Has labeled behaviors? (train)\"); plt.xticks([0,1], ['No','Yes']); plt.tight_layout(); plt.show()\n\n    # ---------- Visualize several videos (static frames) ----------\n    print(\"\\n[Frames] sampling a few videos and plotting key frames ...\")\n    vis = Visualizer(train, CFG.INPUT_DIR)\n    for idx in CFG.EDA_SAMPLE_TRAIN_IDXS:\n        if idx < 0 or idx >= len(train): \n            continue\n        try:\n            vis.load_video(idx)\n            L = len(vis)\n            steps = max(1, CFG.EDA_FRAMES_PER_VIDEO)\n            # positions: evenly spaced (include first frame)\n            frame_indices = np.linspace(0, max(0, L-1), num=steps, dtype=int)\n            print(f\"  - {vis.video_name}: total_frames={L} | sampled={list(frame_indices)}\")\n            for fi in frame_indices:\n                vis.plot_frame(fi)\n        except Exception as e:\n            print(f\"  [warn] fail {idx}: {e}\")\n\n    # ---------- One sample: pivot wide + skeleton frame ----------\n    print(\"\\n[Wide view] Pivot one sample video then plot skeleton @ one frame ...\")\n    # pick first available row that has an annotation file\n    sample_row = None\n    for i in range(len(train)):\n        lab_id, vid = train.iloc[i].lab_id, train.iloc[i].video_id\n        ap = CFG.INPUT_DIR / \"train_annotation\" / lab_id / f\"{vid}.parquet\"\n        tp = CFG.INPUT_DIR / \"train_tracking\" / lab_id / f\"{vid}.parquet\"\n        if ap.exists() and tp.exists():\n            sample_row = train.iloc[i]\n            break\n    if sample_row is not None:\n        lab_id = sample_row.lab_id; video_id = sample_row.video_id\n        tpath = CFG.INPUT_DIR / \"train_tracking\" / lab_id / f\"{video_id}.parquet\"\n        apath = CFG.INPUT_DIR / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n        df_tracking_sample = pd.read_parquet(tpath)\n        df_annot_sample   = pd.read_parquet(apath)\n        print(f\"  sample -> {lab_id}/{video_id} | tracking={df_tracking_sample.shape} annot={df_annot_sample.shape}\")\n        wide = _pivot_wide_xy(df_tracking_sample)\n        # plot a mid frame\n        mid_fr = int(np.clip(len(wide)//2, 0, max(0, len(wide)-1)))\n        _plot_skeleton_frame(wide.iloc[mid_fr], wide.index[mid_fr])\n\n        # optional animation around first annotation\n        if CFG.EDA_DO_ANIM and (len(df_annot_sample) > 0):\n            print(\"  creating short animation around first annotated segment ...\")\n            _animate_behavior_segment(\n                wide_df = wide,\n                annot_row = df_annot_sample.iloc[0],\n                padding = CFG.EDA_ANIM_PADDING_FR,\n                save_html = CFG.EDA_SAVE_ANIM_HTML,\n                out_html_path = CFG.WORK_DIR / f\"anim_{lab_id}_{video_id}.html\"\n            )\n    else:\n        print(\"  [info] no sample with both tracking & annotation found.\")\n\n    # ---------- Build full annotations table ----------\n    print(\"\\n[Annotations] building full annotations table (this may take a bit) ...\")\n    df_annotations_full = _build_full_annotations(\n        train_meta=train,\n        input_dir=CFG.INPUT_DIR,\n        max_videos=CFG.EDA_MAX_ANNOT_VIDEOS\n    )\n    print(f\"  annotations shape: {df_annotations_full.shape}\")\n    _safe_display_df_head(df_annotations_full, 5, \"annotations_full\")\n\n    if len(df_annotations_full) > 0:\n        # Behavior frequency\n        print(\"\\n[Behavior] frequency across all annotations\")\n        beh_counts = df_annotations_full['action'].value_counts()\n        plt.figure(figsize=(10,4)); sns.barplot(x=beh_counts.index, y=beh_counts.values)\n        plt.title(\"Behavior frequency\"); plt.xticks(rotation=45, ha='right'); plt.tight_layout(); plt.show()\n\n        # Duration analysis\n        print(\"\\n[Behavior] duration statistics (frames)\")\n        df_annotations_full['duration_frames'] = df_annotations_full['stop_frame'] - df_annotations_full['start_frame']\n        zero_duration = int((df_annotations_full['duration_frames'] == 0).sum())\n        print(f\"  zero-duration events: {zero_duration}\")\n        print(df_annotations_full['duration_frames'].describe())\n\n        plt.figure(figsize=(10,4))\n        vals = (df_annotations_full['duration_frames'] + 1).clip(lower=1)  # avoid log(0)\n        plt.hist(vals, bins=100, log=True)\n        plt.title(\"Distribution of behavior durations (log scale, +1)\")\n        plt.xlabel(\"duration+1 (frames)\"); plt.ylabel(\"count (log)\"); plt.tight_layout(); plt.show()\n\n        order = df_annotations_full.groupby('action')['duration_frames'].median().sort_values(ascending=False).index\n        plt.figure(figsize=(11,6))\n        sns.boxplot(x=(df_annotations_full['duration_frames']+1), y='action', data=df_annotations_full, order=order)\n        plt.xscale('log'); plt.title(\"Duration per behavior (log scale, +1)\")\n        plt.xlabel(\"duration+1 (frames)\"); plt.ylabel(\"behavior\"); plt.tight_layout(); plt.show()\n\n        # Lab variability (proportions per lab)\n        print(\"\\n[Lab] behavior composition by lab (proportions)\")\n        df_lab = train[['video_id','lab_id']].copy()\n        ann_lab = df_annotations_full.merge(df_lab, on='video_id', how='left')\n        crosstab = pd.crosstab(ann_lab['lab_id'], ann_lab['action'], normalize='index')\n        plt.figure(figsize=(12,7))\n        sns.heatmap(crosstab, cmap='viridis')\n        plt.title(\"Proportion of behaviors by lab\")\n        plt.xlabel(\"behavior\"); plt.ylabel(\"lab_id\"); plt.tight_layout(); plt.show()\n    else:\n        print(\"  [info] no annotations found; skip behavior stats.\")\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T12:44:24.338026Z","iopub.execute_input":"2025-10-25T12:44:24.338198Z","iopub.status.idle":"2025-10-25T12:44:24.382002Z","shell.execute_reply.started":"2025-10-25T12:44:24.338186Z","shell.execute_reply":"2025-10-25T12:44:24.381416Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# 8) Thresholds & robustify & inference helpers\n# ---------------------------\nfrom collections import defaultdict\naction_thresholds = {\n    \"default\": 0.27,\n    \"single_default\": 0.27,\n    \"pair_default\": 0.27,\n    \"single\": {\"rear\": 0.30},\n}\n\ndef _select_threshold_map(thresholds, mode: str):\n    if isinstance(thresholds, dict):\n        if any(k in thresholds for k in (\"single\",\"pair\",\"single_default\",\"pair_default\")):\n            base_default = float(thresholds.get(\"default\", 0.27))\n            mode_default = float(thresholds.get(f\"{mode}_default\", base_default))\n            mode_overrides = thresholds.get(mode, {}) or {}\n            out = defaultdict(lambda: mode_default)\n            out.update({str(k): float(v) for k, v in mode_overrides.items()})\n            return out\n        out = defaultdict(lambda: float(thresholds.get(\"default\", 0.27)))\n        out.update({str(k): float(v) for k, v in thresholds.items() if k != \"default\"})\n        return out\n    return defaultdict(lambda: 0.27)\n\ndef predict_multiclass_adaptive(pred_df: pd.DataFrame, meta_df: pd.DataFrame, thresholds):\n    pred_smoothed = pred_df.rolling(window=5, min_periods=1, center=True).mean()\n    mode = \"single\" if (\"target_id\" in meta_df.columns and meta_df[\"target_id\"].eq(\"self\").all()) else \"pair\"\n    ama = np.argmax(pred_smoothed.values, axis=1)\n    th_map = _select_threshold_map(thresholds, mode)\n    max_probs = pred_smoothed.max(axis=1).values\n    threshold_mask = np.zeros(len(pred_smoothed), dtype=bool)\n    for i, action in enumerate(pred_smoothed.columns):\n        action_mask = (ama == i)\n        thr = th_map[action]\n        threshold_mask |= (action_mask & (max_probs >= thr))\n    ama = np.where(threshold_mask, ama, -1)\n    ama = pd.Series(ama, index=meta_df.video_frame)\n\n    changes_mask = (ama != ama.shift(1)).values\n    ama_changes = ama[changes_mask]; meta_changes = meta_df[changes_mask]\n    mask = ama_changes.values >= 0\n    if len(mask) == 0: \n        return pd.DataFrame(columns=[\"video_id\",\"agent_id\",\"target_id\",\"action\",\"start_frame\",\"stop_frame\"])\n    mask[-1] = False\n\n    sub = pd.DataFrame({\n        \"video_id\": meta_changes[\"video_id\"][mask].values,\n        \"agent_id\": meta_changes[\"agent_id\"][mask].values,\n        \"target_id\": meta_changes[\"target_id\"][mask].values,\n        \"action\": pred_smoothed.columns[ama_changes[mask].values],\n        \"start_frame\": ama_changes.index[mask],\n        \"stop_frame\": ama_changes.index[1:][mask[:-1]]\n    })\n\n    stop_vid = meta_changes[\"video_id\"][1:][mask[:-1]].values\n    stop_agent = meta_changes[\"agent_id\"][1:][mask[:-1]].values\n    stop_target = meta_changes[\"target_id\"][1:][mask[:-1]].values\n    for i in range(len(sub)):\n        vid = sub.video_id.iloc[i]; ag = sub.agent_id.iloc[i]; tg = sub.target_id.iloc[i]\n        if i < len(stop_vid):\n            if (stop_vid[i] != vid) or (stop_agent[i] != ag) or (stop_target[i] != tg):\n                new_stop = meta_df.query(\"(video_id == @vid)\").video_frame.max() + 1\n                sub.iat[i, sub.columns.get_loc(\"stop_frame\")] = new_stop\n        else:\n            new_stop = meta_df.query(\"(video_id == @vid)\").video_frame.max() + 1\n            sub.iat[i, sub.columns.get_loc(\"stop_frame\")] = new_stop\n\n    dur = sub.stop_frame - sub.start_frame\n    sub = sub[dur >= 3].reset_index(drop=True)\n    if len(sub) > 0:\n        assert (sub.stop_frame > sub.start_frame).all()\n    return sub\n\ndef robustify(submission: pd.DataFrame, dataset: pd.DataFrame, traintest: str, traintest_directory=None):\n    if traintest_directory is None:\n        traintest_directory = f\"{CFG.INPUT_DIR}/{traintest}_tracking\"\n    submission = submission[submission.start_frame < submission.stop_frame]\n    group_list = []\n    for _, g in submission.groupby([\"video_id\",\"agent_id\",\"target_id\"]):\n        g = g.sort_values(\"start_frame\")\n        mask = np.ones(len(g), dtype=bool); last_stop = -1\n        for i, (_, row) in enumerate(g.iterrows()):\n            if row[\"start_frame\"] < last_stop: mask[i] = False\n            else: last_stop = row[\"stop_frame\"]\n        group_list.append(g[mask])\n    submission = pd.concat(group_list) if group_list else submission\n\n    s_list = []\n    for _, row in dataset.iterrows():\n        lab_id = row[\"lab_id\"]; video_id = row[\"video_id\"]\n        if (submission.video_id == video_id).any(): continue\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n        vb = json.loads(row[\"behaviors_labeled\"])\n        vb = sorted(list({b.replace(\"'\", \"\") for b in vb}))\n        vb = [b.split(\",\") for b in vb]\n        vb = pd.DataFrame(vb, columns=[\"agent\",\"target\",\"action\"])\n        start_frame = vid.video_frame.min(); stop_frame  = vid.video_frame.max() + 1\n        for (agent, target), actions in vb.groupby([\"agent\",\"target\"]):\n            batch_len = int(np.ceil((stop_frame - start_frame) / len(actions)))\n            for i, (_, arow) in enumerate(actions.iterrows()):\n                b_start = start_frame + i * batch_len\n                b_stop  = min(b_start + batch_len, stop_frame)\n                s_list.append((video_id, agent, target, arow[\"action\"], b_start, b_stop))\n    if len(s_list) > 0:\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    return submission.reset_index(drop=True)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T12:44:24.382957Z","iopub.execute_input":"2025-10-25T12:44:24.383174Z","iopub.status.idle":"2025-10-25T12:44:24.402097Z","shell.execute_reply.started":"2025-10-25T12:44:24.383159Z","shell.execute_reply":"2025-10-25T12:44:24.401416Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ----- model I/O for inference -----\ndef _read_feature_columns(cols_path: Path):\n    with open(cols_path, \"r\") as f:\n        cols = [line.strip() for line in f.readlines() if len(line.strip())]\n    return cols\n\ndef _align_features(X_df: pd.DataFrame, cols_order):\n    for c in cols_order:\n        if c not in X_df.columns:\n            X_df[c] = 0.0\n    return X_df[cols_order]\n\ndef _load_xgb(path: Path):\n    from xgboost import XGBClassifier\n    clf = XGBClassifier()\n    try:\n        clf.load_model(str(path)); return clf\n    except Exception:\n        return joblib.load(path)\n\ndef _load_cat(path: Path):\n    try:\n        from catboost import CatBoostClassifier\n        model = CatBoostClassifier(); model.load_model(str(path)); return model\n    except Exception:\n        return joblib.load(path)\n\ndef _load_lgbm(path: Path):\n    try:\n        booster = lgb.Booster(model_file=str(path))\n        return (\"booster\", booster)\n    except Exception:\n        model = joblib.load(path)\n        return (\"sk\", model)\n\ndef _predict_proba_any(model_entry, X_np: np.ndarray):\n    if isinstance(model_entry, tuple) and model_entry[0] in (\"booster\",\"sk\"):\n        kind, obj = model_entry\n        if kind == \"booster\":\n            p = obj.predict(X_np)\n            return p[:,1] if (p.ndim == 2 and p.shape[1] == 2) else p\n        else:\n            try:\n                return obj.predict_proba(X_np)[:,1]\n            except Exception:\n                p = obj.predict(X_np)\n                return p[:,1] if (p.ndim == 2 and p.shape[1] == 2) else p\n    try:\n        return model_entry.predict_proba(X_np)[:,1]\n    except Exception:\n        p = model_entry.predict(X_np)\n        return p[:,1] if (p.ndim == 2 and p.shape[1] == 2) else p\n\ndef _tta_predict_models(models, X_np: np.ndarray, n: int, std: float, seed: int):\n    if n <= 0 or std <= 0:\n        probs = [ _predict_proba_any(m, X_np) for m in models ]\n        return np.mean(np.stack(probs, 0), 0)\n    rng = np.random.default_rng(seed)\n    outs = []\n    for _ in range(n):\n        Xn = X_np + rng.normal(0.0, std, size=X_np.shape).astype(np.float32)\n        probs = [ _predict_proba_any(m, Xn) for m in models ]\n        outs.append(np.mean(np.stack(probs, 0), 0))\n    return np.mean(np.stack(outs, 0), 0)\n\ndef _gather_models(action_dir: Path):\n    models = []\n    for k in range(CFG.N_FOLDS_XGB):\n        p = action_dir / f\"xgb_fold{k}.json\"\n        if p.exists() and XGBOOST_AVAILABLE: models.append(_load_xgb(p))\n    p = action_dir / \"xgb_full.json\"\n    if p.exists() and XGBOOST_AVAILABLE: models.append(_load_xgb(p))\n    p = action_dir / \"cat_full.cbm\"\n    if p.exists() and CATBOOST_AVAILABLE: models.append(_load_cat(p))\n    for i in range(1, 4):\n        p = action_dir / f\"lgbm_full_{i}.txt\"\n        if p.exists(): models.append(_load_lgbm(p))\n    return models\n\ndef inference_for_section_mode(section_idx: int, bpt_str: str, mode: str, submission_list: list):\n    bslug = bpt_slug(bpt_str)\n    meta_path = CFG.META_DIR / f\"{bslug}_{mode}.json\"\n    if not meta_path.exists():\n        print(f\"  [WARN] meta not found for {bslug} {mode}, skip.\")\n        return\n    with open(meta_path, \"r\") as f:\n        meta_info = json.load(f)\n    body_parts = meta_info[\"body_parts_tracked\"]\n    actions     = meta_info[\"actions\"]\n\n    test_subset = test[test.body_parts_tracked == bpt_str]\n    if len(test_subset) == 0:\n        print(\"  -> No test rows.\"); return\n    fps_lookup = (\n        test_subset[[\"video_id\",\"frames_per_second\"]]\n        .drop_duplicates(\"video_id\").set_index(\"video_id\")[\"frames_per_second\"].to_dict()\n    )\n\n    # load models & feature columns per action\n    action_models, action_cols = {}, {}\n    for action in actions:\n        action_dir = CFG.MODEL_DIR / bslug / mode / action\n        cols_file  = action_dir / \"feature_columns.txt\"\n        if not (action_dir.exists() and cols_file.exists()): continue\n        models = _gather_models(action_dir)\n        if len(models) == 0: continue\n        action_models[action] = models\n        action_cols[action]   = _read_feature_columns(cols_file)\n    if len(action_models) == 0:\n        print(\"  -> No models found for this section/mode\"); return\n\n    gen = generate_mouse_data(test_subset, \"test\",\n                              generate_single=(mode==\"single\"),\n                              generate_pair=(mode==\"pair\"),\n                              verbose=False)\n    for switch_te, data_te, meta_te, actions_te in gen:\n        if switch_te != mode: continue\n        try:\n            fps_i = _fps_from_meta(meta_te, fps_lookup, default_fps=30.0)\n            if mode == \"single\":\n                X_te = transform_single(data_te, body_parts, fps_i).astype(np.float32)\n            else:\n                X_te = transform_pair(data_te, body_parts, fps_i).astype(np.float32)\n\n            pred_df = pd.DataFrame(index=meta_te.video_frame)\n            X_cache = {}\n            for action in actions_te:\n                if action not in action_models: continue\n                cols = action_cols[action]\n                if action not in X_cache:\n                    Xa = _align_features(X_te.copy(), cols)\n                    X_cache[action] = Xa.to_numpy(np.float32, copy=False)\n                models = action_models[action]\n                probs = _tta_predict_models(models, X_cache[action], CFG.TTA_N, CFG.TTA_NOISE_STD, CFG.SEED+123)\n                pred_df[action] = probs\n\n            if pred_df.shape[1] > 0:\n                sub_part = predict_multiclass_adaptive(pred_df, meta_te, action_thresholds)\n                if len(sub_part): submission_list.append(sub_part)\n\n            del X_te, data_te; gc.collect()\n        except Exception as e:\n            print(f\"  [ERR] inference error: {str(e)[:120]}\")\n            try: del data_te\n            except: pass\n            gc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T12:44:24.402745Z","iopub.execute_input":"2025-10-25T12:44:24.402947Z","iopub.status.idle":"2025-10-25T12:44:24.422027Z","shell.execute_reply.started":"2025-10-25T12:44:24.402930Z","shell.execute_reply":"2025-10-25T12:44:24.421366Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---------------------------\n# 9) MAIN (single)\n# ---------------------------\nif __name__ == \"__main__\":\n    print(\"CFG:\", {k:getattr(CFG,k) for k in dir(CFG) if k.isupper()})\n\n    if CFG.VIS:\n        run_eda()\n\n    if CFG.TRAIN:\n        print(\"\\n=== TRAIN START ===\")\n        for section in range(1, len(body_parts_tracked_list)):\n            bpt_str = body_parts_tracked_list[section]\n            try:\n                body_parts = json.loads(bpt_str)\n                print(f\">> Section {section}: {len(body_parts)} body parts\")\n            except Exception:\n                print(f\">> Section {section}: [invalid body_parts json] -> skip\")\n                continue\n            train_for_section(section, bpt_str); gc.collect()\n        print(\"=== TRAIN DONE ===\\n\")\n\n    if CFG.SUB:\n        print(\"\\n=== INFERENCE / SUB START ===\")\n        submission_list = []\n        for section in range(1, len(body_parts_tracked_list)):\n            bpt_str = body_parts_tracked_list[section]\n            try:\n                _ = json.loads(bpt_str)\n            except Exception:\n                print(f\">> Section {section}: invalid body_parts json -> skip\")\n                continue\n            print(f\">> Section {section}: ensemble inference\")\n            for mode in (\"single\",\"pair\"):\n                print(f\"   - mode={mode}\")\n                inference_for_section_mode(section, bpt_str, mode, submission_list)\n                gc.collect()\n\n        if len(submission_list) > 0:\n            submission = pd.concat(submission_list, ignore_index=True)\n        else:\n            submission = pd.DataFrame({\n                \"video_id\":[438887472],\n                \"agent_id\":[\"mouse1\"],\n                \"target_id\":[\"self\"],\n                \"action\":[\"rear\"],\n                \"start_frame\":[278],\n                \"stop_frame\":[500],\n            })\n\n        submission_robust = robustify(submission, test, \"test\")\n        submission_robust.index.name = \"row_id\"\n        submission_robust.to_csv(CFG.SUB_PATH)\n        print(f\"\\nSubmission created: {len(submission_robust)} rows -> {CFG.SUB_PATH}\")\n        print(submission_robust.head())\n        print(\"=== INFERENCE / SUB DONE ===\")\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-25T12:44:24.423019Z","iopub.execute_input":"2025-10-25T12:44:24.423173Z","iopub.status.idle":"2025-10-25T12:48:35.347114Z","shell.execute_reply.started":"2025-10-25T12:44:24.423161Z","shell.execute_reply":"2025-10-25T12:48:35.346463Z"}},"outputs":[],"execution_count":null}]}