{"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,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Thanks to Taylor S. Amarel for the code this submission is based off,\n\nWhat changed from the versions.\nincreased n-estimstors and the learning rate:\nn_estimators: 150 -> 200\nlearning rate: 0.03 -> 0.06","metadata":{}},{"cell_type":"code","source":"# MABe Challenge - Social Action Recognition in Mice\n# Complete code with long-range temporal dependencies\n\nvalidate_or_submit = 'submit' # 'validate' or 'submit' or 'stresstest'\nverbose = True\n\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom tqdm import trange, tqdm\nimport itertools\nimport warnings\nimport json\nimport os\nimport lightgbm\nfrom collections import defaultdict\nimport polars as pl\n\nfrom sklearn.base import ClassifierMixin, BaseEstimator, clone\nfrom sklearn.model_selection import cross_val_predict, GroupKFold, train_test_split\nfrom sklearn.pipeline import make_pipeline\nfrom sklearn.impute import SimpleImputer\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.neighbors import KNeighborsClassifier\nfrom sklearn.metrics import f1_score\n\n# Custom classifier for training on subset\nclass TrainOnSubsetClassifier(ClassifierMixin, BaseEstimator):\n    \"\"\"Fit estimator to a subset of the training data.\"\"\"\n    def __init__(self, estimator, n_samples):\n        self.estimator = estimator\n        self.n_samples = n_samples\n\n    def fit(self, X, y):\n        downsample = len(X) // self.n_samples\n        downsample = max(downsample, 1)\n        self.estimator.fit(np.array(X, copy=False)[::downsample],\n                           np.array(y, copy=False)[::downsample])\n        self.classes_ = self.estimator.classes_\n        return self\n\n    def predict_proba(self, X):\n        if len(self.classes_) == 1:\n            return np.full((len(X), 1), 1.0)\n        probs = self.estimator.predict_proba(np.array(X))\n        return probs\n        \n    def predict(self, X):\n        return self.estimator.predict(np.array(X))\n\n# F-Beta scoring functions\nclass HostVisibleError(Exception):\n    pass\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\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\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# Load data\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)\ntrain_without_mabe22 = train.query(\"~ lab_id.str.startswith('MABe22_')\")\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\n# Solution dataframe creation function for validation\ndef create_solution_df(dataset):\n    solution = []\n    for _, row in tqdm(dataset.iterrows(), total=len(dataset)):\n        lab_id = row['lab_id']\n        if lab_id.startswith('MABe22'): continue\n        video_id = row['video_id']\n        path = f\"/kaggle/input/MABe-mouse-behavior-detection/train_annotation/{lab_id}/{video_id}.parquet\"\n        try:\n            annot = pd.read_parquet(path)\n        except FileNotFoundError:\n            if verbose: print(f\"No annotations for {path}\")\n            continue\n    \n        annot['lab_id'] = lab_id\n        annot['video_id'] = video_id\n        annot['behaviors_labeled'] = row['behaviors_labeled']\n        annot['target_id'] = np.where(annot.target_id != annot.agent_id, annot['target_id'].apply(lambda s: f\"mouse{s}\"), 'self')\n        annot['agent_id'] = annot['agent_id'].apply(lambda s: f\"mouse{s}\")\n        solution.append(annot)\n    \n    solution = pd.concat(solution)\n    return solution\n\nif validate_or_submit == 'validate':\n    solution = create_solution_df(train_without_mabe22)\n\n# Stress test code (only runs if validate_or_submit == 'stresstest')\nif validate_or_submit == 'stresstest':\n    n_videos_per_lab = 2\n    \n    try:\n        os.mkdir(f\"stresstest_tracking\")\n    except FileExistsError:\n        pass\n    \n    rng = np.random.default_rng()\n    stresstest = pd.concat(\n        [train.query(\"video_id == 1459695188\")]\n        + [df.sample(min(n_videos_per_lab, len(df)), random_state=1) for (_, df) in train.groupby('lab_id')])\n    for _, row in tqdm(stresstest.iterrows(), total=len(stresstest)):\n        lab_id = row['lab_id']\n        video_id = row['video_id']\n        \n        path = f\"/kaggle/input/MABe-mouse-behavior-detection/train_tracking/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n    \n        if video_id == 1459695188:\n            vid = pd.concat([vid] * 3)\n            vid['video_frame'] = np.arange(len(vid))\n    \n        dropped_frames = list(rng.choice(np.unique(vid.video_frame), size=100, replace=False))\n        vid = vid.query(\"~ video_frame.isin(@dropped_frames)\")\n        \n        if rng.uniform() < 0.2:\n            dropped_bodypart = rng.choice(np.unique(vid.bodypart), size=1, replace=False)[0]\n            vid = vid.query(\"bodypart != @dropped_bodypart\")\n        \n        if rng.uniform() < 0.1:\n            vid = vid.query(\"mouse_id != 1\")\n        \n        if rng.uniform() < 0.7:\n            mask = np.ones(len(vid), dtype=bool)\n            mask[:int(0.4 * len(mask))] = False\n            rng.shuffle(mask)\n            vid = vid[mask]\n    \n        if rng.uniform() < 0.7:\n            mask = np.ones(len(vid), dtype=bool)\n            mask[:int(0.2 * len(mask))] = False\n            rng.shuffle(mask)\n            vid.loc[:, 'x'] = np.where(mask, np.nan, vid.loc[:, 'x'])\n            rng.shuffle(mask)\n            vid.loc[:, 'y'] = np.where(mask, np.nan, vid.loc[:, 'y'])\n    \n        try:\n            os.mkdir(f\"stresstest_tracking/{lab_id}\")\n        except FileExistsError:\n            pass\n        new_path = f\"stresstest_tracking/{lab_id}/{video_id}.parquet\"\n        vid.to_parquet(new_path)\n\n# Body parts to drop for memory efficiency\ndrop_body_parts = ['headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', \n                   'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft', 'headpiece_topfrontright', \n                   'spine_1', 'spine_2',\n                   'tail_middle_1', 'tail_middle_2', 'tail_midpoint']\n\n# Generate mouse data function\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        if lab_id.startswith('MABe22'): continue\n        video_id = row.video_id\n\n        if type(row.behaviors_labeled) != str:\n            print('No labeled behaviors:', lab_id, video_id, type(row.behaviors_labeled), row.behaviors_labeled)\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': print('video with missing values', video_id, traintest, len(vid), 'frames')\n        else:\n            if verbose and traintest == 'test': 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                    })\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: 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                    })\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: print('- test pair', video_id, agent, target)\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n\n# Threshold for multiclass prediction\nthreshold = 0.27\n\n# Predict multiclass function\ndef predict_multiclass(pred, meta):\n    ama = np.argmax(pred, axis=1)\n    ama = np.where(pred.max(axis=1) >= threshold, ama, -1)\n    ama = pd.Series(ama, index=meta.video_frame)\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    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    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        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    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# Transform functions with LONG-RANGE TEMPORAL DEPENDENCIES\ndef transform_single(single_mouse, body_parts_tracked):\n    available_body_parts = single_mouse.columns.get_level_values(0)\n    X = pd.DataFrame({\n            f\"{part1}+{part2}\": np.square(single_mouse[part1] - single_mouse[part2]).sum(axis=1, skipna=False)\n            for part1, part2 in itertools.combinations(body_parts_tracked, 2) if part1 in available_body_parts and part2 in available_body_parts\n        })\n    X = X.reindex(columns=[f\"{part1}+{part2}\" for part1, part2 in itertools.combinations(body_parts_tracked, 2)], copy=False)\n\n    if 'ear_left' in single_mouse.columns and 'ear_right' in single_mouse.columns and 'tail_base' in single_mouse.columns:\n        shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(10)\n        X = pd.concat([\n            X, \n            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            })\n        ], axis=1)\n    \n    # Add relative distances\n    if 'nose+tail_base' in X.columns and 'ear_left+ear_right' in X.columns:\n        X['elongation_ratio'] = X['nose+tail_base'] / (X['ear_left+ear_right'] + 1e-6)\n    \n    # Body angle/curvature\n    if 'nose' in available_body_parts and 'body_center' in available_body_parts and 'tail_base' in available_body_parts:\n        v1 = single_mouse['nose'] - single_mouse['body_center']\n        v2 = single_mouse['tail_base'] - single_mouse['body_center']\n        \n        dot_product = (v1['x'] * v2['x'] + v1['y'] * v2['y'])\n        norm_v1 = np.sqrt(v1['x']**2 + v1['y']**2)\n        norm_v2 = np.sqrt(v2['x']**2 + v2['y']**2)\n        \n        X['body_angle_cos'] = dot_product / (norm_v1 * norm_v2 + 1e-6)\n    \n    # LONG-RANGE TEMPORAL DEPENDENCIES\n    # Multiple time windows for capturing patterns at different scales\n    time_windows = [5, 15, 30, 60]  # Short, medium, long, very long range\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 over different time windows\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 in window\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            # Cumulative displacement over window\n            X[f'cumulative_disp_{window}'] = np.sqrt(\n                center_x.diff().rolling(window, min_periods=1).sum()**2 + \n                center_y.diff().rolling(window, min_periods=1).sum()**2\n            )\n            \n            # Activity level (movement variance) over window\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 at multiple time offsets\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    # Temporal context features\n    if 'ear_left' in available_body_parts and 'ear_right' in available_body_parts:\n        ear_dist = np.sqrt(\n            (single_mouse['ear_left']['x'] - single_mouse['ear_right']['x'])**2 + \n            (single_mouse['ear_left']['y'] - single_mouse['ear_right']['y'])**2\n        )\n        \n        # Look both forward and backward in time\n        for offset in [-20, -10, 10, 20]:\n            X[f'ear_dist_offset_{offset}'] = ear_dist.shift(-offset)\n        \n        # Temporal consistency (how stable is the feature over time)\n        X['ear_dist_consistency_30'] = ear_dist.rolling(30, min_periods=1, center=True).std() / (ear_dist.rolling(30, min_periods=1, center=True).mean() + 1e-6)\n    \n    return X#*1.04\n\ndef transform_pair(mouse_pair, body_parts_tracked):\n    available_body_parts_A = mouse_pair['A'].columns.get_level_values(0)\n    available_body_parts_B = mouse_pair['B'].columns.get_level_values(0)\n    X = pd.DataFrame({\n            f\"12+{part1}+{part2}\": np.square(mouse_pair['A'][part1] - mouse_pair['B'][part2]).sum(axis=1, skipna=False)\n            for part1, part2 in itertools.product(body_parts_tracked, repeat=2) if part1 in available_body_parts_A and part2 in available_body_parts_B\n        })\n    X = X.reindex(columns=[f\"12+{part1}+{part2}\" for part1, part2 in itertools.product(body_parts_tracked, repeat=2)], copy=False)\n\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        X = pd.concat([\n            X,\n            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            })\n        ], axis=1)\n    \n    # Add relative distances\n    if 'nose+tail_base' in X.columns and 'ear_left+ear_right' in X.columns:\n        X['elongation_ratio'] = X['nose+tail_base'] / (X['ear_left+ear_right'] + 1e-6)\n    \n    # Relative orientation between mice\n    if 'nose' in available_body_parts_A and 'tail_base' in available_body_parts_A and \\\n       'nose' in available_body_parts_B and 'tail_base' in available_body_parts_B:\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        \n        dot_product = (dir_A['x'] * dir_B['x'] + dir_A['y'] * dir_B['y'])\n        norm_A = np.sqrt(dir_A['x']**2 + dir_A['y']**2)\n        norm_B = np.sqrt(dir_B['x']**2 + dir_B['y']**2)\n        \n        X['relative_orientation_cos'] = dot_product / (norm_A * norm_B + 1e-6)\n    \n    # Approach rate\n    if 'nose' in available_body_parts_A and 'nose' in available_body_parts_B:\n        current_dist = np.square(mouse_pair['A']['nose'] - mouse_pair['B']['nose']).sum(axis=1, skipna=False)\n        \n        shifted_A_nose = mouse_pair['A']['nose'].shift(10)\n        shifted_B_nose = mouse_pair['B']['nose'].shift(10)\n        past_dist = np.square(shifted_A_nose - shifted_B_nose).sum(axis=1, skipna=False)\n        \n        X['approach_rate'] = current_dist - past_dist\n    \n    # Social zone indicators\n    if 'body_center' in available_body_parts_A and 'body_center' in available_body_parts_B:\n        center_dist = np.sqrt(\n            (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        )\n        \n        X['very_close'] = (center_dist < 5.0).astype(float)\n        X['close'] = ((center_dist >= 5.0) & (center_dist < 15.0)).astype(float)\n        X['medium'] = ((center_dist >= 15.0) & (center_dist < 30.0)).astype(float)\n        X['far'] = (center_dist >= 30.0).astype(float)\n    \n    # LONG-RANGE TEMPORAL DEPENDENCIES FOR PAIRS\n    time_windows = [5, 15, 30, 60]\n    \n    if 'body_center' in available_body_parts_A and 'body_center' in available_body_parts_B:\n        # Inter-mouse distance over time\n        center_dist_full = np.square(mouse_pair['A']['body_center'] - mouse_pair['B']['body_center']).sum(axis=1, skipna=False)\n        \n        for window in time_windows:\n            # Distance statistics over windows\n            X[f'dist_mean_{window}'] = center_dist_full.rolling(window, min_periods=1, center=True).mean()\n            X[f'dist_std_{window}'] = center_dist_full.rolling(window, min_periods=1, center=True).std()\n            X[f'dist_min_{window}'] = center_dist_full.rolling(window, min_periods=1, center=True).min()\n            X[f'dist_max_{window}'] = center_dist_full.rolling(window, min_periods=1, center=True).max()\n            \n            # Interaction intensity (inverse of distance variance)\n            dist_var = center_dist_full.rolling(window, min_periods=1, center=True).var()\n            X[f'interaction_intensity_{window}'] = 1 / (1 + dist_var)\n            \n            # Coordinated movement over window\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            X[f'coord_movement_std_{window}'] = coord_movement.rolling(window, min_periods=1, center=True).std()\n    \n    # Lag features for interaction patterns\n    if 'nose' in available_body_parts_A and 'nose' in available_body_parts_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 (how long they stay close)\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    # Temporal context for social behaviors\n    if 'body_center' in available_body_parts_A and 'body_center' in available_body_parts_B:\n        # Relative velocity alignment over time\n        A_vel_x = mouse_pair['A']['body_center']['x'].diff()\n        A_vel_y = mouse_pair['A']['body_center']['y'].diff()\n        B_vel_x = mouse_pair['B']['body_center']['x'].diff()\n        B_vel_y = mouse_pair['B']['body_center']['y'].diff()\n        \n        vel_alignment = (A_vel_x * B_vel_x + A_vel_y * B_vel_y) / (\n            np.sqrt(A_vel_x**2 + A_vel_y**2) * np.sqrt(B_vel_x**2 + B_vel_y**2) + 1e-6\n        )\n        \n        # Look at alignment patterns over different time horizons\n        for offset in [-20, -10, 0, 10, 20]:\n            X[f'vel_alignment_offset_{offset}'] = vel_alignment.shift(-offset)\n        \n        # Temporal consistency of interaction\n        X['interaction_consistency_30'] = center_dist_full.rolling(30, min_periods=1, center=True).std() / (center_dist_full.rolling(30, min_periods=1, center=True).mean() + 1e-6)\n    \n    return X#*1.04\n\n# Cross-validation function\ndef cross_validate_classifier(binary_classifier, X, label, meta):\n    oof = pd.DataFrame(index=meta.video_frame)\n    for action in label.columns:\n        action_mask = ~ label[action].isna().values\n        X_action = X[action_mask]\n        y_action = label[action][action_mask].values.astype(int)\n        p = y_action.mean()\n        baseline_score = p / (1 + p)\n        groups_action = meta.video_id[action_mask]\n        if len(np.unique(groups_action)) < 5:\n            continue\n\n        if not (y_action == 0).all():\n            with warnings.catch_warnings():\n                warnings.filterwarnings('ignore', category=RuntimeWarning)\n                oof_action = cross_val_predict(binary_classifier, X_action, y_action, groups=groups_action, cv=GroupKFold(), method='predict_proba')\n            oof_action = oof_action[:, 1]\n        else:\n            oof_action = np.zeros(len(y_action))\n        f1 = f1_score(y_action, (oof_action >= threshold), zero_division=0)\n        ch = '>' if f1 > baseline_score else '=' if f1 == baseline_score else '<'\n        print(f\"  F1: {f1:.3f} {ch} ({baseline_score:.3f}) {action}\")\n        f1_list.append((body_parts_tracked_str, action, f1))\n        oof_column = np.zeros(len(label))\n        oof_column[action_mask] = oof_action\n        oof[action] = oof_column\n\n    submission_part = predict_multiclass(oof, meta)\n    submission_list.append(submission_part)\n\n# Submit function\ndef submit(body_parts_tracked_str, switch_tr, binary_classifier, X_tr, label, meta):\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\n        if not (y_action == 0).all():\n            model = clone(binary_classifier)\n            model.fit(X_tr[action_mask], y_action)\n            assert len(model.classes_) == 2\n            model_list.append((action, model))\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    if validate_or_submit == 'submit':\n        test_subset = test[test.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    else:\n        test_subset = stresstest.query(\"body_parts_tracked == @body_parts_tracked_str\")\n        generator = generate_mouse_data(test_subset, 'test',\n                                        traintest_directory='stresstest_tracking',\n                                        generate_single=(switch_tr == 'single'),\n                                        generate_pair=(switch_tr == 'pair'))\n    if verbose: print(f\"n_videos: {len(test_subset)}\")\n    for switch_te, data_te, meta_te, actions_te in generator:\n        assert switch_te == switch_tr\n        try:\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            if verbose and len(X_te) == 0: 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            if pred.shape[1] != 0:\n                submission_part = predict_multiclass(pred, meta_te)\n                submission_list.append(submission_part)\n            else:\n                if verbose: print(f\"  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# Robustify function\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    old_submission = submission.copy()\n    submission = submission[submission.start_frame < submission.stop_frame]\n    if len(submission) != len(old_submission):\n        print(\"ERROR: Dropped frames with start >= stop\")\n    \n    old_submission = submission.copy()\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_frame = 0\n        for i, (_, row) in enumerate(group.iterrows()):\n            if row['start_frame'] < last_stop_frame:\n                mask[i] = False\n            else:\n                last_stop_frame = row['stop_frame']\n        group_list.append(group[mask])\n    submission = pd.concat(group_list)\n    if len(submission) != len(old_submission):\n        print(\"ERROR: Dropped duplicate frames\")\n\n    s_list = []\n    for idx, row in dataset.iterrows():\n        lab_id = row['lab_id']\n        if 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: 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_length = 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_length\n                batch_stop = min(batch_start + batch_length, 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        print(\"ERROR: Filled empty videos\")\n\n    submission = submission.reset_index(drop=True)\n    return submission\n\n# Main processing loop\nf1_list = []\nsubmission_list = []\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 videos with {body_parts_tracked}\")\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        single_mouse_list = []\n        single_mouse_label_list = []\n        single_mouse_meta_list = []\n        mouse_pair_list = []\n        mouse_pair_label_list = []\n        mouse_pair_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_mouse_meta_list.append(meta)\n                single_mouse_label_list.append(label)\n            else:\n                mouse_pair_list.append(data)\n                mouse_pair_meta_list.append(meta)\n                mouse_pair_label_list.append(label)\n    \n        binary_classifier = make_pipeline(\n            SimpleImputer(),\n            TrainOnSubsetClassifier(\n                lightgbm.LGBMClassifier(\n                    n_estimators=200, \n                    learning_rate=0.06,\n                    min_child_samples=40,\n                    num_leaves=31,\n                    max_depth=-1,\n                    subsample=0.8,\n                    colsample_bytree=0.8,\n                    verbose=-1),\n                100000)\n        )\n    \n        if len(single_mouse_list) > 0:\n            single_mouse = pd.concat(single_mouse_list)\n            single_mouse_label = pd.concat(single_mouse_label_list)\n            single_mouse_meta = pd.concat(single_mouse_meta_list)\n            del single_mouse_list, single_mouse_label_list, single_mouse_meta_list\n            assert len(single_mouse) == len(single_mouse_label)\n            assert len(single_mouse) == len(single_mouse_meta)\n            \n            X_tr = transform_single(single_mouse, body_parts_tracked)\n            del single_mouse\n            print(f\"{X_tr.shape=}\")\n    \n            if validate_or_submit == 'validate':\n                cross_validate_classifier(binary_classifier, X_tr, single_mouse_label, single_mouse_meta)\n            else:\n                submit(body_parts_tracked_str, 'single', binary_classifier, X_tr, single_mouse_label, single_mouse_meta)\n            del X_tr\n                \n        if len(mouse_pair_list) > 0:\n            mouse_pair = pd.concat(mouse_pair_list)\n            mouse_pair_label = pd.concat(mouse_pair_label_list)\n            mouse_pair_meta = pd.concat(mouse_pair_meta_list)\n            del mouse_pair_list, mouse_pair_label_list, mouse_pair_meta_list\n            assert len(mouse_pair) == len(mouse_pair_label)\n            assert len(mouse_pair) == len(mouse_pair_meta)\n        \n            X_tr = transform_pair(mouse_pair, body_parts_tracked)\n            del mouse_pair\n            print(f\"{X_tr.shape=}\")\n    \n            if validate_or_submit == 'validate':\n                cross_validate_classifier(binary_classifier, X_tr, mouse_pair_label, mouse_pair_meta)\n            else:\n                submit(body_parts_tracked_str, 'pair', binary_classifier, X_tr, mouse_pair_label, mouse_pair_meta)\n            del X_tr\n                \n    except Exception as e:\n        print(f'***Exception*** {e}')\n    print()\n\n# Final submission creation\nif validate_or_submit == 'validate':\n    submission = pd.concat(submission_list)\n    submission_robust = robustify(submission, train, 'train')\n    print(f\"# OOF score with competition metric: {score(solution, submission_robust, ''):.4f}\")\n\n    f1_df = pd.DataFrame(f1_list, columns=['body_parts_tracked_str', 'action', 'binary F1 score'])\n    print(f\"# Average of {len(f1_df)} binary F1 scores {f1_df['binary F1 score'].mean():.4f}\")\n\nif validate_or_submit != 'validate':\n    if len(submission_list) > 0:\n        submission = pd.concat(submission_list)\n    else:\n        submission = pd.DataFrame(\n            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    if validate_or_submit == 'submit':\n        submission_robust = robustify(submission, test, 'test')\n    else:\n        submission_robust = robustify(submission, stresstest, 'stresstest', 'stresstest_tracking')\n    submission_robust.index.name = 'row_id'\n    submission_robust.to_csv('submission.csv')\n    print(\"Submission file created: submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T02:45:33.174931Z","iopub.execute_input":"2025-09-30T02:45:33.175168Z","execution_failed":"2025-09-30T02:59:25.351Z"}},"outputs":[],"execution_count":null}]}