{"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,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =========================================================================================\n# MABe Challenge: Social Action Recognition (Fixed v19)\n# =========================================================================================\nimport sys\nimport os\nimport gc\nimport re\nimport ast\nimport json\nimport joblib\nimport itertools\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport xgboost as xgb\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nfrom scipy.ndimage import gaussian_filter1d, binary_closing\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom sklearn.metrics import f1_score\n\n# -----------------------------------------------------------------------------------------\n# 1. CONFIGURATION & SETUP\n# -----------------------------------------------------------------------------------------\n# Attempt offline install, but fallback gracefully if dataset is missing\nif os.path.exists(\"/kaggle/input/mabe-package\"):\n    !pip install -q --no-index --find-links=/kaggle/input/mabe-package xgboost==3.1.1\nelse:\n    print(\"⚠️ 'mabe-package' not found. Using pre-installed XGBoost.\")\n\nclass Config:\n    seed = 42\n    n_folds = 3\n    negative_ratio = 10\n    max_samples = float('inf') # Process everything\n    \n    # Paths\n    input_dir = Path(\"/kaggle/input/MABe-mouse-behavior-detection\")\n    train_trk = input_dir / \"train_tracking\"\n    train_ant = input_dir / \"train_annotation\"\n    test_trk = input_dir / \"test_tracking\"\n    work_dir = Path(\"/kaggle/working\")\n    \n    # Behaviors\n    body_parts = [\"ear_left\", \"ear_right\", \"nose\", \"neck\", \"body_center\",\n                  \"lateral_left\", \"lateral_right\", \"hip_left\", \"hip_right\", \"tail_base\", \"tail_tip\"]\n    \n    self_behaviors = [\"biteobject\", \"climb\", \"dig\", \"exploreobject\", \"freeze\", \n                      \"genitalgroom\", \"huddle\", \"rear\", \"rest\", \"run\", \"selfgroom\"]\n    \n    pair_behaviors = [\"allogroom\", \"approach\", \"attack\", \"attemptmount\", \"avoid\", \"chase\",\n                      \"chaseattack\", \"defend\", \"disengage\", \"dominance\", \"dominancegroom\",\n                      \"dominancemount\", \"ejaculate\", \"escape\", \"flinch\", \"follow\", \"intromit\",\n                      \"mount\", \"reciprocalsniff\", \"shepherd\", \"sniff\", \"sniffbody\",\n                      \"sniffface\", \"sniffgenital\", \"submit\", \"tussle\"]\n\n    # XGBoost Params (The 0.44 Settings)\n    xgb_params = {\n        \"objective\": \"binary:logistic\",\n        \"eval_metric\": \"logloss\",\n        \"tree_method\": \"hist\",\n        \"learning_rate\": 0.05,\n        \"max_depth\": 6,\n        \"min_child_weight\": 5,\n        \"subsample\": 0.8,\n        \"colsample_bytree\": 0.8,\n        \"max_bin\": 64,\n        \"seed\": 42\n    }\n\n# Ensure directories exist\n(Config.work_dir / \"models\").mkdir(parents=True, exist_ok=True)\n(Config.work_dir / \"features\").mkdir(parents=True, exist_ok=True)\n\n# -----------------------------------------------------------------------------------------\n# 2. FEATURE ENGINEERING ENGINE\n# -----------------------------------------------------------------------------------------\nclass FeatureEngineer:\n    \"\"\"Encapsulates all logic for generating behavioral features.\"\"\"\n    \n    @staticmethod\n    def _cast_float(df):\n        return df.with_columns(pl.col(c).cast(pl.Float32) for c in df.columns if c not in [\"video_frame\", \"mouse_id\", \"bodypart\"])\n\n    @staticmethod\n    def _dist(p1, b1, p2, b2, scale):\n        return ((pl.col(f\"{p1}_x_{b1}\") - pl.col(f\"{p2}_x_{b2}\")).pow(2) + \n                (pl.col(f\"{p1}_y_{b1}\") - pl.col(f\"{p2}_y_{b2}\")).pow(2)).sqrt() / scale\n\n    @staticmethod\n    def _speed(p, b, fps, ms, scale):\n        w = max(1, int(round(ms * fps / 1000.0)))\n        return (((pl.col(f\"{p}_x_{b}\").diff()).pow(2) + \n                 (pl.col(f\"{p}_y_{b}\").diff()).pow(2)).sqrt() / scale * fps).rolling_mean(w, center=True)\n\n    @staticmethod\n    def make_self(meta, tracking):\n        tracking = FeatureEngineer._cast_float(tracking)\n        scale, fps = meta[\"pix_per_cm_approx\"], meta[\"frames_per_second\"]\n        \n        # Pivot\n        start, end = tracking[\"video_frame\"].min(), tracking[\"video_frame\"].max()\n        pivot = tracking.pivot(on=[\"bodypart\"], index=[\"video_frame\", \"mouse_id\"], values=[\"x\", \"y\"]).sort([\"mouse_id\", \"video_frame\"])\n        \n        n_mice = sum(1 for i in range(1,5) if f\"mouse{i}_strain\" in meta and meta[f\"mouse{i}_strain\"] is not None)\n        track_map = {m: pivot.filter(pl.col(\"mouse_id\")==m) for m in range(1, n_mice+1)}\n        \n        results = []\n        vid_id = str(meta[\"video_id\"])\n        \n        for agent in range(1, n_mice+1):\n            if agent not in track_map: continue\n            \n            # Base Frame\n            base = pl.DataFrame({\n                \"video_id\": vid_id, \"agent_mouse_id\": agent, \"target_mouse_id\": -1,\n                \"video_frame\": pl.arange(start, end+1, eager=True).cast(pl.Int32)\n            }).with_columns([pl.col(\"video_id\").cast(pl.Utf8), pl.col(\"agent_mouse_id\").cast(pl.Int32), pl.col(\"target_mouse_id\").cast(pl.Int32)])\n            \n            # Agent Data\n            p_data = track_map[agent].select(pl.col(\"video_frame\"), pl.exclude(\"video_frame\").name.prefix(\"agent_\"))\n            \n            # Fill Missing\n            for bp in Config.body_parts:\n                if f\"agent_x_{bp}\" not in p_data.columns: \n                    p_data = p_data.with_columns([pl.lit(None, pl.Float32).alias(f\"agent_x_{bp}\"), pl.lit(None, pl.Float32).alias(f\"agent_y_{bp}\")])\n\n            # Features\n            feats = p_data.with_columns([\n                *[FeatureEngineer._dist(\"agent\", b1, \"agent\", b2, scale).alias(f\"aa_{b1}_{b2}_dist\") \n                  for b1, b2 in itertools.combinations(Config.body_parts, 2)],\n                *[FeatureEngineer._speed(\"agent\", bp, fps, t, scale).alias(f\"ag_{bp}_spd_{t}\") \n                  for bp, t in itertools.product([\"ear_left\", \"tail_base\"], [500, 1000])]\n            ]).select(pl.exclude([c for c in p_data.columns if \"_x_\" in c or \"_y_\" in c]))\n            \n            results.append(base.join(feats, on=\"video_frame\", how=\"left\"))\n            \n        return pl.concat(results, how=\"vertical\") if results else pl.DataFrame()\n\n    @staticmethod\n    def make_pair(meta, tracking):\n        tracking = FeatureEngineer._cast_float(tracking)\n        scale, fps = meta[\"pix_per_cm_approx\"], meta[\"frames_per_second\"]\n        \n        start, end = tracking[\"video_frame\"].min(), tracking[\"video_frame\"].max()\n        pivot = tracking.pivot(on=[\"bodypart\"], index=[\"video_frame\", \"mouse_id\"], values=[\"x\", \"y\"]).sort([\"mouse_id\", \"video_frame\"])\n        n_mice = sum(1 for i in range(1,5) if f\"mouse{i}_strain\" in meta and meta[f\"mouse{i}_strain\"] is not None)\n        track_map = {m: pivot.filter(pl.col(\"mouse_id\")==m) for m in range(1, n_mice+1)}\n        \n        results = []\n        vid_id = str(meta[\"video_id\"])\n        \n        for ag, tg in itertools.permutations(range(1, n_mice+1), 2):\n            if ag not in track_map or tg not in track_map: continue\n            \n            base = pl.DataFrame({\n                \"video_id\": vid_id, \"agent_mouse_id\": ag, \"target_mouse_id\": tg,\n                \"video_frame\": pl.arange(start, end+1, eager=True).cast(pl.Int32)\n            }).with_columns([pl.col(\"video_id\").cast(pl.Utf8), pl.col(\"agent_mouse_id\").cast(pl.Int32), pl.col(\"target_mouse_id\").cast(pl.Int32)])\n            \n            a_dat = track_map[ag].select(pl.col(\"video_frame\"), pl.exclude(\"video_frame\").name.prefix(\"agent_\"))\n            t_dat = track_map[tg].select(pl.col(\"video_frame\"), pl.exclude(\"video_frame\").name.prefix(\"target_\"))\n            merged = a_dat.join(t_dat, on=\"video_frame\", how=\"inner\")\n            \n            # Features\n            feats = merged.with_columns([\n                # 1. Distances\n                *[FeatureEngineer._dist(\"agent\", b1, \"target\", b2, scale).alias(f\"at_{b1}_{b2}_dist\") \n                  for b1, b2 in itertools.product([\"nose\", \"body_center\", \"tail_base\"], repeat=2)],\n                \n                # 2. Rolling Aggregates (THE SECRET SAUCE)\n                *[FeatureEngineer._dist(\"agent\", \"body_center\", \"target\", \"body_center\", scale)\n                  .rolling_mean(int(round(t*fps/1000)), center=True).alias(f\"dist_mean_{t}\") for t in [500, 1000]],\n                *[FeatureEngineer._dist(\"agent\", \"body_center\", \"target\", \"body_center\", scale)\n                  .rolling_std(int(round(t*fps/1000)), center=True).alias(f\"dist_std_{t}\") for t in [500, 1000]],\n                *[FeatureEngineer._dist(\"agent\", \"body_center\", \"target\", \"body_center\", scale)\n                  .rolling_min(int(round(t*fps/1000)), center=True).alias(f\"dist_min_{t}\") for t in [500, 1000]],\n                \n                # 3. Speeds\n                *[FeatureEngineer._speed(\"agent\", bp, fps, t, scale).alias(f\"ag_{bp}_spd_{t}\") \n                  for bp, t in itertools.product([\"ear_left\", \"tail_base\"], [500, 1000])],\n                *[FeatureEngineer._speed(\"target\", bp, fps, t, scale).alias(f\"tg_{bp}_spd_{t}\") \n                  for bp, t in itertools.product([\"ear_left\", \"tail_base\"], [500, 1000])]\n            ]).select(pl.exclude([c for c in merged.columns if \"_x_\" in c or \"_y_\" in c]))\n            \n            results.append(base.join(feats, on=\"video_frame\", how=\"left\"))\n            \n        return pl.concat(results, how=\"vertical\") if results else pl.DataFrame()\n\n# -----------------------------------------------------------------------------------------\n# 3. TRAINER CLASS (FIXED)\n# -----------------------------------------------------------------------------------------\nclass Trainer:\n    def __init__(self):\n        self.models = {}\n        \n    def tune_threshold(self, preds, y_true):\n        best_th, best_f1 = 0.5, 0.0\n        for th in np.arange(0.1, 0.9, 0.05):\n            score = f1_score(y_true, (preds >= th), zero_division=0)\n            if score > best_f1: best_f1, best_th = score, th\n        return best_th\n\n    def train_behavior(self, lab_id, behavior, indices, features, labels):\n        # Setup Output Dir\n        out_dir = Config.work_dir / \"models\" / lab_id / behavior\n        out_dir.mkdir(parents=True, exist_ok=True)\n        \n        # --- FIXED POLARS CHECK HERE ---\n        # Polars returns a DF on sum(), so we need .item() to get the scalar 0 or >0\n        if labels.select(pl.sum(\"label\")).item() == 0:\n            return 0.0\n        \n        kf = StratifiedGroupKFold(n_splits=Config.n_folds, shuffle=True, random_state=Config.seed)\n        oof = np.zeros(len(labels), dtype=np.float32)\n        \n        # To Pandas for XGBoost\n        X = features.to_pandas()\n        y = labels.to_pandas()\n        groups = indices.get_column(\"video_id\").to_numpy()\n        \n        for fold, (t_idx, v_idx) in enumerate(kf.split(X, y, groups=groups)):\n            X_tr, y_tr = X.iloc[t_idx], y.iloc[t_idx]\n            X_val, y_val = X.iloc[v_idx], y.iloc[v_idx]\n            \n            # Scale Pos Weight\n            # Use iloc[0] to access value safely if it's a dataframe wrapper\n            pos_count = y_tr.sum().iloc[0] if hasattr(y_tr.sum(), 'iloc') else y_tr.sum()\n            ratio = (len(y_tr) - pos_count) / (pos_count + 1e-6)\n            \n            params = Config.xgb_params.copy()\n            params[\"scale_pos_weight\"] = min(ratio, 100.0)\n            \n            dtrain = xgb.DMatrix(X_tr, label=y_tr)\n            dvalid = xgb.DMatrix(X_val, label=y_val)\n            \n            model = xgb.train(\n                params, dtrain, num_boost_round=1000,\n                evals=[(dtrain, \"train\"), (dvalid, \"valid\")],\n                early_stopping_rounds=50, verbose_eval=False\n            )\n            \n            # Predict & Threshold\n            p_val = model.predict(dvalid)\n            th = self.tune_threshold(p_val, y_val)\n            oof[v_idx] = p_val\n            \n            # Save\n            model.save_model(out_dir / f\"fold_{fold}.json\")\n            with open(out_dir / f\"fold_{fold}_th.txt\", \"w\") as f: f.write(str(th))\n            \n        return f1_score(y, (oof >= 0.5).astype(int))\n\n# -----------------------------------------------------------------------------------------\n# 4. EXECUTION PIPELINE\n# -----------------------------------------------------------------------------------------\ndef run_training():\n    print(\">>> 1. Loading Training Metadata...\")\n    train_df = pl.read_csv(Config.input_dir / \"train.csv\")\n    train_beh = (\n        train_df.filter(pl.col(\"behaviors_labeled\").is_not_null())\n        .select(pl.col(\"lab_id\"), pl.col(\"video_id\"), pl.col(\"behaviors_labeled\").map_elements(eval, return_dtype=pl.List(pl.Utf8)).alias(\"b\"))\n        .explode(\"b\")\n        .select(\n            pl.col(\"lab_id\"), pl.col(\"video_id\").cast(pl.Utf8),\n            pl.col(\"b\").str.split(\",\").list[0].str.replace_all(\"'\",\"\").alias(\"agent\"),\n            pl.col(\"b\").str.split(\",\").list[1].str.replace_all(\"'\",\"\").alias(\"target\"),\n            pl.col(\"b\").str.split(\",\").list[2].str.replace_all(\"'\",\"\").alias(\"action\")\n        )\n    )\n    \n    # --- A. FEATURE GENERATION (SERIAL) ---\n    print(\">>> 2. Generating Features (Serial for Safety)...\")\n    videos = list(train_df.filter(pl.col(\"behaviors_labeled\").is_not_null()).rows(named=True))\n    \n    for row in tqdm(videos):\n        vid_id = str(row[\"video_id\"])\n        lab = row[\"lab_id\"]\n        path = Config.train_trk / f\"{lab}/{vid_id}.parquet\"\n        \n        if not path.exists(): continue\n        # Check if already done (if restarting)\n        if (Config.work_dir / \"features\" / f\"s_{vid_id}.parquet\").exists(): continue\n        \n        try:\n            trk = pl.read_parquet(path)\n            sf = FeatureEngineer.make_self(row, trk)\n            pf = FeatureEngineer.make_pair(row, trk)\n            \n            if sf.height > 0: sf.write_parquet(Config.work_dir / \"features\" / f\"s_{vid_id}.parquet\")\n            if pf.height > 0: pf.write_parquet(Config.work_dir / \"features\" / f\"p_{vid_id}.parquet\")\n        except: pass\n\n    # --- B. TRAINING ---\n    trainer = Trainer()\n    \n    # 1. Self Behaviors\n    print(\">>> 3. Training Self Behaviors...\")\n    for (lab, action), grp in tqdm(train_beh.filter(pl.col(\"action\").is_in(Config.self_behaviors)).group_by(\"lab_id\", \"action\")):\n        idxs, feats, lbls = [], [], []\n        for r in grp.rows(named=True):\n            vid, ag = str(r[\"video_id\"]), r[\"agent\"]\n            fpath = Config.work_dir / \"features\" / f\"s_{vid}.parquet\"\n            if not fpath.exists(): continue\n            \n            ag_id = int(re.search(r\"mouse(\\d+)\", ag).group(1))\n            dat = pl.scan_parquet(fpath).filter(pl.col(\"agent_mouse_id\")==ag_id)\n            \n            # Load Annotation\n            a_path = Config.train_ant / lab / f\"{vid}.parquet\"\n            ant = pl.scan_parquet(a_path).filter((pl.col(\"action\")==action) & (pl.col(\"agent_id\")==ag_id)).collect() if a_path.exists() else pl.DataFrame()\n            \n            pos_frames = set()\n            for ar in ant.rows(named=True): pos_frames.update(range(ar[\"start_frame\"], ar[\"stop_frame\"]))\n            \n            # Get Data & Label\n            d = dat.collect()\n            l = d.select(pl.col(\"video_frame\").is_in(pos_frames).cast(pl.Int8).alias(\"label\"))\n            \n            # Downsample Negatives\n            pos = l.filter(pl.col(\"label\")==1)\n            neg = l.filter(pl.col(\"label\")==0)\n            \n            if pos.height == 0: continue\n            \n            n_keep = min(neg.height, pos.height * Config.negative_ratio)\n            neg_keep_idx = np.random.choice(np.where(l[\"label\"]==0)[0], n_keep, replace=False)\n            pos_keep_idx = np.where(l[\"label\"]==1)[0]\n            keep = np.sort(np.concatenate([pos_keep_idx, neg_keep_idx]))\n            \n            idxs.append(d[keep].select([\"video_id\"]))\n            feats.append(d[keep].select(pl.exclude([\"video_id\", \"agent_mouse_id\", \"target_mouse_id\", \"video_frame\"])))\n            lbls.append(l[keep])\n            \n        if idxs:\n            trainer.train_behavior(lab, action, pl.concat(idxs), pl.concat(feats), pl.concat(lbls))\n            gc.collect()\n\n    # 2. Pair Behaviors\n    print(\">>> 4. Training Pair Behaviors...\")\n    for (lab, action), grp in tqdm(train_beh.filter(pl.col(\"action\").is_in(Config.pair_behaviors)).group_by(\"lab_id\", \"action\")):\n        idxs, feats, lbls = [], [], []\n        for r in grp.rows(named=True):\n            vid, ag, tg = str(r[\"video_id\"]), r[\"agent\"], r[\"target\"]\n            fpath = Config.work_dir / \"features\" / f\"p_{vid}.parquet\"\n            if not fpath.exists(): continue\n            \n            ag_id = int(re.search(r\"mouse(\\d+)\", ag).group(1))\n            tg_id = int(re.search(r\"mouse(\\d+)\", tg).group(1))\n            \n            dat = pl.scan_parquet(fpath).filter((pl.col(\"agent_mouse_id\")==ag_id) & (pl.col(\"target_mouse_id\")==tg_id))\n            \n            a_path = Config.train_ant / lab / f\"{vid}.parquet\"\n            ant = pl.scan_parquet(a_path).filter((pl.col(\"action\")==action) & (pl.col(\"agent_id\")==ag_id) & (pl.col(\"target_id\")==tg_id)).collect() if a_path.exists() else pl.DataFrame()\n            \n            pos_frames = set()\n            for ar in ant.rows(named=True): pos_frames.update(range(ar[\"start_frame\"], ar[\"stop_frame\"]))\n            \n            d = dat.collect()\n            l = d.select(pl.col(\"video_frame\").is_in(pos_frames).cast(pl.Int8).alias(\"label\"))\n            \n            pos = l.filter(pl.col(\"label\")==1)\n            if pos.height == 0: continue\n            \n            n_keep = min(l.height - pos.height, pos.height * Config.negative_ratio)\n            neg_keep_idx = np.random.choice(np.where(l[\"label\"]==0)[0], n_keep, replace=False)\n            keep = np.sort(np.concatenate([np.where(l[\"label\"]==1)[0], neg_keep_idx]))\n            \n            idxs.append(d[keep].select([\"video_id\"]))\n            feats.append(d[keep].select(pl.exclude([\"video_id\", \"agent_mouse_id\", \"target_mouse_id\", \"video_frame\"])))\n            lbls.append(l[keep])\n            \n        if idxs:\n            trainer.train_behavior(lab, action, pl.concat(idxs), pl.concat(feats), pl.concat(lbls))\n            gc.collect()\n\ndef run_inference():\n    print(\">>> 5. Starting Inference...\")\n    test_df = pl.read_csv(Config.input_dir / \"test.csv\")\n    preds = []\n    \n    # Preload Models\n    models = {} # {lab: {beh: [(model, th), ...]}}\n    for lab_dir in (Config.work_dir / \"models\").glob(\"*\"):\n        lab = lab_dir.name\n        models[lab] = {}\n        for beh_dir in lab_dir.glob(\"*\"):\n            beh = beh_dir.name\n            fold_models = []\n            for fold_f in beh_dir.glob(\"*.json\"):\n                try:\n                    bst = xgb.Booster()\n                    bst.load_model(str(fold_f))\n                    with open(str(fold_f).replace(\".json\", \"_th.txt\")) as f: th = float(f.read().strip())\n                    fold_models.append((bst, bst.feature_names, th))\n                except: pass\n            if fold_models: models[lab][beh] = fold_models\n            \n    for row in tqdm(test_df.rows(named=True)):\n        vid = str(row[\"video_id\"])\n        lab = row[\"lab_id\"]\n        path = Config.test_trk / f\"{lab}/{vid}.parquet\"\n        if not path.exists() or lab not in models: continue\n        \n        # Generate on Fly\n        trk = pl.read_parquet(path)\n        \n        # Self\n        sf = FeatureEngineer.make_self(row, trk)\n        if sf.height > 0:\n            pdf = sf.to_pandas()\n            for c in pdf.columns: \n                if pdf[c].dtype=='object': pdf[c]=pdf[c].astype('float32')\n                \n            for ag_id in pdf[\"agent_mouse_id\"].unique():\n                sub = pdf[pdf[\"agent_mouse_id\"]==ag_id]\n                frames = sub[\"video_frame\"].values\n                for beh in Config.self_behaviors:\n                    if beh not in models[lab]: continue\n                    ensemble = models[lab][beh]\n                    req_col = ensemble[0][1]\n                    \n                    # Align cols\n                    missing = [c for c in req_col if c not in sub.columns]\n                    if missing: \n                        for c in missing: sub[c] = np.nan\n                    \n                    dtest = xgb.DMatrix(sub[req_col])\n                    avg_p, avg_t = 0, 0\n                    for m, _, t in ensemble:\n                        avg_p += m.predict(dtest)\n                        avg_t += t\n                    \n                    avg_p /= len(ensemble)\n                    avg_t /= len(ensemble)\n                    \n                    # Post Process\n                    avg_p = gaussian_filter1d(avg_p, 2.0)\n                    active = binary_closing(avg_p > (avg_t * 0.95), structure=[1]*10)\n                    \n                    for idx in np.where(active)[0]:\n                        preds.append((vid, f\"mouse{ag_id}\", \"self\", beh, frames[idx]))\n                        \n        # Pair\n        pf = FeatureEngineer.make_pair(row, trk)\n        if pf.height > 0:\n            pdf = pf.to_pandas()\n            for c in pdf.columns: \n                if pdf[c].dtype=='object': pdf[c]=pdf[c].astype('float32')\n                \n            for ag_id in pdf[\"agent_mouse_id\"].unique():\n                for tg_id in pdf[pdf[\"agent_mouse_id\"]==ag_id][\"target_mouse_id\"].unique():\n                    sub = pdf[(pdf[\"agent_mouse_id\"]==ag_id) & (pdf[\"target_mouse_id\"]==tg_id)]\n                    frames = sub[\"video_frame\"].values\n                    for beh in Config.pair_behaviors:\n                        if beh not in models[lab]: continue\n                        ensemble = models[lab][beh]\n                        req_col = ensemble[0][1]\n                        \n                        missing = [c for c in req_col if c not in sub.columns]\n                        if missing: \n                            for c in missing: sub[c] = np.nan\n                        \n                        dtest = xgb.DMatrix(sub[req_col])\n                        avg_p, avg_t = 0, 0\n                        for m, _, t in ensemble:\n                            avg_p += m.predict(dtest)\n                            avg_t += t\n                        \n                        avg_p /= len(ensemble)\n                        avg_t /= len(ensemble)\n                        \n                        avg_p = gaussian_filter1d(avg_p, 2.0)\n                        active = binary_closing(avg_p > (avg_t * 0.95), structure=[1]*10)\n                        \n                        for idx in np.where(active)[0]:\n                            preds.append((vid, f\"mouse{ag_id}\", f\"mouse{tg_id}\", beh, frames[idx]))\n        gc.collect()\n\n    # Save\n    if not preds:\n        pl.DataFrame({'video_id': [0], 'agent_id':['m1'], 'target_id':['m2'], 'action':['x'], 'start_frame':[0], 'stop_frame':[10]}).write_csv(Config.work_dir / \"submission.csv\")\n    else:\n        df = pl.DataFrame(preds, schema=[\"video_id\", \"agent_id\", \"target_id\", \"action\", \"start_frame\"], orient=\"row\")\n        df = df.sort([\"video_id\", \"agent_id\", \"target_id\", \"start_frame\"])\n        \n        # RLE\n        df = df.with_columns([\n            pl.col(\"video_id\").shift(1).alias(\"pv\"), pl.col(\"agent_id\").shift(1).alias(\"pa\"), \n            pl.col(\"target_id\").shift(1).alias(\"pt\"), pl.col(\"action\").shift(1).alias(\"pact\"), \n            pl.col(\"start_frame\").shift(1).alias(\"pf\")\n        ])\n        df = df.with_columns(\n            ((pl.col(\"video_id\")!=pl.col(\"pv\")) | (pl.col(\"agent_id\")!=pl.col(\"pa\")) | \n             (pl.col(\"target_id\")!=pl.col(\"pt\")) | (pl.col(\"action\")!=pl.col(\"pact\")) | \n             ((pl.col(\"start_frame\")-pl.col(\"pf\"))>5)).fill_null(True).cum_sum().alias(\"seg\")\n        )\n        final = df.group_by([\"seg\", \"video_id\", \"agent_id\", \"target_id\", \"action\"]).agg([\n            pl.col(\"start_frame\").min(), pl.col(\"start_frame\").max().alias(\"stop_frame\")\n        ]).select([\"video_id\", \"agent_id\", \"target_id\", \"action\", \"start_frame\", \"stop_frame\"])\n        \n        final.filter(pl.col(\"stop_frame\") > pl.col(\"start_frame\")).with_row_index(\"row_id\").write_csv(Config.work_dir / \"submission.csv\")\n    \n    print(\">>> 6. DONE.\")\n\n# -----------------------------------------------------------------------------------------\n# 5. ENTRY POINT\n# -----------------------------------------------------------------------------------------\nif __name__ == \"__main__\":\n    run_training()\n    run_inference()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-15T05:48:50.437225Z","iopub.execute_input":"2025-12-15T05:48:50.438141Z"}},"outputs":[],"execution_count":null}]}