{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## I used some more features taken from MABe Nearest Neighbors: Testing New Features section (0.34 -> 0.38)\n\n## References\n- [MABe Nearest Neighbors: Testing New Features](https://www.kaggle.com/code/taylorsamarel/mabe-nearest-neighbors-testing-new-features)\n  \n","metadata":{}},{"cell_type":"code","source":"from __future__ import annotations\nimport os\nimport gc\nimport json\nimport itertools\nimport warnings\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Any, Dict, Iterable, List, Optional, Tuple\n\nimport numpy as np\nimport pandas as pd\n\nfrom sklearn.base import BaseEstimator, ClassifierMixin, clone\nfrom sklearn.impute import SimpleImputer\nfrom sklearn.pipeline import make_pipeline\nfrom sklearn.model_selection import GroupKFold\nfrom sklearn.metrics import f1_score\n\ntry:\n    from lightgbm import LGBMClassifier\nexcept Exception as e:\n    raise RuntimeError(\"LightGBM is required.\") from e\n\n\n# ========================\n# Global Configuration\n# ========================\nvalidate_or_submit = 'submit'\nverbose = True\nUSE_HYSTERESIS_DECODER = False  \n\n# ========================\n# Logging\n# ========================\ndef log(msg: str):\n    if verbose:\n        print(msg, flush=True)\n\n# ========================\n# Enhanced decoder parameters per action\n# ========================\nSTATE_ACTIONS = {\n    'rest', 'sniff', 'sniffface', 'sniffgenital', 'mount', 'intromit', 'reciprocalsniff'\n}\nEVENT_ACTIONS = {\n    'approach', 'avoid', 'follow', 'chase', 'chaseattack', 'attack', 'defend', 'flinch',\n    'tussle', 'dominance', 'escape', 'submit', 'attemptmount'\n}\n\n# Labs with many short labels (1-2 frames)\nSHORT_LABEL_LABS = {'CRIM13', 'CalMS21_task1', 'CalMS21_task2', 'CalMS21_supplemental'}\n\nDEFAULT_PER_ACTION = {\n    'sniff':        dict(t_on=0.55, t_off=0.40, min_len=6,  merge_gap=6,  smooth_k=11),\n    'sniffface':    dict(t_on=0.55, t_off=0.40, min_len=6,  merge_gap=6,  smooth_k=11),\n    'sniffgenital': dict(t_on=0.60, t_off=0.45, min_len=8,  merge_gap=8,  smooth_k=13),\n    'rest':         dict(t_on=0.60, t_off=0.45, min_len=12, merge_gap=8,  smooth_k=13),\n    'mount':        dict(t_on=0.60, t_off=0.45, min_len=8,  merge_gap=6,  smooth_k=11),\n    'intromit':     dict(t_on=0.65, t_off=0.50, min_len=6,  merge_gap=4,  smooth_k=9),\n    'approach':     dict(t_on=0.60, t_off=0.45, min_len=4,  merge_gap=4,  smooth_k=9),\n    'avoid':        dict(t_on=0.60, t_off=0.45, min_len=4,  merge_gap=4,  smooth_k=9),\n    'follow':       dict(t_on=0.60, t_off=0.45, min_len=5,  merge_gap=5,  smooth_k=9),\n    'chase':        dict(t_on=0.60, t_off=0.45, min_len=5,  merge_gap=5,  smooth_k=9),\n    'chaseattack':  dict(t_on=0.62, t_off=0.48, min_len=5,  merge_gap=4,  smooth_k=9),\n    'attack':       dict(t_on=0.65, t_off=0.50, min_len=3,  merge_gap=3,  smooth_k=7),\n    'defend':       dict(t_on=0.62, t_off=0.48, min_len=3,  merge_gap=3,  smooth_k=7),\n    'flinch':       dict(t_on=0.62, t_off=0.48, min_len=2,  merge_gap=2,  smooth_k=7),\n}\n\n# ========================\n# Config\n# ========================\n@dataclass\nclass Config:\n    data_root: Path = Path(\"/kaggle/input/MABe-mouse-behavior-detection\")\n    submission_file: str = \"submission.csv\"\n    row_id_col: str = \"row_id\"\n    max_train_samples_per_action: int = 100_000\n\n    @property\n    def train_csv(self) -> Path: return self.data_root / \"train.csv\"\n    @property\n    def test_csv(self) -> Path: return self.data_root / \"test.csv\"\n    @property\n    def train_track_dir(self) -> Path: return self.data_root / \"train_tracking\"\n    @property\n    def train_ann_dir(self) -> Path: return self.data_root / \"train_annotation\"\n    @property\n    def test_track_dir(self) -> Path: return self.data_root / \"test_tracking\"\n\n\nCFG = Config()\n\n# ========================\n# Body parts to drop\n# ========================\ndrop_body_parts = [\n    'headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft',\n    'headpiece_bottomfrontright', 'headpiece_topbackleft', 'headpiece_topbackright',\n    'headpiece_topfrontleft', 'headpiece_topfrontright', 'spine_1', 'spine_2',\n    'tail_middle_1', 'tail_middle_2', 'tail_midpoint'\n]\n\n# ========================\n# Hysteresis Decoder\n# ========================\ndef _smooth_probs(x, k):\n    if k <= 1: return x\n    pad = k // 2\n    xx = np.pad(x, (pad, pad), mode='edge')\n    w = np.ones(k) / k\n    return np.convolve(xx, w, mode='valid')\n\ndef _decode_one_series(scores_df, actions, per_action):\n    frames = scores_df.index.values\n    current = None\n    events = []\n    last_end = {}\n\n    t_on  = {a: per_action.get(a, {}).get('t_on',  0.55) for a in actions}\n    t_off = {a: per_action.get(a, {}).get('t_off', 0.40) for a in actions}\n    minl  = {a: per_action.get(a, {}).get('min_len', 4)  for a in actions}\n    mgap  = {a: per_action.get(a, {}).get('merge_gap', 4) for a in actions}\n\n    for i, f in enumerate(frames):\n        row = scores_df.iloc[i].values\n        a_idx = int(np.argmax(row))\n        a = actions[a_idx]\n        s = row[a_idx]\n\n        if current is None:\n            if s >= t_on[a]:\n                current = [a, f]\n        else:\n            act_on, st = current\n            if a == act_on and s >= t_off[act_on]:\n                pass\n            else:\n                en = f\n                if en > st and (en - st) >= minl[act_on]:\n                    if act_on in last_end and (st - last_end[act_on][1]) <= mgap[act_on]:\n                        last_end[act_on][1] = en\n                    else:\n                        last_end[act_on] = [st, en]\n                current = None\n                if s >= t_on[a]:\n                    current = [a, f]\n\n    if current is not None:\n        a, st = current\n        en = frames[-1] + 1\n        if (en - st) >= minl[a]:\n            if a in last_end and (st - last_end[a][1]) <= mgap[a]:\n                last_end[a][1] = en\n            else:\n                last_end[a] = [st, en]\n\n    out = []\n    for a, (st, en) in last_end.items():\n        out.append((a, st, en))\n    return out\n\ndef _lab_params(lab_id: str, actions: List[str]) -> dict:\n    per_action = {a: DEFAULT_PER_ACTION.get(a, dict(t_on=0.60, t_off=0.45, min_len=4, merge_gap=4, smooth_k=9))\n                  for a in actions}\n    if lab_id in SHORT_LABEL_LABS:\n        for a in actions:\n            if a in EVENT_ACTIONS:\n                p = per_action[a].copy()\n                p['min_len'] = max(1, min(p.get('min_len', 4), 2))\n                p['smooth_k'] = max(5, p.get('smooth_k', 9) - 2)\n                p['t_on']  = max(0.50, p.get('t_on', 0.60) - 0.05)\n                p['t_off'] = max(0.35, p.get('t_off', 0.45) - 0.05)\n                per_action[a] = p\n    return per_action\n\ndef predict_multiclass_hysteresis(pred: pd.DataFrame, meta: pd.DataFrame) -> pd.DataFrame:\n    \"\"\"Enhanced decoder with hysteresis and lab-aware parameters\"\"\"\n    if pred.shape[1] == 0:\n        return pd.DataFrame(columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n\n    pieces = []\n    gcols = ['video_id', 'agent_id', 'target_id']\n    idx_df = meta.assign(_idx=np.arange(len(meta)))\n    \n    for keys, sub_idx in idx_df.groupby(gcols)['_idx']:\n        ii = sub_idx.values\n        meta_g = meta.iloc[ii]\n        pred_g = pred.iloc[ii].copy()\n        actions = list(pred_g.columns)\n\n        lab = str(meta_g['lab_id'].iloc[0]) if 'lab_id' in meta_g.columns else ''\n        per_action = _lab_params(lab, actions)\n        \n        # Smooth each action\n        sm = {}\n        for a in actions:\n            k = per_action.get(a, {}).get('smooth_k', 9)\n            sm[a] = _smooth_probs(pred_g[a].values, k)\n        sm = pd.DataFrame(sm, index=pred_g.index, columns=actions)\n        \n        evs = _decode_one_series(sm, actions, per_action)\n        if not evs:\n            continue\n            \n        g = meta_g.iloc[0]\n        out = pd.DataFrame(evs, columns=['action', 'start_frame', 'stop_frame'])\n        out.insert(0, 'target_id', g['target_id'])\n        out.insert(0, 'agent_id', g['agent_id'])\n        out.insert(0, 'video_id', g['video_id'])\n        pieces.append(out)\n\n    return pd.concat(pieces, ignore_index=True) if pieces else \\\n        pd.DataFrame(columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n\n# ========================\n# Generate mouse data \n# ========================\ndef generate_mouse_data(dataset: pd.DataFrame,\n                        traintest: str,\n                        traintest_directory: Optional[str] = None,\n                        generate_single: bool = True,\n                        generate_pair: bool = True):\n    assert traintest in ['train', 'test']\n    if traintest_directory is None:\n        traintest_directory = f\"{CFG.data_root}/{traintest}_tracking\"\n\n    for _, row in dataset.iterrows():\n        lab_id = row.lab_id\n        if traintest == 'train' and str(lab_id).startswith('MABe22'):\n            continue\n\n        video_id = row.video_id\n        behaviors = row.behaviors_labeled\n        if type(behaviors) != str:\n            if verbose and traintest == 'test':\n                print('No labeled behaviors:', lab_id, video_id, type(behaviors), behaviors)\n            continue\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n        \n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        if len(np.unique(vid.bodypart)) > 5:\n            try:\n                pvid = pvid.drop(columns=[('x', slice(None), b) for b in drop_body_parts], errors='ignore')\n                pvid = pvid.drop(columns=[('y', slice(None), b) for b in drop_body_parts], errors='ignore')\n            except Exception:\n                pass\n\n        if pvid.isna().any().any():\n            if verbose and traintest == 'test':\n                print('video with missing values', video_id, traintest, len(vid), 'frames')\n        else:\n            if verbose and traintest == 'test':\n                print('video with all values', video_id, traintest, len(vid), 'frames')\n\n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n        pvid = pvid / row.pix_per_cm_approx\n\n        vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in json.loads(row.behaviors_labeled)}))\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        # SINGLE\n        if generate_single:\n            vid_b_single = vid_behaviors.query(\"target == 'self'\")\n            for mouse_id_str in np.unique(vid_b_single.agent):\n                try:\n                    mouse_id = int(mouse_id_str[-1])\n                    vid_agent_actions = np.unique(vid_b_single.query(\"agent == @mouse_id_str\").action)\n                    single_mouse = pvid.loc[:, mouse_id]\n                    single_mouse_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': mouse_id_str,\n                        'target_id': 'self',\n                        'video_frame': single_mouse.index,\n                        'lab_id': lab_id,  # ADDED: lab_id for decoder\n                    })\n                    if traintest == 'train':\n                        single_mouse_label = pd.DataFrame(0.0, columns=vid_agent_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                            r = annot_subset.iloc[i]\n                            single_mouse_label.loc[r['start_frame']:r['stop_frame'], r.action] = 1.0\n                        yield 'single', single_mouse, single_mouse_meta, single_mouse_label\n                    else:\n                        if verbose: print('- test single', video_id, mouse_id)\n                        yield 'single', single_mouse, single_mouse_meta, vid_agent_actions\n                except KeyError:\n                    pass\n\n        # PAIR\n        if generate_pair:\n            vid_b_pair = vid_behaviors.query(\"target != 'self'\")\n            if len(vid_b_pair) > 0:\n                avail = np.unique(pvid.columns.get_level_values('mouse_id'))\n                for agent, target in itertools.permutations(avail, 2):\n                    agent_str = f\"mouse{agent}\"\n                    target_str = f\"mouse{target}\"\n                    vid_agent_actions = np.unique(vid_b_pair.query(\"(agent == @agent_str) & (target == @target_str)\").action)\n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\n                    mouse_pair_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': agent_str,\n                        'target_id': target_str,\n                        'video_frame': mouse_pair.index,\n                        'lab_id': lab_id,  # ADDED: lab_id for decoder\n                    })\n                    if traintest == 'train':\n                        annot_subset = annot.query(\"(agent_id == @agent) & (target_id == @target)\")\n                        mouse_pair_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=mouse_pair.index)\n                        for i in range(len(annot_subset)):\n                            r = annot_subset.iloc[i]\n                            mouse_pair_label.loc[r['start_frame']:r['stop_frame'], r.action] = 1.0\n                        yield 'pair', mouse_pair, mouse_pair_meta, mouse_pair_label\n                    else:\n                        if verbose: print('- test pair', video_id, agent, target)\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n\n# ========================\n# ENHANCED Feature Engineering\n# ========================\ndef transform_single(single_mouse: pd.DataFrame, body_parts_tracked: List[str]) -> pd.DataFrame:\n    \"\"\"Enhanced feature engineering with long-range temporal dependencies\"\"\"\n    available_body_parts = single_mouse.columns.get_level_values(0)\n    \n    # Base pairwise distances\n    feats = {}\n    for part1, part2 in itertools.combinations(body_parts_tracked, 2):\n        if part1 in available_body_parts and part2 in available_body_parts:\n            dif = single_mouse[part1] - single_mouse[part2]\n            feats[f\"{part1}+{part2}\"] = np.square(dif).sum(axis=1, skipna=False)\n\n    X = pd.DataFrame(feats, index=single_mouse.index)\n    \n    # Basic speed features\n    try:\n        if ('ear_left' in available_body_parts) and ('ear_right' in available_body_parts) and ('tail_base' in available_body_parts):\n            shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(10)\n            X2 = pd.DataFrame({\n                'speed_left':  np.square(single_mouse['ear_left']  - shifted['ear_left']).sum(axis=1, skipna=False),\n                'speed_right': np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False),\n                'speed_left2': np.square(single_mouse['ear_left']  - shifted['tail_base']).sum(axis=1, skipna=False),\n                'speed_right2': np.square(single_mouse['ear_right'] - shifted['tail_base']).sum(axis=1, skipna=False),\n            }, index=single_mouse.index)\n            X = pd.concat([X, X2], axis=1)\n    except Exception:\n        pass\n    \n    # LONG-RANGE TEMPORAL DEPENDENCIES (from files 2-3)\n    time_windows = [5, 15, 30, 60]\n    \n    if 'body_center' in available_body_parts:\n        center_x = single_mouse['body_center']['x']\n        center_y = single_mouse['body_center']['y']\n        \n        for window in time_windows:\n            # Rolling statistics\n            X[f'center_x_mean_{window}'] = center_x.rolling(window, min_periods=1, center=True).mean()\n            X[f'center_y_mean_{window}'] = center_y.rolling(window, min_periods=1, center=True).mean()\n            X[f'center_x_std_{window}'] = center_x.rolling(window, min_periods=1, center=True).std()\n            X[f'center_y_std_{window}'] = center_y.rolling(window, min_periods=1, center=True).std()\n            \n            # Movement range\n            X[f'x_range_{window}'] = center_x.rolling(window, min_periods=1, center=True).max() - center_x.rolling(window, min_periods=1, center=True).min()\n            X[f'y_range_{window}'] = center_y.rolling(window, min_periods=1, center=True).max() - center_y.rolling(window, min_periods=1, center=True).min()\n            \n            # Activity level\n            X[f'activity_level_{window}'] = np.sqrt(\n                center_x.diff().rolling(window, min_periods=1).var() + \n                center_y.diff().rolling(window, min_periods=1).var()\n            )\n    \n    # Lag features\n    if 'nose' in available_body_parts and 'tail_base' in available_body_parts:\n        nose_tail_dist = np.sqrt(\n            (single_mouse['nose']['x'] - single_mouse['tail_base']['x'])**2 + \n            (single_mouse['nose']['y'] - single_mouse['tail_base']['y'])**2\n        )\n        \n        for lag in [10, 20, 40]:\n            X[f'nose_tail_dist_lag_{lag}'] = nose_tail_dist.shift(lag)\n            X[f'nose_tail_dist_diff_{lag}'] = nose_tail_dist - nose_tail_dist.shift(lag)\n    \n    return X\n\ndef transform_pair(mouse_pair: pd.DataFrame, body_parts_tracked: List[str]) -> pd.DataFrame:\n    \"\"\"Enhanced feature engineering for pairs with social interaction features\"\"\"\n    available_A = mouse_pair['A'].columns.get_level_values(0)\n    available_B = mouse_pair['B'].columns.get_level_values(0)\n\n    # Cross-mouse pairwise distances\n    feats = {}\n    for p1, p2 in itertools.product(body_parts_tracked, repeat=2):\n        if p1 in available_A and p2 in available_B:\n            dif = mouse_pair['A'][p1] - mouse_pair['B'][p2]\n            feats[f\"12+{p1}+{p2}\"] = np.square(dif).sum(axis=1, skipna=False)\n    X = pd.DataFrame(feats, index=mouse_pair.index)\n\n    # Basic speed features\n    try:\n        if ('A', 'ear_left') in mouse_pair.columns and ('B', 'ear_left') in mouse_pair.columns:\n            shifted_A = mouse_pair['A']['ear_left'].shift(10)\n            shifted_B = mouse_pair['B']['ear_left'].shift(10)\n            X2 = pd.DataFrame({\n                'speed_left_A':  np.square(mouse_pair['A']['ear_left'] - shifted_A).sum(axis=1, skipna=False),\n                'speed_left_AB': np.square(mouse_pair['A']['ear_left'] - shifted_B).sum(axis=1, skipna=False),\n                'speed_left_B':  np.square(mouse_pair['B']['ear_left'] - shifted_B).sum(axis=1, skipna=False),\n            }, index=mouse_pair.index)\n            X = pd.concat([X, X2], axis=1)\n    except Exception:\n        pass\n    \n    # SOCIAL INTERACTION FEATURES (from files 2-3)\n    time_windows = [5, 15, 30, 60]\n    \n    if 'body_center' in available_A and 'body_center' in available_B:\n        center_dist = np.square(mouse_pair['A']['body_center'] - mouse_pair['B']['body_center']).sum(axis=1, skipna=False)\n        \n        # Social zone indicators\n        X['very_close'] = (np.sqrt(center_dist) < 5.0).astype(float)\n        X['close'] = ((np.sqrt(center_dist) >= 5.0) & (np.sqrt(center_dist) < 15.0)).astype(float)\n        \n        for window in time_windows:\n            # Distance statistics\n            X[f'dist_mean_{window}'] = center_dist.rolling(window, min_periods=1, center=True).mean()\n            X[f'dist_std_{window}'] = center_dist.rolling(window, min_periods=1, center=True).std()\n            \n            # Interaction intensity\n            dist_var = center_dist.rolling(window, min_periods=1, center=True).var()\n            X[f'interaction_intensity_{window}'] = 1 / (1 + dist_var)\n            \n            # Coordinated movement\n            A_x_diff = mouse_pair['A']['body_center']['x'].diff()\n            A_y_diff = mouse_pair['A']['body_center']['y'].diff()\n            B_x_diff = mouse_pair['B']['body_center']['x'].diff()\n            B_y_diff = mouse_pair['B']['body_center']['y'].diff()\n            \n            coord_movement = A_x_diff * B_x_diff + A_y_diff * B_y_diff\n            X[f'coord_movement_mean_{window}'] = coord_movement.rolling(window, min_periods=1, center=True).mean()\n    \n    # Lag features for interactions\n    if 'nose' in available_A and 'nose' in available_B:\n        nose_nose_dist = np.sqrt(\n            (mouse_pair['A']['nose']['x'] - mouse_pair['B']['nose']['x'])**2 +\n            (mouse_pair['A']['nose']['y'] - mouse_pair['B']['nose']['y'])**2\n        )\n        \n        for lag in [10, 20, 40]:\n            X[f'nose_nose_dist_lag_{lag}'] = nose_nose_dist.shift(lag)\n            X[f'nose_nose_dist_change_{lag}'] = nose_nose_dist - nose_nose_dist.shift(lag)\n            \n            # Interaction persistence\n            close_threshold = 10.0\n            is_close = (nose_nose_dist < close_threshold).astype(float)\n            X[f'close_persistence_{lag}'] = is_close.rolling(lag, min_periods=1).mean()\n    \n    return X\n\n# ========================\n# TrainOnSubsetClassifier\n# ========================\nclass TrainOnSubsetClassifier(BaseEstimator, ClassifierMixin):\n    def __init__(self, model: Any, max_samples: int = 100_000, random_state: int = 42):\n        self.model = model\n        self.max_samples = int(max_samples)\n        self.random_state = random_state\n\n    def fit(self, X, y):\n        n = X.shape[0]\n        if self.max_samples and n > self.max_samples:\n            rng = np.random.default_rng(self.random_state)\n            idx = rng.choice(n, size=self.max_samples, replace=False)\n            Xs = X.iloc[idx] if hasattr(X, \"iloc\") else X[idx]\n            ys = y[idx]\n            self.model.fit(Xs, ys)\n        else:\n            self.model.fit(X, y)\n        return self\n\n    def predict_proba(self, X):\n        return self.model.predict_proba(X)\n\n    def predict(self, X):\n        return self.model.predict(X)\n\n    @property\n    def classes_(self):\n        return self.model.classes_\n\n# ========================\n# Submit function\n# ========================\ndef submit(body_parts_tracked_str, switch_tr, binary_classifier, X_tr, label, meta, test_df):\n    train_cols = list(X_tr.columns)\n\n    # Fit one-vs-rest models\n    model_list = []\n    for action in label.columns:\n        action_mask = ~ label[action].isna().values\n        y_action = label[action][action_mask].values.astype(int)\n        if not (y_action == 0).all():\n            model = clone(binary_classifier)\n            model.fit(X_tr[action_mask], y_action)\n            model_list.append((action, model))\n\n    # Inference\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\n    test_subset = test_df[test_df.body_parts_tracked == body_parts_tracked_str]\n    generator = generate_mouse_data(test_subset, 'test',\n                                    generate_single=(switch_tr == 'single'),\n                                    generate_pair=(switch_tr == 'pair'))\n    if verbose: print(f\"n_videos: {len(test_subset)}\")\n\n    parts = []\n\n    for switch_te, data_te, meta_te, actions_te in generator:\n        assert switch_te == switch_tr\n        try:\n            # Transform\n            if switch_te == 'single':\n                X_te = transform_single(data_te, body_parts_tracked)\n            else:\n                X_te = transform_pair(data_te, body_parts_tracked)\n\n            X_te = X_te.reindex(columns=train_cols)\n\n            if verbose and len(X_te) == 0:\n                print(\"ERROR: X_te is empty\")\n            del data_te\n\n            pred = pd.DataFrame(index=meta_te.video_frame)\n            for action, model in model_list:\n                if action in actions_te:\n                    pred[action] = model.predict_proba(X_te)[:, 1]\n            del X_te\n\n            if pred.shape[1] != 0:\n                if USE_HYSTERESIS_DECODER:\n                    submission_part = predict_multiclass_hysteresis(pred, meta_te)\n                else:\n                    # Fallback to simple decoder\n                    submission_part = predict_multiclass_simple(pred, meta_te)\n                parts.append(submission_part)\n            else:\n                if verbose: print(\"  ERROR: no useful training data\")\n        except KeyError:\n            if verbose: print(f'  ERROR: KeyError because of missing bodypart ({switch_tr})')\n            del data_te\n\n    if len(parts) == 0:\n        return pd.DataFrame(columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n    return pd.concat(parts, ignore_index=True)\n\n# ========================\n# Simple decoder fallback\n# ========================\ndef predict_multiclass_simple(pred: pd.DataFrame, meta: pd.DataFrame, threshold: float = 0.231) -> pd.DataFrame:\n    \"\"\"Simple argmax decoder with threshold\"\"\"\n    ama = np.argmax(pred.values, axis=1)\n    ama = np.where(pred.max(axis=1).values >= threshold, ama, -1)\n    ama = pd.Series(ama, index=meta.video_frame.values)\n\n    changes_mask = (ama != ama.shift(1)).values\n    ama_changes = ama[changes_mask]\n    meta_changes = meta[changes_mask]\n\n    mask = ama_changes.values >= 0\n    mask[-1] = False\n    submission_part = 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.columns[ama_changes[mask].values],\n        'start_frame': ama_changes.index[mask],\n        'stop_frame':   ama_changes.index[1:][mask[:-1]]\n    })\n\n    # Repair stop frame if group changed\n    stop_video_id = meta_changes['video_id'][1:][mask[:-1]].values\n    stop_agent_id = meta_changes['agent_id'][1:][mask[:-1]].values\n    stop_target_id = meta_changes['target_id'][1:][mask[:-1]].values\n    for i in range(len(submission_part)):\n        v = submission_part.video_id.iloc[i]\n        a = submission_part.agent_id.iloc[i]\n        t = submission_part.target_id.iloc[i]\n        if stop_video_id[i] != v or stop_agent_id[i] != a or stop_target_id[i] != t:\n            new_stop_frame = meta.query(\"(video_id == @v)\").video_frame.max() + 1\n            submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n\n    assert (submission_part.stop_frame > submission_part.start_frame).all(), 'stop <= start'\n    if verbose: print('  actions found:', len(submission_part))\n    return submission_part\n\n# ========================\n# Robustify function\n# ========================\ndef robustify(submission: pd.DataFrame,\n              dataset: pd.DataFrame,\n              traintest: str,\n              traintest_directory: Optional[str] = None) -> pd.DataFrame:\n    \"\"\"Clean and fill empty videos in submission\"\"\"\n    if traintest_directory is None:\n        traintest_directory = f\"{CFG.data_root}/{traintest}_tracking\"\n\n    for c in ['start_frame', 'stop_frame']:\n        submission[c] = pd.to_numeric(submission[c], errors='coerce')\n    submission = submission.dropna(subset=['start_frame', 'stop_frame'])\n    submission[['start_frame', 'stop_frame']] = submission[['start_frame', 'stop_frame']].astype(int)\n\n    old_len = len(submission)\n    submission = submission[submission.start_frame < submission.stop_frame]\n    if len(submission) != old_len:\n        print(\"ERROR: Dropped frames with start >= stop\")\n\n    old_len = len(submission)\n    group_list = []\n    for _, group in submission.groupby(['video_id', 'agent_id', 'target_id'], sort=False):\n        group = group.sort_values('start_frame').reset_index(drop=True)\n        mask = np.ones(len(group), dtype=bool)\n        last_stop_frame = -1\n        for i, row in group.iterrows():\n            st = int(row['start_frame'])\n            en = int(row['stop_frame'])\n            if st < last_stop_frame:\n                mask[i] = False\n            else:\n                last_stop_frame = en\n        group_list.append(group[mask])\n    if len(group_list) > 0:\n        submission = pd.concat(group_list, ignore_index=True)\n    if len(submission) != old_len:\n        print(\"ERROR: Dropped duplicate/overlapping frames\")\n\n    # Fill empty videos with minimal stub\n    s_list = []\n    for _, row in dataset.iterrows():\n        lab_id = row.get('lab_id', '')\n        if str(lab_id).startswith('MABe22'):\n            continue\n        video_id = row['video_id']\n        if (submission.video_id == video_id).any():\n            continue\n\n        if verbose:\n            print(f\"Video {video_id} has no predictions.\")\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        try:\n            vid = pd.read_parquet(path)\n            start_frame = int(vid.video_frame.min())\n            stop_bound = int(vid.video_frame.max()) + 1\n            if not np.isfinite(start_frame) or not np.isfinite(stop_bound):\n                start_frame, stop_bound = 0, 1\n        except Exception:\n            start_frame, stop_bound = 0, 1\n\n        # Minimal 1-frame stub\n        stub_start = start_frame\n        stub_stop = min(stub_start + 1, stop_bound)\n        if stub_stop > stub_start:\n            s_list.append((video_id, 'mouse1', 'self', 'sniff', stub_start, stub_stop))\n\n    if len(s_list) > 0:\n        fill_df = pd.DataFrame(\n            s_list,\n            columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\n        )\n        fill_df[['start_frame', 'stop_frame']] = fill_df[['start_frame', 'stop_frame']].astype(int)\n        submission = pd.concat([submission, fill_df], ignore_index=True)\n        print(f\"Filled {len(s_list)} empty videos\")\n\n    submission = submission.reset_index(drop=True)\n    return submission\n\n# ========================\n# Runner\n# ========================\nclass SubmitRunner:\n    def __init__(self, cfg: Config):\n        self.cfg = cfg\n        self.submission_parts: List[pd.DataFrame] = []\n\n    def load_metadata(self) -> Tuple[pd.DataFrame, pd.DataFrame, List[str]]:\n        log(\"[INFO] Loading train/test metadata…\")\n        train = pd.read_csv(self.cfg.train_csv)\n        test = pd.read_csv(self.cfg.test_csv)\n\n        train['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\n        body_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\n        return train, test, body_parts_tracked_list\n\n    def build_binary_classifier(self) -> Any:\n        \"\"\"Improved model configuration based on files 2-3\"\"\"\n        clf = make_pipeline(\n            SimpleImputer(),\n            TrainOnSubsetClassifier(\n                model=LGBMClassifier(\n                    n_estimators=200,  # Increased from 150\n                    learning_rate=0.025,  # Balanced between 0.003 and 0.03\n                    min_child_samples=40,\n                    num_leaves=31,\n                    max_depth=-1,\n                    subsample=0.8,\n                    colsample_bytree=0.8,\n                    random_state=42,\n                    verbose=-1\n                ),\n                max_samples=self.cfg.max_train_samples_per_action,\n                random_state=42\n            )\n        )\n        return clf\n\n    def run(self):\n        assert validate_or_submit == 'submit', \"This script is submit-only by design.\"\n\n        train, test, body_parts_tracked_list = self.load_metadata()\n\n        for section in range(1, len(body_parts_tracked_list)):\n            body_parts_tracked_str = body_parts_tracked_list[section]\n            try:\n                body_parts = json.loads(body_parts_tracked_str)\n                log(f\"{section}. Processing videos with {body_parts}\")\n\n                if len(body_parts) > 5:\n                    body_parts = [b for b in body_parts if b not in drop_body_parts]\n\n                train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n\n                # Collect batches\n                single_mouse_list, single_label_list, single_meta_list = [], [], []\n                mouse_pair_list, mouse_label_list, mouse_meta_list = [], [], []\n\n                for switch, data, meta, label in generate_mouse_data(train_subset, 'train'):\n                    if switch == 'single':\n                        single_mouse_list.append(data)\n                        single_meta_list.append(meta)\n                        single_label_list.append(label)\n                    else:\n                        mouse_pair_list.append(data)\n                        mouse_meta_list.append(meta)\n                        mouse_label_list.append(label)\n\n                binary_classifier = self.build_binary_classifier()\n\n                # SINGLE head\n                if len(single_mouse_list) > 0:\n                    single_mouse = pd.concat(single_mouse_list)\n                    single_mouse_label = pd.concat(single_label_list)\n                    single_mouse_meta = pd.concat(single_meta_list)\n                    del single_mouse_list, single_label_list, single_meta_list\n                    assert len(single_mouse) == len(single_mouse_label) == len(single_mouse_meta)\n\n                    X_tr = transform_single(single_mouse, body_parts)\n                    del single_mouse\n                    log(f\"{X_tr.shape=}\")\n\n                    sub_part = submit(\n                        body_parts_tracked_str, 'single',\n                        binary_classifier, X_tr, single_mouse_label, single_mouse_meta,\n                        test_df=test\n                    )\n                    if len(sub_part):\n                        self.submission_parts.append(sub_part)\n                    del X_tr, single_mouse_label, single_mouse_meta\n                    gc.collect()\n\n                # PAIR head\n                if len(mouse_pair_list) > 0:\n                    mouse_pair = pd.concat(mouse_pair_list)\n                    mouse_pair_label = pd.concat(mouse_label_list)\n                    mouse_pair_meta = pd.concat(mouse_meta_list)\n                    del mouse_pair_list, mouse_label_list, mouse_meta_list\n                    assert len(mouse_pair) == len(mouse_pair_label) == len(mouse_pair_meta)\n\n                    X_tr = transform_pair(mouse_pair, body_parts)\n                    del mouse_pair\n                    log(f\"{X_tr.shape=}\")\n\n                    sub_part = submit(\n                        body_parts_tracked_str, 'pair',\n                        binary_classifier, X_tr, mouse_pair_label, mouse_pair_meta,\n                        test_df=test\n                    )\n                    if len(sub_part):\n                        self.submission_parts.append(sub_part)\n                    del X_tr, mouse_pair_label, mouse_pair_meta\n                    gc.collect()\n\n            except Exception as e:\n                print(f'***Exception*** {e}')\n            print()\n\n        # Stitch submission\n        if len(self.submission_parts) > 0:\n            submission = pd.concat(self.submission_parts, ignore_index=True)\n        else:\n            submission = pd.DataFrame(dict(\n                video_id=438887472,\n                agent_id='mouse1',\n                target_id='self',\n                action='rear',\n                start_frame=278,\n                stop_frame=500\n            ), index=[44])\n\n        submission_robust = robustify(submission, test, 'test')\n        submission_robust.index.name = self.cfg.row_id_col\n        submission_robust.to_csv(self.cfg.submission_file)\n        log(f\"[DONE] Wrote {self.cfg.submission_file} with {len(submission_robust):,} rows.\")\n\n\n# ========================\n# Entry\n# ========================\nif __name__ == \"__main__\":\n    warnings.filterwarnings(\"ignore\")\n    SubmitRunner(CFG).run()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}