{"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"},{"sourceId":13130542,"sourceType":"datasetVersion","datasetId":8318122},{"sourceId":279806245,"sourceType":"kernelVersion"},{"sourceId":281545678,"sourceType":"kernelVersion"},{"sourceId":281590623,"sourceType":"kernelVersion"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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\n#for 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":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:15.456468Z","iopub.execute_input":"2025-12-10T18:44:15.457215Z","iopub.status.idle":"2025-12-10T18:44:18.017263Z","shell.execute_reply.started":"2025-12-10T18:44:15.457189Z","shell.execute_reply":"2025-12-10T18:44:18.016520Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pathlib import Path\nimport os\nimport sys\n\n# ------------------------------------------------------------\n# Input dataset checks\n# ------------------------------------------------------------\nCOMP_DIR = Path(\"/kaggle/input/MABe-mouse-behavior-detection\")\nSTARTER_DIR = Path(\"/kaggle/input/mabe-starter-train-ja\")\nMABE_PKG_DIR = Path(\"/kaggle/input/mabe-package\")\n\nif not COMP_DIR.exists():\n    raise FileNotFoundError(\n        \"Competition dataset 'MABe Challenge - Social Action Recognition in Mice' \"\n        \"must be attached as an input.\"\n    )\n\nif not STARTER_DIR.exists():\n    raise FileNotFoundError(\n        \"Dataset 'mabe-starter-train-ja' is not attached. \"\n        \"Click 'Add input' and add it before running.\"\n    )\n\nif not MABE_PKG_DIR.exists():\n    raise FileNotFoundError(\n        \"Dataset 'mabe-package' is not attached. \"\n        \"It provides the offline xgboost wheel used by the starter models.\"\n    )\n#defines the dfs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:18.018525Z","iopub.execute_input":"2025-12-10T18:44:18.018888Z","iopub.status.idle":"2025-12-10T18:44:18.026521Z","shell.execute_reply.started":"2025-12-10T18:44:18.018866Z","shell.execute_reply":"2025-12-10T18:44:18.025868Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Install xgboost from offline wheel (no internet)\n# ------------------------------------------------------------\n!pip install -q --no-index --find-links=/kaggle/input/mabe-package xgboost==3.1.1\n\n# ------------------------------------------------------------\n# Copy helper scripts and trained models from starter dataset\n# ------------------------------------------------------------\n!cp /kaggle/input/mabe-starter-train-ja/self_features.py .\n!cp /kaggle/input/mabe-starter-train-ja/pair_features.py .\n!cp /kaggle/input/mabe-starter-train-ja/robustify.py .\n!cp -r /kaggle/input/mabe-starter-train-ja/results .","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:18.027223Z","iopub.execute_input":"2025-12-10T18:44:18.027454Z","iopub.status.idle":"2025-12-10T18:44:49.812905Z","shell.execute_reply.started":"2025-12-10T18:44:18.027431Z","shell.execute_reply":"2025-12-10T18:44:49.811946Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"INPUT_DIR = COMP_DIR\nTRAIN_TRACKING_DIR = INPUT_DIR / \"train_tracking\"\nTRAIN_ANNOTATION_DIR = INPUT_DIR / \"train_annotation\"\nTEST_TRACKING_DIR = INPUT_DIR / \"test_tracking\"\n\nWORKING_DIR = Path(\"/kaggle/working\")\nWORKING_DIR.mkdir(parents=True, exist_ok=True)\n\nSELF_FEATURE_DIR = WORKING_DIR / \"self_features\"\nPAIR_FEATURE_DIR = WORKING_DIR / \"pair_features\"\nSELF_FEATURE_DIR.mkdir(parents=True, exist_ok=True)\nPAIR_FEATURE_DIR.mkdir(parents=True, exist_ok=True)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:49.814191Z","iopub.execute_input":"2025-12-10T18:44:49.814959Z","iopub.status.idle":"2025-12-10T18:44:49.820156Z","shell.execute_reply.started":"2025-12-10T18:44:49.814925Z","shell.execute_reply":"2025-12-10T18:44:49.819345Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\nimport re\nimport ast\nimport itertools\nfrom pathlib import Path\n\nimport numpy as np\n\n# polars is preinstalled on Kaggle GPU/CPU images\ntry:\n    import polars as pl\nexcept ImportError:\n    raise ImportError(\n        \"polars is not available in this environment. \"\n        \"Use a Kaggle GPU/CPU notebook image where polars is preinstalled.\"\n    )\n\nimport xgboost as xgb\nfrom tqdm.auto import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:49.821899Z","iopub.execute_input":"2025-12-10T18:44:49.822079Z","iopub.status.idle":"2025-12-10T18:44:53.175435Z","shell.execute_reply.started":"2025-12-10T18:44:49.822065Z","shell.execute_reply":"2025-12-10T18:44:53.174716Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Helper scripts from starter notebook\n%run -i self_features.py\n%run -i pair_features.py\n%run -i robustify.py","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:53.176126Z","iopub.execute_input":"2025-12-10T18:44:53.176407Z","iopub.status.idle":"2025-12-10T18:44:53.185240Z","shell.execute_reply.started":"2025-12-10T18:44:53.176391Z","shell.execute_reply":"2025-12-10T18:44:53.184713Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"INDEX_COLS = [\n    \"video_id\",\n    \"agent_mouse_id\",\n    \"target_mouse_id\",\n    \"video_frame\",\n]\n\nBODY_PARTS = [\n    \"ear_left\",\n    \"ear_right\",\n    \"nose\",\n    \"neck\",\n    \"body_center\",\n    \"lateral_left\",\n    \"lateral_right\",\n    \"hip_left\",\n    \"hip_right\",\n    \"tail_base\",\n    \"tail_tip\",\n]\n\nSELF_BEHAVIORS = [\n    \"biteobject\",\n    \"climb\",\n    \"dig\",\n    \"exploreobject\",\n    \"freeze\",\n    \"genitalgroom\",\n    \"huddle\",\n    \"rear\",\n    \"rest\",\n    \"run\",\n    \"selfgroom\",\n]\n\nPAIR_BEHAVIORS = [\n    \"allogroom\",\n    \"approach\",\n    \"attack\",\n    \"attemptmount\",\n    \"avoid\",\n    \"chase\",\n    \"chaseattack\",\n    \"defend\",\n    \"disengage\",\n    \"dominance\",\n    \"dominancegroom\",\n    \"dominancemount\",\n    \"ejaculate\",\n    \"escape\",\n    \"flinch\",\n    \"follow\",\n    \"intromit\",\n    \"mount\",\n    \"reciprocalsniff\",\n    \"shepherd\",\n    \"sniff\",\n    \"sniffbody\",\n    \"sniffface\",\n    \"sniffgenital\",\n    \"submit\",\n    \"tussle\",\n]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:53.186627Z","iopub.execute_input":"2025-12-10T18:44:53.187306Z","iopub.status.idle":"2025-12-10T18:44:53.202793Z","shell.execute_reply.started":"2025-12-10T18:44:53.187282Z","shell.execute_reply":"2025-12-10T18:44:53.202119Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def parse_behaviors_column(behaviors_str):\n    if behaviors_str is None:\n        return []\n    return ast.literal_eval(behaviors_str)\n\ndef build_behavior_dataframe(test_df: pl.DataFrame) -> pl.DataFrame:\n    \"\"\"\n    Expand behaviors_labeled into one row per (lab, video, agent, target, behavior).\n    \"\"\"\n    behavior_df = (\n        test_df\n        .filter(pl.col(\"behaviors_labeled\").is_not_null())\n        .select([\"lab_id\", \"video_id\", \"behaviors_labeled\"])\n        .with_columns(\n            pl.col(\"behaviors_labeled\")\n            .map_elements(\n                parse_behaviors_column,\n                return_dtype=pl.List(pl.Utf8),\n            )\n            .alias(\"behaviors_labeled_list\")\n        )\n        .explode(\"behaviors_labeled_list\")\n        .rename({\"behaviors_labeled_list\": \"behaviors_labeled_element\"})\n        .with_columns(\n            pl.col(\"behaviors_labeled_element\").str.split(\",\").list.get(0)\n            .str.replace_all(\"[()' ]\", \"\")\n            .alias(\"agent\"),\n            pl.col(\"behaviors_labeled_element\").str.split(\",\").list.get(1)\n            .str.replace_all(\"[()' ]\", \"\")\n            .alias(\"target\"),\n            pl.col(\"behaviors_labeled_element\").str.split(\",\").list.get(2)\n            .str.replace_all(\"[()' ]\", \"\")\n            .alias(\"behavior\"),\n        )\n        .select([\"lab_id\", \"video_id\", \"agent\", \"target\", \"behavior\"])\n    )\n    return behavior_df\n\ndef extract_mouse_id(mouse_str: str) -> int:\n    \"\"\"\n    Convert 'mouse1' -> 1, 'mouse2' -> 2, 'self' -> -1.\n    \"\"\"\n    if mouse_str == \"self\":\n        return -1\n    m = re.search(r\"mouse(\\d+)\", mouse_str)\n    if m:\n        return int(m.group(1))\n    raise ValueError(f\"Unexpected mouse id format: {mouse_str}\")\n\ndef load_features_for_group(lab_id, video_id, agent, target):\n    \"\"\"\n    Load per frame features for a given (lab, video, agent, target) group.\n    Returns:\n      index_df   - DataFrame with INDEX_COLS\n      feature_df - DataFrame with feature columns only\n    \"\"\"\n    agent_mouse_id = extract_mouse_id(agent)\n    target_mouse_id = extract_mouse_id(target)\n\n    if target == \"self\":\n        feature_path = SELF_FEATURE_DIR / f\"{video_id}.parquet\"\n        scan = pl.scan_parquet(feature_path).filter(\n            pl.col(\"agent_mouse_id\") == agent_mouse_id\n        )\n    else:\n        feature_path = PAIR_FEATURE_DIR / f\"{video_id}.parquet\"\n        scan = pl.scan_parquet(feature_path).filter(\n            (pl.col(\"agent_mouse_id\") == agent_mouse_id)\n            & (pl.col(\"target_mouse_id\") == target_mouse_id)\n        )\n\n    full_df = scan.collect()\n    if full_df.height == 0:\n        return full_df, full_df\n\n    index_df = full_df.select(INDEX_COLS)\n    feature_df = full_df.select(pl.exclude(INDEX_COLS))\n    return index_df, feature_df\n\ndef load_models_for_behavior(lab_id: str, behavior: str):\n    \"\"\"\n    Load all fold models and thresholds for a given (lab, behavior).\n    Returns list of (model, threshold).\n    \"\"\"\n    behavior_dir = WORKING_DIR / \"results\" / lab_id / behavior\n    fold_dirs = sorted(behavior_dir.glob(\"fold_*\"))\n    models = []\n    for fold_dir in fold_dirs:\n        model_file = fold_dir / \"model.json\"\n        thr_file = fold_dir / \"threshold.txt\"\n        if not model_file.exists() or not thr_file.exists():\n            continue\n        with open(thr_file, \"r\") as f:\n            threshold = float(f.read().strip())\n        model = xgb.Booster(model_file=str(model_file))\n        models.append((model, threshold))\n    return models\n\ndef predict_for_group(\n    lab_id: str,\n    video_id: int,\n    agent: str,\n    target: str,\n    group_behaviors: pl.DataFrame,\n    ):\n    \"\"\"\n    Run inference for one group of (lab_id, video_id, agent, target).\n\n    Improvements:\n      - Aggregate folds per behavior into a single score column\n        (mean of thresholded probabilities).\n      - Pick best behavior per frame using those aggregated scores.\n    \"\"\"\n    index_df, feature_df = load_features_for_group(lab_id, video_id, agent, target)\n\n    if feature_df.height == 0:\n        return None\n\n    # Create XGBoost DMatrix once per group and reuse across behaviors\n    dtest = xgb.DMatrix(feature_df.to_pandas(), feature_names=feature_df.columns)\n\n    prediction_df = index_df.clone()\n    used_cols = []\n    unique_behaviors = (group_behaviors.select(\"behavior\").unique()[\"behavior\"].to_list()\n    )\n\n    for behavior in unique_behaviors:\n        models = load_models_for_behavior(lab_id, behavior)\n        if not models:\n            # No trained model for this (lab, behavior) in the starter models\n            continue\n\n        # Aggregate over folds: mean of thresholded probabilities\n        agg_scores = np.zeros(feature_df.height, dtype=np.float32)\n\n        for model, threshold in models:\n            probs = model.predict(dtest)\n            labels = (probs >= threshold).astype(np.int8)\n            agg_scores += probs * labels\n\n        agg_scores /= max(len(models), 1)\n\n        col_name = behavior\n        prediction_df = prediction_df.with_columns(\n            pl.Series(name=col_name, values=agg_scores)\n        )\n        used_cols.append(col_name)\n\n    if not used_cols:\n        return None\n\n    # Pick best behavior per frame (over behaviors only)\n    cols = used_cols\n\n    prediction_labels_df = (\n        prediction_df\n        .with_columns(\n            pl.struct(pl.col(cols))\n            .map_elements(\n                lambda row: (\n                    \"none\"\n                    if sum(row.values()) == 0\n                    else cols[int(np.argmax(list(row.values())))]\n                ),\n                return_dtype=pl.String,\n            )\n            .alias(\"prediction\")\n        )\n        .select(INDEX_COLS + [\"prediction\"])\n    )\n\n    # Convert per frame labels into time segments\n    agent_mouse_id = extract_mouse_id(agent)\n    target_mouse_id = extract_mouse_id(target) \n    group_submission = (\n        prediction_labels_df\n        .filter(pl.col(\"prediction\") != pl.col(\"prediction\").shift(1))\n        .with_columns(\n            pl.col(\"video_frame\").shift(-1).alias(\"stop_frame\")\n        )\n        .filter(pl.col(\"prediction\") != \"none\")\n        .select(\n            pl.col(\"video_id\"),\n            (pl.lit(\"mouse\") + pl.lit(agent_mouse_id).cast(pl.Utf8)).alias(\"agent_id\"),\n            pl.when(pl.lit(target_mouse_id) == -1)\n            .then(pl.lit(\"self\"))\n            .otherwise(pl.lit(\"mouse\") + pl.lit(target_mouse_id).cast(pl.Utf8))\n            .alias(\"target_id\"),\n            pl.col(\"prediction\").alias(\"action\"),\n            pl.col(\"video_frame\").alias(\"start_frame\"),\n            pl.col(\"stop_frame\"),\n        )\n    )\n\n    return group_submission\n\n\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:53.203652Z","iopub.execute_input":"2025-12-10T18:44:53.203996Z","iopub.status.idle":"2025-12-10T18:44:53.233077Z","shell.execute_reply.started":"2025-12-10T18:44:53.203972Z","shell.execute_reply":"2025-12-10T18:44:53.232579Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print('Load test metadata:')\ntest_df = pl.read_csv(COMP_DIR/'test.csv')\nprint(test_df.head())\n\n#behavior_df, list of behaviors? why not label encoder?\n\nbehavior_df = build_behavior_dataframe(test_df)\n\nprint('behavior_df')\n\nprint(behavior_df.head())\n\ngroups = list(\n    behavior_df.group_by(\"lab_id\", \"video_id\", \"agent\", \"target\", maintain_order=True)\n)\nprint(f\"Number of (lab, video, agent, target) groups: {len(groups)}\")\n\nprint(\"Generating self and pair features for all test videos...\")\n\nrows = test_df.rows(named=True)\n\n\nfor row in tqdm(rows, total=len(rows)): #make self/pair features and it seems like the agent and target is predefined in the make_self_features from the tracking and the metadata...\n    lab_id = test_df['lab_id'].item()\n    video_id = test_df['video_id'].item()\n\n    tracking_path = TEST_TRACKING_DIR / f\"{lab_id}/{video_id}.parquet\"\n    print(tracking_path)\n    tracking = pl.read_parquet(tracking_path)\n\n    self_feat = make_self_features(metadata=row, tracking=tracking)\n\n    pair_feat = make_pair_features(metadata=row, tracking=tracking)\n    \n    print(self_feat)\n    print(pair_feat)\n\n    self_feat.write_parquet(SELF_FEATURE_DIR / f\"{video_id}.parquet\")\n    pair_feat.write_parquet(PAIR_FEATURE_DIR / f\"{video_id}.parquet\")\n\n    del self_feat, pair_feat, tracking\n    gc.collect()\n\nprint('Running inference and creating group submissions:')\n\ngroup_submissions = []\n\nfor (lab_id, video_id, agent, target), group in tqdm(groups, total=len(groups)):\n    group_submission = predict_for_group(\n        lab_id=lab_id,\n        video_id=video_id,\n        agent=agent,\n        target=target,\n        group_behaviors=group,\n    )\n\n    if group_submission is not None and group_submission.height > 0:\n        group_submissions.append(group_submission)\n\nif not group_submissions:\n    raise RuntimeError(\n        \"No submissions were generated. \"\n        \"Check that starter models exist under /kaggle/working/results.\"\n    )\n\nsubmission = pl.concat(group_submissions, how='vertical').sort(\n        \"video_id\",\n        \"agent_id\",\n        \"target_id\",\n        \"action\",\n        \"start_frame\",\n        \"stop_frame\",\n    )\n\nprint(\"Initial submission rows:\", submission.height)\n\n# ============================================================\n# 4. Robustify and final clean up\n# ============================================================\nprint(\"Running robustify on submission...\")\nsubmission = robustify(submission, test_df, train_test=\"test\")\n\n# Keep only valid intervals\nsubmission = submission.filter(pl.col(\"start_frame\") < pl.col(\"stop_frame\"))\n\n# Drop ultra short segments (likely noise)\nsubmission = submission.with_columns(\n    (pl.col(\"stop_frame\") - pl.col(\"start_frame\")).alias(\"duration\")\n).filter(pl.col(\"duration\") >= 2).drop(\"duration\")\n\nprint(\"Rows after robustify, validity check and duration filter:\", submission.height)\n\nsubmission = submission.to_pandas()\n\nsubmission = submission.reset_index()\n\nsubmission.index.name = 'row_id'\n\nprint(submission.head())\n\nsubmission.to_csv('submission.csv')\n\nprint('csv saved to submission.csv')\n\n# Add row_id and save as submission.csv\n#final_submission = submission.with_row_index(\"row_id\")\n#print(final_submission.head())\n#output_dir = 'kaggle/working'\n#Path(output_dir).mkdir(parents=True, exist_ok=True)\n# The parents=True argument creates necessary parent directories ('kaggle', then 'working')\n# The exist_ok=True prevents an error if the directory already exists (good practice for re-runs)\n\n# --- Now write the CSV file ---\n#output_path = os.path.join(output_dir, 'submission.csv')\n#final_submission.write_csv(output_path, encoding=\"utf8\")\n\n#print(f\"CSV file has been created successfully at {output_path}\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T18:44:53.233777Z","iopub.execute_input":"2025-12-10T18:44:53.234016Z","iopub.status.idle":"2025-12-10T18:45:22.789668Z","shell.execute_reply.started":"2025-12-10T18:44:53.233995Z","shell.execute_reply":"2025-12-10T18:45:22.788944Z"}},"outputs":[],"execution_count":null}]}