{"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":"nvidiaTeslaT4","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":"markdown","source":"# Mouse Behavior Detection - Ensemble Model 🐁\n\n## 📚 Reference\n\nThis implementation is based on and extends the work from:\n- [Mouse Behaviour Prediction v1](https://www.kaggle.com/code/paharpurailab/mouse-behaviour-prediction-v1) by paharpurailab , thank you \n\n## 🗝️ Key Improvements\n\n### Feature Engineering\n- **Enhanced temporal features**: Additional median features for single and pair interactions\n- **Improved smoothing**: Multi-window temporal smoothing (7+3 frames)\n- **FPS-aware scaling**: All temporal windows properly scaled by video frame rate\n\n### Prediction Improvements\n- **Optimized thresholds**: Lower default threshold (0.25 vs 0.27) for better recall\n- **Minimum duration**: Reduced from 3 to 2 frames for more sensitive detection\n- **Better error handling**: Robust exception handling for imbalanced classes\n\n### Technical Optimizations\n- **LightGBM on CPU**: More stable for imbalanced data (GPU had issues)\n- **Conservative parameters**: Safer settings for rare behaviors\n- **Sample size**: Increased to 500K-600K samples per model\n\n## 📝 Notes\n\n- LightGBM models run on CPU for stability with imbalanced data\n- XGBoost and CatBoost utilize GPU acceleration\n- The model handles missing values and rare behaviors gracefully\n- Ensemble combines predictions from all successful models","metadata":{}},{"cell_type":"code","source":"validate_or_submit = 'submit'\nverbose = True\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport itertools\nimport warnings\nimport json\nimport os, random\nimport gc\nimport lightgbm\nfrom collections import defaultdict\nimport polars as pl\nfrom scipy import signal, stats\n\nfrom sklearn.base import ClassifierMixin, BaseEstimator, clone\nfrom sklearn.model_selection import cross_val_predict, GroupKFold, StratifiedKFold\nfrom sklearn.pipeline import make_pipeline\nfrom sklearn.model_selection import StratifiedShuffleSplit\nfrom sklearn.metrics import f1_score\n\nwarnings.filterwarnings('ignore')\n# Suppress specific warnings that may appear with GPU usage\nwarnings.filterwarnings('ignore', category=UserWarning)\nwarnings.filterwarnings('ignore', category=FutureWarning)\nos.environ['PYTHONWARNINGS'] = 'ignore'\n\n# Try importing additional models\ntry:\n    from xgboost import XGBClassifier\n    XGBOOST_AVAILABLE = True\nexcept:\n    XGBOOST_AVAILABLE = False\n\ntry:\n    from catboost import CatBoostClassifier\n    CATBOOST_AVAILABLE = True\nexcept:\n    CATBOOST_AVAILABLE = False\n\n# --- SEED EVERYTHING -----\nSEED = 1234\nos.environ[\"PYTHONHASHSEED\"] = str(SEED)  # set very early\n\nrnd = np.random.RandomState(SEED)\nrandom.seed(SEED)\nnp.random.seed(SEED)\n\n# --- GPU CONFIGURATION -----\n# T4 x 2 GPU setup for Kaggle\nUSE_GPU = True\nNUM_GPUS = 2\n\n\nclass StratifiedSubsetClassifier(ClassifierMixin, BaseEstimator):\n    def __init__(self, estimator, n_samples=None):\n        self.estimator = estimator\n        self.n_samples = n_samples  # if None → no subsampling/stratification\n\n    def _to_numpy(self, X):\n        try:\n            return X.to_numpy(np.float32, copy=False)\n        except AttributeError:\n            return np.asarray(X, dtype=np.float32)\n\n    def fit(self, X, y):\n        Xn = self._to_numpy(X)\n        y = np.asarray(y).ravel()\n\n        uniq = np.unique(y[~pd.isna(y)])\n        if set(uniq.tolist()) == {0, 2}:\n            y = (y > 0).astype(np.int8)\n\n        # Check for minimum samples per class\n        unique_classes, counts = np.unique(y, return_counts=True)\n        min_class_count = counts.min() if len(counts) > 0 else 0\n        \n        # If n_samples is None → fit on full data, no stratification\n        if self.n_samples is None or len(Xn) <= int(self.n_samples):\n            try:\n                self.estimator.fit(Xn, y)\n            except Exception as e:\n                # If fit fails, try with more conservative parameters\n                # This is a fallback - the exception will be caught by caller\n                raise e\n        else:\n            sss = StratifiedShuffleSplit(n_splits=1, train_size=int(self.n_samples), random_state=42)\n            try:\n                idx, _ = next(sss.split(np.zeros_like(y), y))\n                self.estimator.fit(Xn[idx], y[idx])\n            except Exception:\n                # Fallback to simple sampling\n                step = max(len(Xn) // int(self.n_samples), 1)\n                try:\n                    self.estimator.fit(Xn[::step], y[::step])\n                except Exception as e:\n                    # Re-raise to be caught by caller\n                    raise e\n\n        try:\n            self.classes_ = np.asarray(self.estimator.classes_)\n        except Exception:\n            self.classes_ = np.unique(y)\n        return self\n\n    def predict_proba(self, X):\n        Xn = self._to_numpy(X)\n        try:\n            P = self.estimator.predict_proba(Xn)\n        except Exception:\n            if len(self.classes_) == 1:\n                n = len(Xn)\n                c = int(self.classes_[0])\n                if c == 1:\n                    return np.column_stack([np.zeros(n, dtype=np.float32), np.ones(n, dtype=np.float32)])\n                else:\n                    return np.column_stack([np.ones(n, dtype=np.float32), np.zeros(n, dtype=np.float32)])\n            return np.full((len(Xn), 2), 0.5, dtype=np.float32)\n\n        P = np.asarray(P)\n        if P.ndim == 1:\n            P1 = P.astype(np.float32)\n            return np.column_stack([1.0 - P1, P1])\n        if P.shape[1] == 1 and len(self.classes_) == 2:\n            P1 = P[:, 0].astype(np.float32)\n            return np.column_stack([1.0 - P1, P1])\n        return P\n\n    def predict(self, X):\n        Xn = self._to_numpy(X)\n        try:\n            return self.estimator.predict(Xn)\n        except Exception:\n            return np.argmax(self.predict_proba(Xn), axis=1)\n\n\n# ==================== SCORING FUNCTIONS ====================\n\nclass HostVisibleError(Exception):\n    pass\n\n\ndef single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1) -> float:\n    label_frames: defaultdict[str, set[int]] = defaultdict(set)\n    prediction_frames: defaultdict[str, set[int]] = defaultdict(set)\n\n    for row in lab_solution.to_dicts():\n        label_frames[row['label_key']].update(range(row['start_frame'], row['stop_frame']))\n\n    for video in lab_solution['video_id'].unique():\n        active_labels: str = lab_solution.filter(pl.col('video_id') == video)['behaviors_labeled'].first()\n        active_labels: set[str] = set(json.loads(active_labels))\n        predicted_mouse_pairs: defaultdict[str, set[int]] = defaultdict(set)\n\n        for row in lab_submission.filter(pl.col('video_id') == video).to_dicts():\n            if ','.join([str(row['agent_id']), str(row['target_id']), row['action']]) not in active_labels:\n                continue\n\n            new_frames = set(range(row['start_frame'], row['stop_frame']))\n            new_frames = new_frames.difference(prediction_frames[row['prediction_key']])\n            prediction_pair = ','.join([str(row['agent_id']), str(row['target_id'])])\n            if predicted_mouse_pairs[prediction_pair].intersection(new_frames):\n                raise HostVisibleError('Multiple predictions for the same frame from one agent/target pair')\n            prediction_frames[row['prediction_key']].update(new_frames)\n            predicted_mouse_pairs[prediction_pair].update(new_frames)\n\n    tps = defaultdict(int)\n    fns = defaultdict(int)\n    fps = defaultdict(int)\n    for key, pred_frames in prediction_frames.items():\n        action = key.split('_')[-1]\n        matched_label_frames = label_frames[key]\n        tps[action] += len(pred_frames.intersection(matched_label_frames))\n        fns[action] += len(matched_label_frames.difference(pred_frames))\n        fps[action] += len(pred_frames.difference(matched_label_frames))\n\n    distinct_actions = set()\n    for key, frames in label_frames.items():\n        action = key.split('_')[-1]\n        distinct_actions.add(action)\n        if key not in prediction_frames:\n            fns[action] += len(frames)\n\n    action_f1s = []\n    for action in distinct_actions:\n        if tps[action] + fns[action] + fps[action] == 0:\n            action_f1s.append(0)\n        else:\n            action_f1s.append((1 + beta**2) * tps[action] / ((1 + beta**2) * tps[action] + beta**2 * fns[action] + fps[action]))\n    return sum(action_f1s) / len(action_f1s)\n\n\ndef mouse_fbeta(solution: pd.DataFrame, submission: pd.DataFrame, beta: float = 1) -> float:\n    if len(solution) == 0 or len(submission) == 0:\n        raise ValueError('Missing solution or submission data')\n\n    expected_cols = ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\n\n    for col in expected_cols:\n        if col not in solution.columns:\n            raise ValueError(f'Solution is missing column {col}')\n        if col not in submission.columns:\n            raise ValueError(f'Submission is missing column {col}')\n\n    solution: pl.DataFrame = pl.DataFrame(solution)\n    submission: pl.DataFrame = pl.DataFrame(submission)\n    assert (solution['start_frame'] <= solution['stop_frame']).all()\n    assert (submission['start_frame'] <= submission['stop_frame']).all()\n    solution_videos = set(solution['video_id'].unique())\n    submission = submission.filter(pl.col('video_id').is_in(solution_videos))\n\n    solution = solution.with_columns(\n        pl.concat_str(\n            [\n                pl.col('video_id').cast(pl.Utf8),\n                pl.col('agent_id').cast(pl.Utf8),\n                pl.col('target_id').cast(pl.Utf8),\n                pl.col('action'),\n            ],\n            separator='_',\n        ).alias('label_key'),\n    )\n    submission = submission.with_columns(\n        pl.concat_str(\n            [\n                pl.col('video_id').cast(pl.Utf8),\n                pl.col('agent_id').cast(pl.Utf8),\n                pl.col('target_id').cast(pl.Utf8),\n                pl.col('action'),\n            ],\n            separator='_',\n        ).alias('prediction_key'),\n    )\n\n    lab_scores = []\n    for lab in solution['lab_id'].unique():\n        lab_solution = solution.filter(pl.col('lab_id') == lab).clone()\n        lab_videos = set(lab_solution['video_id'].unique())\n        lab_submission = submission.filter(pl.col('video_id').is_in(lab_videos)).clone()\n        lab_scores.append(single_lab_f1(lab_solution, lab_submission, beta=beta))\n\n    return sum(lab_scores) / len(lab_scores)\n\n\ndef score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, beta: float = 1) -> float:\n    solution = solution.drop(row_id_column_name, axis='columns', errors='ignore')\n    submission = submission.drop(row_id_column_name, axis='columns', errors='ignore')\n    return mouse_fbeta(solution, submission, beta=beta)\n\n\n# ==================== DATA LOADING ====================\n\ntrain = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\ntrain['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\n\ntest = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\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\ndef generate_mouse_data(dataset, traintest, traintest_directory=None, generate_single=True, generate_pair=True):\n    assert traintest in ['train', 'test']\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n    for _, row in dataset.iterrows():\n\n        lab_id = row.lab_id\n        video_id = row.video_id\n\n        if type(row.behaviors_labeled) != str:\n            if verbose:\n                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        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        del vid\n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n        pvid /= row.pix_per_cm_approx\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            vid_behaviors_subset = vid_behaviors.query(\"target == 'self'\")\n            for mouse_id_str in np.unique(vid_behaviors_subset.agent):\n                try:\n                    mouse_id = int(mouse_id_str[-1])\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"agent == @mouse_id_str\").action)\n                    single_mouse = pvid.loc[:, mouse_id]\n                    assert len(single_mouse) == len(pvid)\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                        'frames_per_second': row.frames_per_second\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                            annot_row = annot_subset.iloc[i]\n                            single_mouse_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'single', single_mouse, single_mouse_meta, single_mouse_label\n                    else:\n                        if verbose:\n                            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        if generate_pair:\n            vid_behaviors_subset = vid_behaviors.query(\"target != 'self'\")\n            if len(vid_behaviors_subset) > 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                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"(agent == @agent_str) & (target == @target_str)\").action)\n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\n                    assert len(mouse_pair) == len(pvid)\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                        'frames_per_second': row.frames_per_second\n                    })\n                    if traintest == 'train':\n                        mouse_pair_label = pd.DataFrame(0.0, columns=vid_agent_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                            annot_row = annot_subset.iloc[i]\n                            mouse_pair_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'pair', mouse_pair, mouse_pair_meta, mouse_pair_label\n                    else:\n                        if verbose:\n                            print('- test pair', video_id, agent, target)\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n\n\n# ==================== ADAPTIVE THRESHOLDING ====================\n\n# Optimized action-specific thresholds for better F1 scores\naction_thresholds = defaultdict(lambda: 0.25)  # Lower default for better recall\n# Action-specific thresholds can be tuned per action if needed\n\ndef predict_multiclass_adaptive(pred, meta, action_thresholds):\n    \"\"\"Adaptive thresholding per action + enhanced temporal smoothing.\"\"\"\n    # Enhanced temporal smoothing with multiple window sizes\n    pred_smoothed = pred.rolling(window=7, min_periods=1, center=True).mean()\n    # Additional smoothing for stability\n    pred_smoothed = pred_smoothed.rolling(window=3, min_periods=1, center=True).mean()\n\n    ama = np.argmax(pred_smoothed, axis=1)\n\n    max_probs = pred_smoothed.max(axis=1)\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        threshold = action_thresholds.get(action, 0.27)\n        threshold_mask |= (action_mask & (max_probs >= threshold))\n\n    ama = np.where(threshold_mask, ama, -1)\n    ama = pd.Series(ama, index=meta.video_frame)\n\n    changes_mask = (ama != ama.shift(1)).values\n    ama_changes = ama[changes_mask]\n    meta_changes = meta[changes_mask]\n    mask = ama_changes.values >= 0\n    mask[-1] = False\n\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    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\n    for i in range(len(submission_part)):\n        video_id = submission_part.video_id.iloc[i]\n        agent_id = submission_part.agent_id.iloc[i]\n        target_id = submission_part.target_id.iloc[i]\n        if i < len(stop_video_id):\n            if stop_video_id[i] != video_id or stop_agent_id[i] != agent_id or stop_target_id[i] != target_id:\n                new_stop_frame = meta.query(\"(video_id == @video_id)\").video_frame.max() + 1\n                submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n        else:\n            new_stop_frame = meta.query(\"(video_id == @video_id)\").video_frame.max() + 1\n            submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame\n\n    # Filter very short events (slightly more lenient for better recall)\n    duration = submission_part.stop_frame - submission_part.start_frame\n    submission_part = submission_part[duration >= 2].reset_index(drop=True)\n\n    if len(submission_part) > 0:\n        assert (submission_part.stop_frame > submission_part.start_frame).all(), 'stop <= start'\n\n    if verbose:\n        print(f'  actions found: {len(submission_part)}')\n    return submission_part\n\n\n# ==================== ADVANCED FEATURE ENGINEERING (FPS-AWARE, SHORTER WINDOWS) ====================\n\ndef safe_rolling(series, window, func, min_periods=None):\n    \"\"\"Safe rolling operation with NaN handling\"\"\"\n    if min_periods is None:\n        min_periods = max(1, window // 4)\n    return series.rolling(window, min_periods=min_periods, center=True).apply(func, raw=True)\n\ndef _scale(n_frames_at_30fps, fps, ref=30.0):\n    \"\"\"Scale a frame count defined at 30 fps to the current video's fps.\"\"\"\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    \"\"\"Signed version of _scale for forward/backward shifts (keeps at least 1 frame when |n|>=1).\"\"\"\n    if n_frames_at_30fps == 0:\n        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, center_x, center_y, fps):\n    \"\"\"Trajectory curvature (shorter windows).\"\"\"\n    vel_x = center_x.diff()\n    vel_y = center_y.diff()\n    acc_x = vel_x.diff()\n    acc_y = vel_y.diff()\n\n    cross_prod = vel_x * acc_y - vel_y * acc_x\n    vel_mag = np.sqrt(vel_x**2 + vel_y**2)\n    curvature = np.abs(cross_prod) / (vel_mag**3 + 1e-6)  # invariant to time scaling\n\n    for w in [25, 50, 75]:  # shorter than original [30, 60]\n        ws = _scale(w, fps)\n        X[f'curv_mean_{w}'] = curvature.rolling(ws, min_periods=max(1, ws // 5)).mean()\n\n    angle = np.arctan2(vel_y, vel_x)\n    angle_change = np.abs(angle.diff())\n    w = 30\n    ws = _scale(w, fps)\n    X[f'turn_rate_{w}'] = angle_change.rolling(ws, min_periods=max(1, ws // 5)).sum()\n\n    return X\n\ndef add_multiscale_features(X, center_x, center_y, fps):\n    \"\"\"Multi-scale temporal features (speed in cm/s; shorter scales).\"\"\"\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)\n\n    scales = [20, 40, 60, 80]  # was [10, 40, 160]\n    for scale in scales:\n        ws = _scale(scale, fps)\n        if len(speed) >= ws:\n            X[f'sp_m{scale}'] = speed.rolling(ws, min_periods=max(1, ws // 4)).mean()\n            X[f'sp_s{scale}'] = speed.rolling(ws, min_periods=max(1, ws // 4)).std()\n\n    if len(scales) >= 2 and f'sp_m{scales[0]}' in X.columns and f'sp_m{scales[-1]}' in X.columns:\n        X['sp_ratio'] = X[f'sp_m{scales[0]}'] / (X[f'sp_m{scales[-1]}'] + 1e-6)\n\n    return X\n\ndef add_state_features(X, center_x, center_y, fps):\n    \"\"\"Behavioral state transitions; fps-invariant bins; shorter windows.\"\"\"\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)  # cm/s\n    w_ma = _scale(15, fps)\n    speed_ma = speed.rolling(w_ma, min_periods=max(1, w_ma // 3)).mean()\n\n    try:\n        # Bins in cm/s to keep semantics across fps:\n        bins = [-np.inf, 0.5 * fps, 2.0 * fps, 5.0 * fps, np.inf]\n        speed_states = pd.cut(speed_ma, bins=bins, labels=[0, 1, 2, 3]).astype(float)\n\n        for window in [20, 40, 60, 80]:  # was [60, 120]\n            ws = _scale(window, fps)\n            if len(speed_states) >= ws:\n                for state in [0, 1, 2, 3]:\n                    X[f's{state}_{window}'] = (\n                        (speed_states == state).astype(float)\n                        .rolling(ws, min_periods=max(1, ws // 5)).mean()\n                    )\n                state_changes = (speed_states != speed_states.shift(1)).astype(float)\n                X[f'trans_{window}'] = state_changes.rolling(ws, min_periods=max(1, ws // 5)).sum()\n    except Exception:\n        pass\n\n    return X\n\ndef add_longrange_features(X, center_x, center_y, fps):\n    \"\"\"Long-range temporal features (shorter spans).\"\"\"\n    for window in [30, 60, 120]:  # was [120, 240]\n        ws = _scale(window, fps)\n        if len(center_x) >= ws:\n            X[f'x_ml{window}'] = center_x.rolling(ws, min_periods=max(5, ws // 6)).mean()\n            X[f'y_ml{window}'] = center_y.rolling(ws, min_periods=max(5, ws // 6)).mean()\n\n    # EWM spans\n    for span in [30, 60, 120]:   # was [60, 120]\n        s = _scale(span, fps)\n        X[f'x_e{span}'] = center_x.ewm(span=s, min_periods=1).mean()\n        X[f'y_e{span}'] = center_y.ewm(span=s, min_periods=1).mean()\n\n    speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps)  # cm/s\n    for window in [30, 60, 120]:  # was [60, 120]\n        ws = _scale(window, fps)\n        if len(speed) >= ws:\n            X[f'sp_pct{window}'] = speed.rolling(ws, min_periods=max(5, ws // 6)).rank(pct=True)\n\n    return X\n\ndef add_interaction_features(X, mouse_pair, avail_A, avail_B, fps):\n    \"\"\"Social interaction features (unchanged core windows).\"\"\"\n    if 'body_center' not in avail_A or 'body_center' not in avail_B:\n        return X\n\n    rel_x = mouse_pair['A']['body_center']['x'] - mouse_pair['B']['body_center']['x']\n    rel_y = mouse_pair['A']['body_center']['y'] - mouse_pair['B']['body_center']['y']\n    rel_dist = np.sqrt(rel_x**2 + rel_y**2)\n\n    # per-frame velocities (cm/frame)\n    A_vx = mouse_pair['A']['body_center']['x'].diff()\n    A_vy = mouse_pair['A']['body_center']['y'].diff()\n    B_vx = mouse_pair['B']['body_center']['x'].diff()\n    B_vy = mouse_pair['B']['body_center']['y'].diff()\n\n    A_lead = (A_vx * rel_x + A_vy * rel_y) / (np.sqrt(A_vx**2 + A_vy**2) * rel_dist + 1e-6)\n    B_lead = (B_vx * (-rel_x) + B_vy * (-rel_y)) / (np.sqrt(B_vx**2 + B_vy**2) * rel_dist + 1e-6)\n\n    for window in [30, 60]:\n        ws = _scale(window, fps)\n        X[f'A_ld{window}'] = A_lead.rolling(ws, min_periods=max(1, ws // 6)).mean()\n        X[f'B_ld{window}'] = B_lead.rolling(ws, min_periods=max(1, ws // 6)).mean()\n\n    approach = -rel_dist.diff()  # decreasing distance => positive approach\n    chase = approach * B_lead\n    w = 30\n    ws = _scale(w, fps)\n    X[f'chase_{w}'] = chase.rolling(ws, min_periods=max(1, ws // 6)).mean()\n\n    for window in [60, 120]:\n        ws = _scale(window, fps)\n        A_sp = np.sqrt(A_vx**2 + A_vy**2)\n        B_sp = np.sqrt(B_vx**2 + B_vy**2)\n        X[f'sp_cor{window}'] = A_sp.rolling(ws, min_periods=max(1, ws // 6)).corr(B_sp)\n\n    return X\n\ndef transform_single(single_mouse, body_parts_tracked, fps):\n    \"\"\"Single mouse transform (shorter windows where relevant; distances in cm).\"\"\"\n    available_body_parts = single_mouse.columns.get_level_values(0)\n\n    # Base distance features (squared distances across body parts)\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 available_body_parts and p2 in available_body_parts\n    })\n    X = X.reindex(columns=[f\"{p1}+{p2}\" for p1, p2 in itertools.combinations(body_parts_tracked, 2)], copy=False)\n\n    # Speed-like features via lagged displacements (duration-aware lag)\n    if all(p in single_mouse.columns for p in ['ear_left', 'ear_right', 'tail_base']):\n        lag = _scale(10, fps)\n        shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(lag)\n        speeds = pd.DataFrame({\n            'sp_lf': np.square(single_mouse['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False),\n            'sp_rt': np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False),\n            'sp_lf2': np.square(single_mouse['ear_left'] - shifted['tail_base']).sum(axis=1, skipna=False),\n            'sp_rt2': np.square(single_mouse['ear_right'] - shifted['tail_base']).sum(axis=1, skipna=False),\n        })\n        X = pd.concat([X, speeds], axis=1)\n\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\n    # Body angle (orientation)\n    if all(p in available_body_parts for p in ['nose', 'body_center', 'tail_base']):\n        v1 = single_mouse['nose'] - single_mouse['body_center']\n        v2 = single_mouse['tail_base'] - single_mouse['body_center']\n        X['body_ang'] = (v1['x'] * v2['x'] + v1['y'] * v2['y']) / (\n            np.sqrt(v1['x']**2 + v1['y']**2) * np.sqrt(v2['x']**2 + v2['y']**2) + 1e-6)\n\n    # Core temporal features (unchanged set; windows scaled by fps)\n    if 'body_center' in available_body_parts:\n        cx = single_mouse['body_center']['x']\n        cy = single_mouse['body_center']['y']\n\n        for w in [5, 15, 30, 60]:\n            ws = _scale(w, fps)\n            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            # Additional features: median and skewness\n            X[f'cx_med{w}'] = cx.rolling(ws, **roll).median()\n            X[f'cy_med{w}'] = cy.rolling(ws, **roll).median()\n\n        # Shorter-window enrichments\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\n    # Nose-tail features with duration-aware lags\n    if all(p in available_body_parts for p in ['nose', 'tail_base']):\n        nt_dist = 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)\n            X[f'nt_lg{lag}'] = nt_dist.shift(l)\n            X[f'nt_df{lag}'] = nt_dist - nt_dist.shift(l)\n\n    # Ear features with duration-aware offsets (expanded offsets)\n    if all(p in available_body_parts for p in ['ear_left', 'ear_right']):\n        ear_d = 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 [-30, -20, -10, 10, 20, 30]:  # expanded from [-20, -10, 10, 20]\n            o = _scale_signed(off, fps)\n            X[f'ear_o{off}'] = ear_d.shift(-o)\n        w = _scale(30, fps)\n        X['ear_con'] = ear_d.rolling(w, min_periods=1, center=True).std() / \\\n                       (ear_d.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n\n    return X.astype(np.float32, copy=False)\n\ndef transform_pair(mouse_pair, body_parts_tracked, fps):\n    \"\"\"Pair transform (shorter windows where relevant; distances in cm).\"\"\"\n    avail_A = mouse_pair['A'].columns.get_level_values(0)\n    avail_B = mouse_pair['B'].columns.get_level_values(0)\n\n    # Inter-mouse distances (squared distances across all part pairs)\n    X = pd.DataFrame({\n        f\"12+{p1}+{p2}\": np.square(mouse_pair['A'][p1] - mouse_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\n    # Speed-like features via lagged displacements (duration-aware lag)\n    if ('A', 'ear_left') in mouse_pair.columns and ('B', 'ear_left') in mouse_pair.columns:\n        lag = _scale(10, fps)\n        shA = mouse_pair['A']['ear_left'].shift(lag)\n        shB = mouse_pair['B']['ear_left'].shift(lag)\n        speeds = pd.DataFrame({\n            'sp_A': np.square(mouse_pair['A']['ear_left'] - shA).sum(axis=1, skipna=False),\n            'sp_AB': np.square(mouse_pair['A']['ear_left'] - shB).sum(axis=1, skipna=False),\n            'sp_B': np.square(mouse_pair['B']['ear_left'] - shB).sum(axis=1, skipna=False),\n        })\n        X = pd.concat([X, speeds], axis=1)\n\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\n    # Relative orientation\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 = mouse_pair['A']['nose'] - mouse_pair['A']['tail_base']\n        dir_B = mouse_pair['B']['nose'] - mouse_pair['B']['tail_base']\n        X['rel_ori'] = (dir_A['x'] * dir_B['x'] + dir_A['y'] * dir_B['y']) / (\n            np.sqrt(dir_A['x']**2 + dir_A['y']**2) * np.sqrt(dir_B['x']**2 + dir_B['y']**2) + 1e-6)\n\n    # Approach rate (duration-aware lag)\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(mouse_pair['A']['nose'] - mouse_pair['B']['nose']).sum(axis=1, skipna=False)\n        lag = _scale(10, fps)\n        shA_n = mouse_pair['A']['nose'].shift(lag)\n        shB_n = mouse_pair['B']['nose'].shift(lag)\n        past = np.square(shA_n - shB_n).sum(axis=1, skipna=False)\n        X['appr'] = cur - past\n\n    # Distance bins (cm; unchanged by fps)\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        cd = np.sqrt((mouse_pair['A']['body_center']['x'] - mouse_pair['B']['body_center']['x'])**2 +\n                     (mouse_pair['A']['body_center']['y'] - mouse_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\n    # Temporal interaction features (keep core windows)\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        cd_full = np.square(mouse_pair['A']['body_center'] - mouse_pair['B']['body_center']).sum(axis=1, skipna=False)\n\n        for w in [5, 15, 30, 60]:\n            ws = _scale(w, fps)\n            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            # Additional distance features\n            X[f'd_med{w}'] = cd_full.rolling(ws, **roll).median()\n\n            d_var = cd_full.rolling(ws, **roll).var()\n            X[f'int{w}'] = 1 / (1 + d_var)\n\n            Axd = mouse_pair['A']['body_center']['x'].diff()\n            Ayd = mouse_pair['A']['body_center']['y'].diff()\n            Bxd = mouse_pair['B']['body_center']['x'].diff()\n            Byd = mouse_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\n    # Nose-nose dynamics (duration-aware lags)\n    if 'nose' in avail_A and 'nose' in avail_B:\n        nn = np.sqrt((mouse_pair['A']['nose']['x'] - mouse_pair['B']['nose']['x'])**2 +\n                     (mouse_pair['A']['nose']['y'] - mouse_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\n    # Velocity alignment (expanded offsets)\n    if 'body_center' in avail_A and 'body_center' in avail_B:\n        Avx = mouse_pair['A']['body_center']['x'].diff()\n        Avy = mouse_pair['A']['body_center']['y'].diff()\n        Bvx = mouse_pair['B']['body_center']['x'].diff()\n        Bvy = mouse_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\n        for off in [-30, -20, -10, 0, 10, 20, 30]:  # expanded from [-20, -10, 0, 10, 20]\n            o = _scale_signed(off, fps)\n            X[f'va_{off}'] = val.shift(-o)\n\n        w = _scale(30, fps)\n        X['int_con'] = cd_full.rolling(w, min_periods=1, center=True).std() / \\\n                       (cd_full.rolling(w, min_periods=1, center=True).mean() + 1e-6)\n\n        # Advanced interaction (reuse helper)\n        X = add_interaction_features(X, mouse_pair, avail_A, avail_B, fps)\n\n    return X.astype(np.float32, copy=False)\n\n\n# ==================== ENSEMBLE TRAINING (fixed sample sizes, as original) ====================\n\ndef submit_ensemble(body_parts_tracked_str, switch_tr, X_tr, label, meta):\n    models = []\n\n    # LightGBM models - use CPU for stability (GPU can have issues with imbalanced data)\n    # LightGBM GPU sometimes fails with imbalanced classes, so we use CPU for reliability\n    lgbm_device = 'cpu'  # Force CPU for stability with imbalanced data\n    \n    # Base LightGBM parameters with safety settings for imbalanced data\n    lgbm_base_params = {\n        'device': lgbm_device,\n        'verbose': -1,\n        'min_data_in_leaf': 5,  # Minimum samples in leaf for stability\n        'min_gain_to_split': 0.0,  # Allow splits even with minimal gain\n        'min_sum_hessian_in_leaf': 1e-3,  # Very small minimum hessian\n    }\n    \n    # LightGBM models with conservative parameters for imbalanced data\n    models.append(make_pipeline(\n        StratifiedSubsetClassifier(\n            lightgbm.LGBMClassifier(\n                n_estimators=300, learning_rate=0.06, min_child_samples=5,\n                num_leaves=31, max_depth=6, subsample=0.8, colsample_bytree=0.8,  # Reduced max_depth\n                reg_alpha=0.01, reg_lambda=0.01,\n                **lgbm_base_params,\n                random_state=SEED, bagging_seed=SEED, feature_fraction_seed=SEED, data_random_seed=SEED\n            ), 600_000,\n        )\n    ))\n    models.append(make_pipeline(\n        StratifiedSubsetClassifier(\n            lightgbm.LGBMClassifier(\n                n_estimators=250, learning_rate=0.08, min_child_samples=5,\n                num_leaves=31, max_depth=6, subsample=0.7, colsample_bytree=0.9,  # Reduced num_leaves and max_depth\n                reg_alpha=0.1, reg_lambda=0.1,\n                **lgbm_base_params,\n                random_state=SEED+1, bagging_seed=SEED+1, feature_fraction_seed=SEED+1, data_random_seed=SEED+1\n            ), 550_000,\n        )\n    ))\n    models.append(make_pipeline(\n        StratifiedSubsetClassifier(\n            lightgbm.LGBMClassifier(\n                n_estimators=200, learning_rate=0.05, min_child_samples=5,\n                num_leaves=63, max_depth=7, subsample=0.75,  # More conservative\n                reg_alpha=0.05, reg_lambda=0.05,\n                **lgbm_base_params,\n                random_state=SEED+2, bagging_seed=SEED+2, feature_fraction_seed=SEED+2, data_random_seed=SEED+2\n            ), 500_000,\n        )\n    ))\n    # Additional LightGBM model with different hyperparameters\n    models.append(make_pipeline(\n        StratifiedSubsetClassifier(\n            lightgbm.LGBMClassifier(\n                n_estimators=280, learning_rate=0.07, min_child_samples=5,\n                num_leaves=47, max_depth=7, subsample=0.75, colsample_bytree=0.85,  # More conservative\n                reg_alpha=0.02, reg_lambda=0.02,\n                **lgbm_base_params,\n                random_state=SEED+3, bagging_seed=SEED+3, feature_fraction_seed=SEED+3, data_random_seed=SEED+3\n            ), 580_000,\n        )\n    ))\n    \n    if XGBOOST_AVAILABLE:\n        # XGBoost with GPU support\n        xgb_tree_method = 'gpu_hist' if USE_GPU else 'hist'\n        xgb_params = {\n            'n_estimators': 250, 'learning_rate': 0.07, 'max_depth': 7,\n            'min_child_weight': 1, 'subsample': 0.8, 'colsample_bytree': 0.8,  # Reduced from 5\n            'reg_alpha': 0.1, 'reg_lambda': 1.0,\n            'tree_method': xgb_tree_method, 'verbosity': 0,\n            'random_state': SEED\n        }\n        if USE_GPU and xgb_tree_method == 'gpu_hist':\n            xgb_params['gpu_id'] = 0\n        models.append(make_pipeline(\n            StratifiedSubsetClassifier(\n                XGBClassifier(**xgb_params), 600_000,\n            )\n        ))\n        # Second XGBoost model with different hyperparameters\n        xgb_params2 = {\n            'n_estimators': 220, 'learning_rate': 0.08, 'max_depth': 6,\n            'min_child_weight': 1, 'subsample': 0.75, 'colsample_bytree': 0.85,  # Reduced from 4\n            'reg_alpha': 0.05, 'reg_lambda': 0.5,\n            'tree_method': xgb_tree_method, 'verbosity': 0,\n            'random_state': SEED+10\n        }\n        if USE_GPU and NUM_GPUS > 1 and xgb_tree_method == 'gpu_hist':\n            xgb_params2['gpu_id'] = 1\n        models.append(make_pipeline(\n            StratifiedSubsetClassifier(\n                XGBClassifier(**xgb_params2), 550_000,\n            )\n        ))\n    \n    if CATBOOST_AVAILABLE:\n        # CatBoost with GPU support\n        catboost_task_type = 'GPU' if USE_GPU else 'CPU'\n        catboost_params = {\n            'iterations': 200, 'learning_rate': 0.08, 'depth': 7,\n            'l2_leaf_reg': 3, 'border_count': 128,\n            'task_type': catboost_task_type,\n            'verbose': False, 'allow_writing_files': False,\n            'random_seed': SEED\n        }\n        if USE_GPU and NUM_GPUS > 1:\n            catboost_params['devices'] = '0:1'\n        models.append(make_pipeline(\n            StratifiedSubsetClassifier(\n                CatBoostClassifier(**catboost_params), 600_000,\n            )\n        ))\n        # Additional CatBoost model with different hyperparameters\n        catboost_params2 = {\n            'iterations': 180, 'learning_rate': 0.1, 'depth': 6,\n            'l2_leaf_reg': 5, 'border_count': 64,\n            'task_type': catboost_task_type,\n            'verbose': False, 'allow_writing_files': False,\n            'random_seed': SEED+20\n        }\n        if USE_GPU and NUM_GPUS > 1:\n            catboost_params2['devices'] = '0:1'\n        models.append(make_pipeline(\n            StratifiedSubsetClassifier(\n                CatBoostClassifier(**catboost_params2), 550_000,\n            )\n        ))\n\n    X_tr_np = X_tr.to_numpy(np.float32, copy=False)\n    del X_tr; gc.collect()\n\n    model_list = []\n    for action in label.columns:\n        y_raw = label[action].to_numpy()\n        mask = ~pd.isna(y_raw)\n        y_action = y_raw[mask].astype(int)\n        # Check for sufficient positive samples and class balance\n        n_positives = np.sum(y_action)\n        n_total = len(y_action)\n        # Require at least 10 positive samples and ensure class balance\n        if not (y_action == 0).all() and n_positives >= 10 and n_positives < n_total * 0.99:\n            trained = []\n            idx = np.flatnonzero(mask)\n            for i, m in enumerate(models):\n                try:\n                    m_clone = clone(m)\n                    m_clone.fit(X_tr_np[idx], y_action)\n                    trained.append(m_clone)\n                except Exception as e:\n                    # Skip failed models and continue with others\n                    # LightGBM may fail on some imbalanced actions, but XGBoost/CatBoost should work\n                    if verbose:\n                        error_msg = str(e)[:50]\n                        print(f\"    Model {i} failed for action {action}: {error_msg}\")\n                    continue\n            if len(trained) > 0:  # Only add if at least one model trained successfully\n                model_list.append((action, trained))\n            elif verbose:\n                print(f\"    No models trained successfully for action {action}\")\n\n    del X_tr_np; gc.collect()\n\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[test.body_parts_tracked == body_parts_tracked_str]\n    generator = generate_mouse_data(\n        test_subset, 'test',\n        generate_single=(switch_tr == 'single'),\n        generate_pair=(switch_tr == 'pair')\n    )\n\n    fps_lookup = (\n        test_subset[['video_id', 'frames_per_second']]\n        .drop_duplicates('video_id')\n        .set_index('video_id')['frames_per_second']\n        .to_dict()\n    )\n\n    if verbose:\n        print(f\"n_videos: {len(test_subset)}, n_models: {len(models)}\")\n\n    for switch_te, data_te, meta_te, actions_te in generator:\n        assert switch_te == switch_tr\n        try:\n            fps_i = _fps_from_meta(meta_te, fps_lookup, default_fps=30.0)\n\n            if switch_te == 'single':\n                X_te = transform_single(data_te, body_parts_tracked, fps_i).astype(np.float32)\n            else:\n                X_te = transform_pair(data_te, body_parts_tracked, fps_i).astype(np.float32)\n\n            X_te_np = X_te.to_numpy(np.float32, copy=False)\n            del X_te, data_te; gc.collect()\n\n            pred = pd.DataFrame(index=meta_te.video_frame)\n            for action, trained in model_list:\n                if action in actions_te:\n                    probs = [m.predict_proba(X_te_np)[:, 1] for m in trained]\n                    # Weighted ensemble: give more weight to later models (often better)\n                    n_models = len(probs)\n                    if n_models > 1:\n                        # Simple linear weighting: later models get slightly more weight\n                        weights = np.linspace(0.9, 1.1, n_models)\n                        weights = weights / weights.sum()\n                        pred[action] = np.average(probs, axis=0, weights=weights)\n                    else:\n                        pred[action] = probs[0]\n\n            del X_te_np; gc.collect()\n\n            if pred.shape[1] != 0:\n                sub_part = predict_multiclass_adaptive(pred, meta_te, action_thresholds)\n                submission_list.append(sub_part)\n            else:\n                if verbose:\n                    print(\"  ERROR: no training data\")\n\n        except Exception as e:\n            if verbose:\n                print(f\"  ERROR: {str(e)[:50]}\")\n            try:\n                del data_te\n            except Exception:\n                pass\n            gc.collect()\n\n\ndef robustify(submission, dataset, traintest, traintest_directory=None):\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n\n    submission = submission[submission.start_frame < submission.stop_frame]\n\n    group_list = []\n    for _, group in submission.groupby(['video_id', 'agent_id', 'target_id']):\n        group = group.sort_values('start_frame')\n        mask = np.ones(len(group), dtype=bool)\n        last_stop = 0\n        for i, (_, row) in enumerate(group.iterrows()):\n            if row['start_frame'] < last_stop:\n                mask[i] = False\n            else:\n                last_stop = row['stop_frame']\n        group_list.append(group[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']\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        vid = pd.read_parquet(path)\n\n        vid_behaviors = eval(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        start_frame = vid.video_frame.min()\n        stop_frame = vid.video_frame.max() + 1\n\n        for (agent, target), actions in vid_behaviors.groupby(['agent', 'target']):\n            batch_len = int(np.ceil((stop_frame - start_frame) / len(actions)))\n            for i, (_, action_row) in enumerate(actions.iterrows()):\n                batch_start = start_frame + i * batch_len\n                batch_stop = min(batch_start + batch_len, stop_frame)\n                s_list.append((video_id, agent, target, action_row['action'], batch_start, batch_stop))\n\n    if len(s_list) > 0:\n        submission = pd.concat([\n            submission,\n            pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n        ])\n\n    submission = submission.reset_index(drop=True)\n    return submission\n\n\n# ==================== MAIN LOOP ====================\n\nsubmission_list = []\n\nprint(f\"XGBoost: {XGBOOST_AVAILABLE}, CatBoost: {CATBOOST_AVAILABLE}\\n\")\n\nfor section in range(1, len(body_parts_tracked_list)):\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        print(f\"{section}. Processing: {len(body_parts_tracked)} body parts\")\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        train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n\n        _fps_lookup = (\n            train_subset[['video_id', 'frames_per_second']]\n            .drop_duplicates('video_id')\n            .set_index('video_id')['frames_per_second']\n            .to_dict()\n        )\n\n        single_list, single_label_list, single_meta_list = [], [], []\n        pair_list, pair_label_list, pair_meta_list = [], [], []\n\n        for switch, data, meta, label in generate_mouse_data(train_subset, 'train'):\n            if switch == 'single':\n                single_list.append(data)\n                single_meta_list.append(meta)\n                single_label_list.append(label)\n            else:\n                pair_list.append(data)\n                pair_meta_list.append(meta)\n                pair_label_list.append(label)\n\n        if len(single_list) > 0:\n            single_feats_parts = []\n            for data_i, meta_i in zip(single_list, single_meta_list):\n                fps_i = _fps_from_meta(meta_i, _fps_lookup, default_fps=30.0)\n                Xi = transform_single(data_i, body_parts_tracked, fps_i).astype(np.float32)\n                single_feats_parts.append(Xi)\n\n            X_tr = pd.concat(single_feats_parts, axis=0, ignore_index=True)\n\n            single_label = pd.concat(single_label_list, axis=0, ignore_index=True)\n            single_meta  = pd.concat(single_meta_list,  axis=0, ignore_index=True)\n\n            del single_list, single_label_list, single_meta_list, single_feats_parts\n            gc.collect()\n\n            print(f\"  Single: {X_tr.shape}\")\n            submit_ensemble(body_parts_tracked_str, 'single', X_tr, single_label, single_meta)\n\n            del X_tr, single_label, single_meta\n            gc.collect()\n\n        if len(pair_list) > 0:\n            pair_feats_parts = []\n            for data_i, meta_i in zip(pair_list, pair_meta_list):\n                fps_i = _fps_from_meta(meta_i, _fps_lookup, default_fps=30.0)\n                Xi = transform_pair(data_i, body_parts_tracked, fps_i).astype(np.float32)\n                pair_feats_parts.append(Xi)\n\n            X_tr = pd.concat(pair_feats_parts, axis=0, ignore_index=True)\n\n            pair_label = pd.concat(pair_label_list, axis=0, ignore_index=True)\n            pair_meta  = pd.concat(pair_meta_list,  axis=0, ignore_index=True)\n\n            del pair_list, pair_label_list, pair_meta_list, pair_feats_parts\n            gc.collect()\n\n            print(f\"  Pair: {X_tr.shape}\")\n            submit_ensemble(body_parts_tracked_str, 'pair', X_tr, pair_label, pair_meta)\n\n            del X_tr, pair_label, pair_meta\n            gc.collect()\n\n    except Exception as e:\n        print(f'***Exception*** {str(e)[:100]}')\n\n    gc.collect()\n    print()\n\nif len(submission_list) > 0:\n    submission = pd.concat(submission_list, ignore_index=True)\nelse:\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\nsubmission_robust = robustify(submission, test, 'test')\nsubmission_robust.index.name = 'row_id'\nsubmission_robust.to_csv('submission.csv')\nprint(f\"\\nSubmission created: {len(submission_robust)} predictions\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}