{"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":13719397,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport gc\nimport numpy as np\nimport pandas as pd\nfrom glob import glob\nfrom tqdm.auto import tqdm\nimport lightgbm as lgb\nfrom joblib import dump, load\nimport warnings\n\nwarnings.filterwarnings('ignore')\n\n# Global constants for feature engineering\nBODY_PARTS = ['nose', 'ear_left', 'ear_right', 'neck', 'hip_left', 'hip_right', 'tail_base']","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CFG:\n    # --- Data Paths ---\n    DATA_ROOT = \"/kaggle/input/MABe-mouse-behavior-detection\"\n    TRAIN_TRACKING_DIR = os.path.join(DATA_ROOT, \"train_tracking\")\n    TEST_TRACKING_DIR = os.path.join(DATA_ROOT, \"test_tracking\")\n    TRAIN_ANNOTATION_DIR = os.path.join(DATA_ROOT, \"train_annotation\")\n    \n    # --- Behaviors to Model ---\n    BEHAVIORS = ['attack', 'chase', 'investigation', 'mount', 'sniff', 'huddle', 'submit']\n    \n    # --- Post-Processing Parameters ---\n    PROB_THRESHOLD = 0.5\n    MIN_EVENT_LENGTH_FRAMES = 5\n    \n    # --- Training Parameters ---\n    VIDEOS_PER_BEHAVIOR = 40 \n    \n    LGB_PARAMS = {\n        'objective': 'binary', 'metric': 'binary_logloss', 'boosting_type': 'gbdt',\n        'n_estimators': 500, 'learning_rate': 0.05, 'num_leaves': 31,\n        'max_depth': -1, 'seed': 42, 'n_jobs': -1, 'verbose': -1,\n    }","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def create_features(df_tracking):\n    \"\"\"\n    Transforms raw tracking data into a standardized egocentric feature set.\n    Feature names are based on 'agent' and 'target' roles, not raw mouse IDs.\n    \"\"\"\n    all_pair_features = []\n    mice_present = sorted(df_tracking['mouse_id'].unique())\n    if len(mice_present) < 2: return None\n\n    try:\n        df_wide = df_tracking.pivot_table(\n            index='video_frame', columns=['mouse_id', 'bodypart'], values=['x', 'y']\n        )\n        df_wide.columns = [f'{val}_{mid}_{bp}' for val, mid, bp in df_wide.columns]\n        df_wide = df_wide.sort_index().ffill().bfill()\n    except Exception:\n        return None\n\n    for agent_id in mice_present:\n        for target_id in mice_present:\n            if agent_id == target_id: continue\n            \n            agent_cols = [f'x_{agent_id}_neck', f'y_{agent_id}_neck', f'x_{agent_id}_nose', f'y_{agent_id}_nose']\n            if not all(c in df_wide.columns for c in agent_cols):\n                continue\n\n            agent_neck_x = df_wide[f'x_{agent_id}_neck'].values\n            agent_neck_y = df_wide[f'y_{agent_id}_neck'].values\n            agent_nose_x = df_wide[f'x_{agent_id}_nose'].values - agent_neck_x\n            agent_nose_y = df_wide[f'y_{agent_id}_nose'].values - agent_neck_y\n            \n            angles = np.arctan2(agent_nose_y, agent_nose_x) - np.pi / 2\n            cos_angles, sin_angles = np.cos(angles), np.sin(angles)\n            \n            pair_features = pd.DataFrame({'video_frame': df_wide.index})\n            \n            # Use STANDARDIZED names: 'agent' and 'target'\n            for role, mouse_id in [('agent', agent_id), ('target', target_id)]:\n                for bp in BODY_PARTS:\n                    x_col, y_col = f'x_{mouse_id}_{bp}', f'y_{mouse_id}_{bp}'\n                    if x_col in df_wide.columns and y_col in df_wide.columns:\n                        orig_x = df_wide[x_col].values - agent_neck_x\n                        orig_y = df_wide[y_col].values - agent_neck_y\n                        rotated_x = orig_x * cos_angles + orig_y * sin_angles\n                        rotated_y = -orig_x * sin_angles + orig_y * cos_angles\n                    else:\n                        rotated_x, rotated_y = np.zeros_like(agent_neck_x), np.zeros_like(agent_neck_y)\n                    \n                    pair_features[f'{role}_{bp}_x_ego'] = rotated_x\n                    pair_features[f'{role}_{bp}_y_ego'] = rotated_y\n                    pair_features[f'{role}_{bp}_vx_ego'] = np.diff(rotated_x, prepend=0)\n                    pair_features[f'{role}_{bp}_vy_ego'] = np.diff(rotated_y, prepend=0)\n            \n            pair_features['agent_id'] = agent_id\n            pair_features['target_id'] = target_id\n            all_pair_features.append(pair_features)\n\n    return pd.concat(all_pair_features) if all_pair_features else None","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"IS_SUBMISSION = len(glob(os.path.join(CFG.TEST_TRACKING_DIR, \"*/*.parquet\"))) > 100\n\nif not IS_SUBMISSION:\n    print(\"Detected Training Environment. Starting model training...\")\n    \n    annotation_files = glob(os.path.join(CFG.TRAIN_ANNOTATION_DIR, \"*/*.parquet\"))\n    all_labels_list = []\n    for f in tqdm(annotation_files, desc=\"Reading annotations\"):\n        video_id = int(os.path.basename(f).split('.')[0])\n        df = pd.read_parquet(f)\n        df['video_id'] = video_id\n        all_labels_list.append(df)\n    all_labels_df = pd.concat(all_labels_list, ignore_index=True)\n\n    for behavior in CFG.BEHAVIORS:\n        print(f\"--- Training model for: {behavior} ---\")\n        behavior_labels = all_labels_df[all_labels_df['action'] == behavior]\n        if behavior_labels.empty:\n            print(f\"No annotations for {behavior}. Skipping.\")\n            continue\n\n        X_train_all, y_train_all = [], []\n        train_video_ids = behavior_labels['video_id'].unique()\n        if len(train_video_ids) > CFG.VIDEOS_PER_BEHAVIOR:\n            train_video_ids = np.random.choice(train_video_ids, CFG.VIDEOS_PER_BEHAVIOR, replace=False)\n        \n        for video_id in tqdm(train_video_ids, desc=f\"Featurizing for {behavior}\"):\n            tracking_path = glob(os.path.join(CFG.TRAIN_TRACKING_DIR, f\"*/{video_id}.parquet\"))\n            if not tracking_path: continue\n            \n            df_tracking = pd.read_parquet(tracking_path[0])\n            df_features = create_features(df_tracking)\n            if df_features is None: continue\n            \n            n_frames = df_tracking['video_frame'].max() + 1\n            vid_labels = behavior_labels[behavior_labels['video_id'] == video_id]\n            \n            for (agent_id, target_id), pair_df in df_features.groupby(['agent_id', 'target_id']):\n                target = np.zeros(n_frames)\n                pair_annotations = vid_labels[(vid_labels['agent_id'] == agent_id) & (vid_labels['target_id'] == target_id)]\n                for _, row in pair_annotations.iterrows():\n                    target[row['start_frame']:row['stop_frame']] = 1\n                \n                pair_df['target'] = pair_df['video_frame'].map(lambda f: target[f] if f < n_frames else 0)\n                \n                positive_samples = pair_df[pair_df['target'] == 1]\n                negative_samples = pair_df[pair_df['target'] == 0]\n                neg_sample_size = min(len(positive_samples) * 2, len(negative_samples))\n                if neg_sample_size > 0:\n                    negative_samples = negative_samples.sample(n=neg_sample_size, random_state=42)\n                \n                train_sample = pd.concat([positive_samples, negative_samples])\n                feature_cols = [c for c in train_sample.columns if c.startswith(('agent_', 'target_'))]\n                \n                X_train_all.append(train_sample[feature_cols])\n                y_train_all.append(train_sample['target'])\n            \n        if X_train_all:\n            X_train, y_train = pd.concat(X_train_all), pd.concat(y_train_all)\n            model = lgb.LGBMClassifier(**CFG.LGB_PARAMS)\n            model.fit(X_train, y_train)\n            dump(model, f'{behavior}_model.joblib')\n            print(f\"Model for {behavior} trained and saved.\")\nelse:\n    print(\"Detected Submission Environment. Skipping training.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Starting Inference...\")\n\nmodels = {}\nfor behavior in CFG.BEHAVIORS:\n    MODEL_DIR = \".\"\n    if IS_SUBMISSION:\n        # Example path. Make sure your training notebook output folder is named 'mabe-final-train'\n        # or update this path accordingly.\n        MODEL_DIR = \"/kaggle/input/mabe-final-train\" \n        \n    model_path = os.path.join(MODEL_DIR, f'{behavior}_model.joblib')\n    if os.path.exists(model_path):\n        models[behavior] = load(model_path)\n    else:\n        models[behavior] = None \n        print(f\"Warning: Model for {behavior} not found at {model_path}.\")\n\nall_events = []\ntest_files = glob(os.path.join(CFG.TEST_TRACKING_DIR, \"*/*.parquet\"))\n\nfor f in tqdm(test_files, desc=\"Processing test files\"):\n    video_id = int(os.path.basename(f).split('.')[0])\n    \n    df_tracking = pd.read_parquet(f)\n    if 'mouse_id' not in df_tracking.columns: continue\n    df_tracking['mouse_id'] = df_tracking['mouse_id'].astype(str).str.replace('mouse', '').astype(int)\n    \n    df_features = create_features(df_tracking)\n    if df_features is None: continue\n        \n    feature_cols = [c for c in df_features.columns if c.startswith(('agent_', 'target_'))]\n    \n    for (agent_id, target_id), pair_df in df_features.groupby(['agent_id', 'target_id']):\n        if pair_df.empty: continue\n        \n        X_test = pair_df[feature_cols]\n\n        for behavior in CFG.BEHAVIORS:\n            if models.get(behavior) is None: continue\n            model = models[behavior]\n            \n            preds = model.predict_proba(X_test)[:, 1]\n            binary_preds = (preds > CFG.PROB_THRESHOLD).astype(int)\n            \n            if np.sum(binary_preds) > 0:\n                diffs = np.diff(binary_preds, prepend=0, append=0)\n                starts = pair_df['video_frame'].values[np.where(diffs == 1)[0]]\n                stops = pair_df['video_frame'].values[np.where(diffs == -1)[0]]\n                \n                for start, stop in zip(starts, stops):\n                    if stop - start >= CFG.MIN_EVENT_LENGTH_FRAMES:\n                        all_events.append({\n                            'video_id': video_id, 'agent_id': f'mouse{agent_id}', 'target_id': f'mouse{target_id}',\n                            'action': behavior, 'start_frame': start, 'stop_frame': stop\n                        })\n\n# --- Create Submission File ---\nsubmission_df = pd.DataFrame(all_events)\nif not submission_df.empty:\n    submission_df['row_id'] = submission_df.index\n    submission_df = submission_df[['row_id', 'video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']]\nelse:\n    sample_sub = pd.read_csv(os.path.join(CFG.DATA_ROOT, \"sample_submission.csv\"))\n    submission_df = pd.DataFrame(columns=sample_sub.columns)\n\nsubmission_df.to_csv('submission.csv', index=False)\nprint(\"Submission file created successfully!\")\nprint(f\"Generated {len(submission_df)} events.\")\nprint(submission_df.head())","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}