{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13719397,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport json\nimport itertools\nimport warnings\nimport os\nfrom pathlib import Path\nfrom sklearn.model_selection import cross_val_predict, GroupKFold\nfrom sklearn.preprocessing import StandardScaler, LabelEncoder\nfrom sklearn.impute import SimpleImputer\nfrom sklearn.pipeline import make_pipeline\nfrom sklearn.metrics import f1_score\nfrom sklearn.base import clone\nfrom sklearn.ensemble import VotingClassifier, RandomForestClassifier\nfrom sklearn.neighbors import KNeighborsClassifier\nfrom lightgbm import LGBMClassifier\nimport catboost as ctb\nfrom tqdm.auto import tqdm\n\nwarnings.filterwarnings('ignore')\n\nvalidate_or_submit = 'submit'\nverbose = True\n\nprint(\"Loading training and test metadata...\")\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_')\")\ntest = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\nclass TrainOnSubsetClassifier:\n    def __init__(self, clf, max_samples):\n        self.clf = clf\n        self.max_samples = max_samples\n    \n    def fit(self, X, y):\n        if len(X) > self.max_samples:\n            idx = np.random.choice(len(X), self.max_samples, replace=False)\n            X_subset = X[idx] if hasattr(X, 'iloc') else X[idx]\n            y_subset = y[idx] if hasattr(y, 'iloc') else y[idx]\n        else:\n            X_subset, y_subset = X, y\n        self.clf.fit(X_subset, y_subset)\n        return self\n    \n    def predict_proba(self, X):\n        return self.clf.predict_proba(X)\n    \n    @property\n    def classes_(self):\n        return self.clf.classes_\n\ndef transform_single_enhanced(single_mouse, body_parts_tracked):\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)\n    })\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        speed_features = 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        \n        shifted_20 = single_mouse[['ear_left', 'ear_right']].shift(20)\n        if not shifted_20.isna().all().all():\n            accel_features = pd.DataFrame({\n                'accel_left': speed_features['speed_left'] - np.square(shifted['ear_left'] - shifted_20['ear_left']).sum(axis=1, skipna=False),\n                'accel_right': speed_features['speed_right'] - np.square(shifted['ear_right'] - shifted_20['ear_right']).sum(axis=1, skipna=False),\n            })\n            X = pd.concat([X, speed_features, accel_features], axis=1)\n        else:\n            X = pd.concat([X, speed_features], axis=1)\n    \n    return X\n\ndef transform_pair_enhanced(mouse_pair, body_parts_tracked):\n    drop_body_parts = ['ear_left', 'ear_right',\n                      'headpiece_bottombackleft', 'headpiece_bottombackright', \n                      'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', \n                      'headpiece_topbackleft', 'headpiece_topbackright', \n                      'headpiece_topfrontleft', 'headpiece_topfrontright', \n                      'tail_midpoint']\n    \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    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)\n    })\n\n    if 'nose' in body_parts_tracked and 'tail_base' in body_parts_tracked:\n        X['face_distance'] = np.square(mouse_pair['A']['nose'] - mouse_pair['B']['nose']).sum(axis=1, skipna=False)\n        X['following'] = np.square(mouse_pair['A']['nose'] - mouse_pair['B']['tail_base']).sum(axis=1, skipna=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    return X\n\ndef transform_single_original(single_mouse, body_parts_tracked):\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)\n        })\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    return X\n\ndef transform_pair_original(mouse_pair, body_parts_tracked):\n    drop_body_parts =  ['ear_left', 'ear_right',\n                        'headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', \n                        'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft', 'headpiece_topfrontright', \n                        'tail_midpoint']\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    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)\n        })\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    return X\n\ndef find_optimal_thresholds(oof_predictions, labels, default_threshold=0.27):\n    optimal_thresholds = {}\n    \n    for action in oof_predictions.columns:\n        if action in labels.columns:\n            mask = ~labels[action].isna()\n            if mask.sum() > 100:\n                y_true = labels[action][mask].values.astype(int)\n                y_pred_proba = oof_predictions[action][mask].values\n                \n                best_f1 = 0\n                best_thresh = default_threshold\n                \n                for thresh in [0.15, 0.2, 0.25, 0.27, 0.3, 0.35, 0.4, 0.45, 0.5]:\n                    f1 = f1_score(y_true, y_pred_proba >= thresh, zero_division=0)\n                    if f1 > best_f1:\n                        best_f1 = f1\n                        best_thresh = thresh\n                        \n                optimal_thresholds[action] = best_thresh\n            else:\n                optimal_thresholds[action] = default_threshold\n    \n    return optimal_thresholds\n\ndef create_ensemble_model():\n    lgb = LGBMClassifier(\n        n_estimators=500,\n        max_depth=6, \n        learning_rate=0.05,\n        random_state=42,\n        verbosity=-1,\n        force_row_wise=True\n    )\n    \n    rf = RandomForestClassifier(\n        n_estimators=200,\n        max_depth=8,\n        random_state=42,\n        n_jobs=-1\n    )\n    \n    ensemble = VotingClassifier([('lgb', lgb), ('rf', rf)], voting='soft')\n    \n    return make_pipeline(\n        SimpleImputer(),\n        StandardScaler(), \n        TrainOnSubsetClassifier(ensemble, 20000)\n    )\n\ndef create_knn_model():\n    return make_pipeline(\n        SimpleImputer(),\n        StandardScaler(),\n        TrainOnSubsetClassifier(KNeighborsClassifier(n_neighbors=11), 20000)\n    )\n\ndef predict_multiclass_optimized(pred, meta, thresholds=None):\n    if thresholds is None:\n        thresholds = {col: 0.27 for col in pred.columns}\n    \n    ama = np.argmax(pred.values, axis=1)\n    max_proba = pred.max(axis=1).values\n    \n    threshold_array = np.array([thresholds.get(col, 0.27) for col in pred.columns])\n    action_thresholds = threshold_array[ama]\n    \n    ama = np.where(max_proba >= action_thresholds, 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    \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 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    \n    assert (submission_part.stop_frame > submission_part.start_frame).all(), 'stop <= start'\n    if verbose: \n        print('  actions found:', len(submission_part))\n        \n    return submission_part\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        \n    for _, 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        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        \n        vid = pd.read_parquet(path)\n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        \n        if pvid.isna().any().any():\n            if verbose and traintest == 'test': \n                print('video with missing values', video_id, traintest, len(vid), 'frames')\n        else:\n            if verbose and traintest == 'test': \n                print('video with all values', video_id, traintest, len(vid), 'frames')\n        \n        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                    \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                    \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                    \n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\n                    mouse_pair_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': agent_str,\n                        'target_id': target_str,\n                        'video_frame': mouse_pair.index\n                    })\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\ndef submit_method1_and_2(body_parts_tracked_str, switch_tr, models_dict, features_dict, labels, meta):\n    predictions_dict = {'ensemble': [], 'knn': []}\n    \n    model_lists = {'ensemble': [], 'knn': []}\n    \n    for action in labels.columns:\n        action_mask = ~labels[action].isna().values\n        y_action = labels[action][action_mask].values.astype(int)\n\n        if ~(y_action == 0).all():\n            model_ensemble = clone(models_dict['ensemble'])\n            model_ensemble.fit(features_dict['enhanced'][action_mask], y_action)\n            model_lists['ensemble'].append((action, model_ensemble))\n            \n            model_knn = clone(models_dict['knn'])\n            model_knn.fit(features_dict['original'][action_mask], y_action)\n            model_lists['knn'].append((action, model_knn))\n\n    body_parts_tracked = json.loads(body_parts_tracked_str)\n    \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    \n    if verbose: \n        print(f\"n_videos: {len(test_subset)}\")\n    \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_enhanced = transform_single_enhanced(data_te, body_parts_tracked)\n                X_te_original = transform_single_original(data_te, body_parts_tracked)\n            else:\n                X_te_enhanced = transform_pair_enhanced(data_te, body_parts_tracked)\n                X_te_original = transform_pair_original(data_te, body_parts_tracked)\n                \n            if verbose and (len(X_te_enhanced) == 0 or len(X_te_original) == 0): \n                print(\"ERROR: X_te is empty\")\n            del data_te\n\n            pred_ensemble = pd.DataFrame(index=meta_te.video_frame)\n            pred_knn = pd.DataFrame(index=meta_te.video_frame)\n            \n            for action, model in model_lists['ensemble']:\n                if action in actions_te:\n                    pred_ensemble[action] = model.predict_proba(X_te_enhanced)[:, 1]\n                    \n            for action, model in model_lists['knn']:\n                if action in actions_te:\n                    pred_knn[action] = model.predict_proba(X_te_original)[:, 1]\n            \n            del X_te_enhanced, X_te_original\n\n            if pred_ensemble.shape[1] != 0:\n                submission_part_ensemble = predict_multiclass_optimized(pred_ensemble, meta_te)\n                predictions_dict['ensemble'].append(submission_part_ensemble)\n            \n            if pred_knn.shape[1] != 0:\n                submission_part_knn = predict_multiclass_optimized(pred_knn, meta_te)\n                predictions_dict['knn'].append(submission_part_knn)\n                    \n        except KeyError:\n            if verbose: \n                print(f'  ERROR: KeyError because of missing bodypart ({switch_tr})')\n            if 'data_te' in locals():\n                del data_te\n    \n    return predictions_dict\n\ndef create_mouse_features(mouse_data, suffix):\n    if len(mouse_data) == 0:\n        return {}\n    \n    features = {}\n    \n    features[f'mean_x_{suffix}'] = mouse_data['x'].mean()\n    features[f'mean_y_{suffix}'] = mouse_data['y'].mean()\n    features[f'std_x_{suffix}'] = mouse_data['x'].std()\n    features[f'std_y_{suffix}'] = mouse_data['y'].std()\n    \n    mouse_data = mouse_data.sort_values('video_frame')\n    dx = mouse_data['x'].diff()\n    dy = mouse_data['y'].diff()\n    speed = np.sqrt(dx**2 + dy**2)\n    \n    features[f'speed_mean_{suffix}'] = speed.mean()\n    features[f'speed_std_{suffix}'] = speed.std()\n    features[f'speed_max_{suffix}'] = speed.max()\n    features[f'total_distance_{suffix}'] = speed.sum()\n    \n    acceleration = speed.diff()\n    features[f'accel_mean_{suffix}'] = acceleration.mean()\n    features[f'accel_std_{suffix}'] = acceleration.std()\n    \n    return features\n\ndef create_interaction_features(agent_data, target_data, agent_id, target_id):\n    if len(agent_data) == 0 or len(target_data) == 0:\n        return None\n    \n    merged = pd.merge(agent_data, target_data, on='video_frame', \n                     suffixes=('_agent', '_target'), how='inner')\n    \n    if len(merged) == 0:\n        return None\n    \n    features = {}\n    \n    agent_features = create_mouse_features(agent_data, 'agent')\n    features.update(agent_features)\n    \n    target_features = create_mouse_features(target_data, 'target')\n    features.update(target_features)\n    \n    merged['distance'] = np.sqrt(\n        (merged['x_agent'] - merged['x_target'])**2 + \n        (merged['y_agent'] - merged['y_target'])**2\n    )\n    \n    agent_speed = np.sqrt(merged['x_agent'].diff()**2 + merged['y_agent'].diff()**2)\n    target_speed = np.sqrt(merged['x_target'].diff()**2 + merged['y_target'].diff()**2)\n    \n    features['distance_mean'] = merged['distance'].mean()\n    features['distance_std'] = merged['distance'].std()\n    features['distance_min'] = merged['distance'].min()\n    features['distance_max'] = merged['distance'].max()\n    \n    features['relative_speed_mean'] = (agent_speed - target_speed).abs().mean()\n    features['speed_ratio'] = agent_speed.mean() / (target_speed.mean() + 1e-8)\n    \n    if len(merged) > 1:\n        agent_dx = merged['x_agent'].iloc[-1] - merged['x_agent'].iloc[0]\n        agent_dy = merged['y_agent'].iloc[-1] - merged['y_agent'].iloc[0]\n        target_dx = merged['x_target'].iloc[-1] - merged['x_target'].iloc[0]\n        target_dy = merged['y_target'].iloc[-1] - merged['y_target'].iloc[0]\n        \n        features['agent_movement'] = np.sqrt(agent_dx**2 + agent_dy**2)\n        features['target_movement'] = np.sqrt(target_dx**2 + target_dy**2)\n    \n    features['agent_id'] = agent_id\n    features['target_id'] = target_id\n    features['frames_count'] = len(merged)\n    \n    return features\n\ndef prepare_test_dataset_method3(segment_size=50):\n    all_test_features = []\n    \n    base_path = Path('/kaggle/input/MABe-mouse-behavior-detection')\n    test_track_path = base_path / 'test_tracking'\n    \n    for lab_id in os.listdir(test_track_path):\n        lab_track_path = test_track_path / lab_id\n        \n        for track_file in os.listdir(lab_track_path):\n            if track_file.endswith('.parquet'):\n                video_id = track_file.replace('.parquet', '')\n                \n                try:\n                    track_df = pd.read_parquet(lab_track_path / track_file)\n                    all_frames = sorted(track_df['video_frame'].unique())\n                    \n                    for i in range(0, len(all_frames), segment_size):\n                        seg_start = all_frames[i]\n                        seg_end = all_frames[min(i + segment_size - 1, len(all_frames) - 1)]\n                        \n                        if seg_end - seg_start < 10:\n                            continue\n                            \n                        segment_data = track_df[\n                            (track_df['video_frame'] >= seg_start) & \n                            (track_df['video_frame'] <= seg_end)\n                        ]\n                        \n                        mice = segment_data['mouse_id'].unique()\n                        if len(mice) == 0:\n                            continue\n                        \n                        for agent in mice:\n                            for target in mice:\n                                agent_data = segment_data[segment_data['mouse_id'] == agent]\n                                target_data = segment_data[segment_data['mouse_id'] == target]\n                                \n                                features = create_interaction_features(agent_data, target_data, agent, target)\n                                \n                                if features is not None:\n                                    features['video_id'] = video_id\n                                    features['start_frame'] = seg_start\n                                    features['stop_frame'] = seg_end\n                                    features['segment_size'] = seg_end - seg_start + 1\n                                    \n                                    all_test_features.append(features)\n                    \n                except Exception as e:\n                    print(f\"Video {video_id} error: {e}\")\n                    continue\n    \n    test_features_df = pd.DataFrame(all_test_features).fillna(0)\n    \n    return test_features_df\n\ndef train_catboost_models():\n    print(\"Training CatBoost models for Method 3...\")\n    \n    base_path = Path('/kaggle/input/MABe-mouse-behavior-detection')\n    \n    if Path('/kaggle/input/mice-dataset/train_features.parquet').exists():\n        train_features = pd.read_parquet('/kaggle/input/mice-dataset/train_features.parquet')\n        train_labels = pd.read_csv('/kaggle/input/mice-dataset/train_labels.csv')\n    else:\n        print(\"Preparing training data for Method 3...\")\n        ann_path = base_path / 'train_annotation'\n        train_track_path = base_path / 'train_tracking'\n        \n        all_features = []\n        all_labels = []\n        \n        for lab_id in tqdm(os.listdir(ann_path)):\n            lab_ann_path = ann_path / lab_id\n            lab_track_path = train_track_path / lab_id\n            \n            if not lab_ann_path.exists() or not lab_track_path.exists():\n                continue\n                \n            for ann_file in os.listdir(lab_ann_path):\n                if ann_file.endswith('.parquet'):\n                    video_id = ann_file.replace('.parquet', '')\n                    \n                    try:\n                        ann_df = pd.read_parquet(lab_ann_path / ann_file)\n                        track_df = pd.read_parquet(lab_track_path / ann_file)\n                        \n                        annotated_segments = set()\n                        for _, ann in ann_df.iterrows():\n                            start, stop = ann['start_frame'], ann['stop_frame']\n                            for seg_start in range(start, stop + 1, 50):\n                                seg_end = min(seg_start + 50 - 1, stop)\n                                if seg_end - seg_start >= 10:\n                                    annotated_segments.add((seg_start, seg_end, ann['agent_id'], ann['target_id'], ann['action']))\n                        \n                        for seg_start, seg_end, agent_id, target_id, action in annotated_segments:\n                            segment_data = track_df[\n                                (track_df['video_frame'] >= seg_start) & \n                                (track_df['video_frame'] <= seg_end)\n                            ]\n                            \n                            agent_data = segment_data[segment_data['mouse_id'] == agent_id]\n                            target_data = segment_data[segment_data['mouse_id'] == target_id]\n                            \n                            features = create_interaction_features(agent_data, target_data, agent_id, target_id)\n                            \n                            if features is not None:\n                                features['video_id'] = video_id\n                                features['start_frame'] = seg_start\n                                features['stop_frame'] = seg_end\n                                features['segment_size'] = seg_end - seg_start + 1\n                                \n                                all_features.append(features)\n                                all_labels.append(action)\n                        \n                        all_frames = sorted(track_df['video_frame'].unique())\n                        annotated_frames = set()\n                        \n                        for seg_start, seg_end, _, _, _ in annotated_segments:\n                            annotated_frames.update(range(seg_start, seg_end + 1))\n                        \n                        for i in range(0, len(all_frames), 50):\n                            if i + 50 > len(all_frames):\n                                break\n                                \n                            seg_start = all_frames[i]\n                            seg_end = all_frames[min(i + 50 - 1, len(all_frames) - 1)]\n                            \n                            segment_frames = set(range(seg_start, seg_end + 1))\n                            if segment_frames & annotated_frames:\n                                continue\n                            \n                            segment_data = track_df[\n                                (track_df['video_frame'] >= seg_start) & \n                                (track_df['video_frame'] <= seg_end)\n                            ]\n                            \n                            mice = segment_data['mouse_id'].unique()\n                            if len(mice) == 0:\n                                continue\n                                \n                            for agent in mice[:2]:\n                                for target in mice[:2]:\n                                    if agent == target:\n                                        agent_data = segment_data[segment_data['mouse_id'] == agent]\n                                        target_data = segment_data[segment_data['mouse_id'] == target]\n                                        \n                                        features = create_interaction_features(agent_data, target_data, agent, target)\n                                        \n                                        if features is not None:\n                                            features['video_id'] = video_id\n                                            features['start_frame'] = seg_start\n                                            features['stop_frame'] = seg_end\n                                            features['segment_size'] = seg_end - seg_start + 1\n                                            \n                                            all_features.append(features)\n                                            all_labels.append('no_action')\n                            \n                    except Exception as e:\n                        print(f\"Video {video_id} error: {e}\")\n                        continue\n        \n        train_features = pd.DataFrame(all_features).fillna(0)\n        train_labels = pd.Series(all_labels)\n    \n    drop_cols = ['video_id']\n    scores = []\n    cv = GroupKFold(n_splits=5)\n    models = []\n    \n    for i, (tr_idx, val_idx) in enumerate(cv.split(train_features, groups=train_features['video_id'])):\n        X_tr, X_val = train_features.iloc[tr_idx], train_features.iloc[val_idx]\n        y_tr, y_val = train_labels.iloc[tr_idx], train_labels.iloc[val_idx]\n        model = ctb.CatBoostClassifier(verbose=100, task_type='GPU', random_state=0, eval_metric='TotalF1:average=Macro', cat_features=['target_id', 'agent_id'])\n        model.fit(X_tr.drop(columns=drop_cols), y_tr, eval_set=(X_val.drop(columns=drop_cols), y_val), early_stopping_rounds=50)\n        models.append(model)\n        scores.append(f1_score(y_val, model.predict(X_val.drop(columns=drop_cols)), average='macro'))\n    \n    print(f\"CatBoost CV Score: {np.mean(scores):.4f} ± {np.std(scores):.4f}\")\n    \n    return models\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    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: \n            print(f\"Video {video_id} has no predictions.\")\n        \n        s_list.append((video_id, 'mouse1', 'self', 'rear', 100, 200))\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\ndef merge_all_predictions(method1_2_predictions, method3_predictions):\n    all_predictions = []\n    \n    if method1_2_predictions['ensemble']:\n        all_predictions.extend(method1_2_predictions['ensemble'])\n    \n    if method1_2_predictions['knn']:\n        all_predictions.extend(method1_2_predictions['knn'])\n    \n    if method3_predictions and len(method3_predictions) > 0:\n        all_predictions.append(method3_predictions)\n    \n    if not all_predictions:\n        return pd.DataFrame()\n    \n    combined = pd.concat(all_predictions, ignore_index=True)\n    \n    grouped = combined.groupby(['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']).size().reset_index(name='count')\n    \n    grouped = grouped.sort_values('count', ascending=False)\n    \n    final_predictions = []\n    seen_segments = set()\n    \n    for _, row in grouped.iterrows():\n        segment_key = (row['video_id'], row['agent_id'], row['target_id'], \n                       tuple(range(row['start_frame'], row['stop_frame'])))\n        \n        overlaps = False\n        for seen_key in seen_segments:\n            if (seen_key[0] == segment_key[0] and \n                seen_key[1] == segment_key[1] and \n                seen_key[2] == segment_key[2]):\n                \n                seen_frames = set(seen_key[3])\n                current_frames = set(segment_key[3])\n                if seen_frames & current_frames:\n                    overlaps = True\n                    break\n        \n        if not overlaps:\n            final_predictions.append(row[['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']])\n            seen_segments.add(segment_key)\n    \n    if final_predictions:\n        return pd.DataFrame(final_predictions)\n    else:\n        return pd.DataFrame()\n\nprint(\"Starting combined processing with all three methods...\")\n\nall_method1_2_predictions = {'ensemble': [], 'knn': []}\nsubmission_list = []\n\nfor section in range(1, len(body_parts_tracked_list)):\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    \n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        print(f\"\\n{section}. Processing videos with {body_parts_tracked}\")\n    \n        train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n        \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        ensemble_model = create_ensemble_model()\n        knn_model = create_knn_model()\n        models_dict = {'ensemble': ensemble_model, 'knn': knn_model}\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            \n            X_enhanced = transform_single_enhanced(single_mouse, body_parts_tracked)\n            X_original = transform_single_original(single_mouse, body_parts_tracked)\n            features_dict = {'enhanced': X_enhanced, 'original': X_original}\n            del single_mouse\n            \n            print(f\"Single mouse features - Enhanced: {X_enhanced.shape}, Original: {X_original.shape}\")\n    \n            if validate_or_submit == 'submit':\n                predictions = submit_method1_and_2(body_parts_tracked_str, 'single', models_dict, features_dict, \n                              single_mouse_label, single_mouse_meta)\n                all_method1_2_predictions['ensemble'].extend(predictions['ensemble'])\n                all_method1_2_predictions['knn'].extend(predictions['knn'])\n            \n            del X_enhanced, X_original\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        \n            X_enhanced = transform_pair_enhanced(mouse_pair, body_parts_tracked)\n            X_original = transform_pair_original(mouse_pair, body_parts_tracked)\n            features_dict = {'enhanced': X_enhanced, 'original': X_original}\n            del mouse_pair\n            \n            print(f\"Mouse pair features - Enhanced: {X_enhanced.shape}, Original: {X_original.shape}\")\n    \n            if validate_or_submit == 'submit':\n                predictions = submit_method1_and_2(body_parts_tracked_str, 'pair', models_dict, features_dict,\n                              mouse_pair_label, mouse_pair_meta)\n                all_method1_2_predictions['ensemble'].extend(predictions['ensemble'])\n                all_method1_2_predictions['knn'].extend(predictions['knn'])\n            \n            del X_enhanced, X_original\n                \n    except Exception as e:\n        print(f'***Exception*** {e}')\n\n    print()\n\nprint(\"\\nProcessing Method 3 (CatBoost)...\")\ncatboost_models = train_catboost_models()\n\ntest_features = prepare_test_dataset_method3(segment_size=50)\n\ntest_preds = []\nfor i in range(5):\n    test_preds.append(catboost_models[i].predict_proba(test_features[catboost_models[i].feature_names_]))\ntest_preds = np.mean(test_preds, 0)\npreds = [catboost_models[0].classes_[x] for x in test_preds.argmax(1)]\nmax_prob = np.array([test_preds[i][x] for i, x in enumerate(test_preds.argmax(1))])\ntest_features['action'] = preds\nmethod3_predictions = test_features[(test_features['action'] != 'no_action') & (max_prob > 0.5)]\nmethod3_predictions['agent_id'] = [\"mouse{}\".format(agent) for agent in method3_predictions['agent_id'].values]\nmethod3_predictions['target_id'] = [\"mouse{}\".format(target) for target in method3_predictions['target_id'].values]\n\nmethod3_submission = method3_predictions[['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']]\n\nprint(\"\\nMerging all predictions...\")\nfinal_submission = merge_all_predictions(all_method1_2_predictions, method3_submission)\n\nprint(\"Finalizing combined submission...\")\n\nif len(final_submission) > 0:\n    submission = final_submission\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')\n\nsubmission_robust.index.name = 'row_id'\nsubmission_robust.to_csv('submission.csv')\n\nprint(f\"Final submission shape: {submission_robust.shape}\")\nprint(\"Submission saved to submission.csv\")\n\nprint(\"\\nFirst few rows of submission:\")\nprint(submission_robust.head(10))\n\nprint(\"\\nSubmission statistics:\")\nprint(f\"Total predictions: {len(submission_robust)}\")\nprint(f\"Unique videos: {submission_robust.video_id.nunique()}\")\nprint(f\"Unique actions: {submission_robust.action.nunique()}\")\nprint(\"\\nAction distribution:\")\nprint(submission_robust.action.value_counts().head(10))","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}