{"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":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13719397,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%writefile baseline.py\n#!/usr/bin/env python3\nimport argparse\nimport json\nimport os\nimport warnings\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Dict, List, Optional, Set, Tuple\nfrom collections import defaultdict\n\nimport polars as pl\nfrom tqdm.auto import tqdm\n\n# ========================\n# Config\n# ========================\n\n@dataclass(frozen=True)\nclass Config:\n    data_root: Path = Path(os.getenv(\"MABE_DATA_ROOT\", \"/kaggle/input/MABe-mouse-behavior-detection\"))\n    submission_file: str = os.getenv(\"MABE_SUBMISSION\", \"submission.csv\")\n    row_id_col: str = os.getenv(\"MABE_ROW_ID_COL\", \"row_id\")\n\n    @property\n    def train_csv(self) -> Path: return self.data_root / \"train.csv\"\n    @property\n    def test_csv(self) -> Path: return self.data_root / \"test.csv\"\n    @property\n    def train_annot_dir(self) -> Path: return self.data_root / \"train_annotation\"\n    @property\n    def train_track_dir(self) -> Path: return self.data_root / \"train_tracking\"\n    @property\n    def test_track_dir(self) -> Path: return self.data_root / \"test_tracking\"\n\n    @property\n    def submission_schema(self) -> Dict[str, pl.DataType]:\n        return {\n            \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n            \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n        }\n\n    @property\n    def solution_schema(self) -> Dict[str, pl.DataType]:\n        return {\n            \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n            \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n            \"lab_id\": pl.Utf8, \"behaviors_labeled\": pl.Utf8,\n        }\n\nclass HostVisibleError(Exception): pass\n\n# ========================\n# Utils\n# ========================\n\ndef safe_json_loads(s: Optional[str]) -> List[str]:\n    if s is None: return []\n    if isinstance(s, list): return [str(x) for x in s]\n    if not isinstance(s, str): return []\n    s = s.strip()\n    if not s: return []\n    try:\n        return json.loads(s)\n    except Exception:\n        try: return json.loads(s.replace(\"'\", '\"'))\n        except Exception: return []\n\ndef validate_schema(df: pl.DataFrame, schema: Dict[str, pl.DataType], name: str) -> pl.DataFrame:\n    casts = []\n    for col, dtype in schema.items():\n        if col in df.columns and df[col].dtype != dtype:\n            casts.append(pl.col(col).cast(dtype))\n    return df.with_columns(casts) if casts else df\n\ndef validate_frame_ranges(df: pl.DataFrame, name: str) -> None:\n    if not (df[\"start_frame\"] <= df[\"stop_frame\"]).all():\n        raise ValueError(f\"{name}: start_frame > stop_frame detected\")\n\ndef _norm_mouse_id(x: str | int) -> str:\n    s = str(x)\n    return s if s.startswith(\"mouse\") else f\"mouse{s}\"\n\ndef _norm_triplet(agent: str | int, target: str | int, action: str) -> str:\n    return f\"{_norm_mouse_id(agent)},{_norm_mouse_id(target)},{action}\"\n\ndef _range_frames(start: int, stop: int):\n    return range(start, stop)\n\n# ========================\n# Metrics\n# ========================\n\ndef single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1.0) -> float:\n    label_frames: Dict[str, Set[int]] = defaultdict(set)\n    for row in lab_solution.to_dicts():\n        label_frames[row[\"label_key\"]].update(_range_frames(row[\"start_frame\"], row[\"stop_frame\"]))\n\n    active_by_video: Dict[int, Set[str]] = {}\n    for row in lab_solution.select([\"video_id\", \"behaviors_labeled\"]).unique().to_dicts():\n        s: Set[str] = set()\n        for item in safe_json_loads(row[\"behaviors_labeled\"]):\n            parts = [p.strip() for p in str(item).replace(\"'\", \"\").split(\",\")]\n            if len(parts) == 3:\n                a, t, act = parts\n                s.add(_norm_triplet(a, t, act))\n        active_by_video[int(row[\"video_id\"])] = s\n\n    prediction_frames: Dict[str, Set[int]] = defaultdict(set)\n    for video_id in lab_solution[\"video_id\"].unique():\n        active = active_by_video.get(int(video_id), set())\n        predicted_mouse_pairs: Dict[str, Set[int]] = defaultdict(set)\n        for row in lab_submission.filter(pl.col(\"video_id\") == video_id).to_dicts():\n            triple_norm = _norm_triplet(row[\"agent_id\"], row[\"target_id\"], row[\"action\"])\n            if triple_norm not in active: continue\n            pred_key = row[\"prediction_key\"]\n            agent_target = f\"{row['agent_id']},{row['target_id']}\"\n            new_frames = set(_range_frames(row[\"start_frame\"], row[\"stop_frame\"]))\n            new_frames -= prediction_frames[pred_key]\n            if predicted_mouse_pairs[agent_target] & new_frames:\n                raise HostVisibleError(\"Multiple predictions for the same frame from one agent/target pair\")\n            prediction_frames[pred_key].update(new_frames)\n            predicted_mouse_pairs[agent_target].update(new_frames)\n\n    tps: Dict[str, int] = defaultdict(int)\n    fns: Dict[str, int] = defaultdict(int)\n    fps: Dict[str, int] = defaultdict(int)\n    distinct_actions: Set[str] = set()\n\n    for key, pred_frames in prediction_frames.items():\n        action = key.split(\"_\")[-1]\n        distinct_actions.add(action)\n        gt_frames = label_frames.get(key, set())\n        tps[action] += len(pred_frames & gt_frames)\n        fns[action] += len(gt_frames - pred_frames)\n        fps[action] += len(pred_frames - gt_frames)\n\n    for key, gt_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(gt_frames)\n\n    if not distinct_actions: return 0.0\n\n    beta2 = beta * beta\n    f_scores: List[float] = []\n    for action in distinct_actions:\n        tp, fn, fp = tps[action], fns[action], fps[action]\n        denom = (1 + beta2) * tp + beta2 * fn + fp\n        f_scores.append(0.0 if denom == 0 else (1 + beta2) * tp / denom)\n    return sum(f_scores) / len(f_scores)\n\ndef mouse_fbeta(solution: pl.DataFrame, submission: pl.DataFrame, beta: float = 1.0, cfg: Optional[Config] = None) -> float:\n    cfg = cfg or Config()\n    solution = validate_schema(solution, cfg.solution_schema, \"Solution\")\n    submission = validate_schema(submission, cfg.submission_schema, \"Submission\")\n    validate_frame_ranges(solution, \"Solution\")\n    validate_frame_ranges(submission, \"Submission\")\n\n    def add_key(df: pl.DataFrame, col_name: str) -> pl.DataFrame:\n        return df.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(col_name)\n        )\n\n    solution = add_key(solution, \"label_key\")\n    submission = add_key(submission, \"prediction_key\")\n\n    lab_scores: List[float] = []\n    for lab_id in solution[\"lab_id\"].unique():\n        lab_solution = solution.filter(pl.col(\"lab_id\") == lab_id)\n        lab_videos = lab_solution[\"video_id\"].unique()\n        lab_submission = submission.filter(pl.col(\"video_id\").is_in(lab_videos))\n        lab_scores.append(single_lab_f1(lab_solution, lab_submission, beta=beta))\n\n    return sum(lab_scores) / len(lab_scores) if lab_scores else 0.0\n\ndef score(solution: pl.DataFrame, submission: pl.DataFrame, row_id_column_name: str = \"\", beta: float = 1.0, cfg: Optional[Config] = None) -> float:\n    if row_id_column_name:\n        solution = solution.drop(row_id_column_name, strict=False)\n        submission = submission.drop(row_id_column_name, strict=False)\n    return mouse_fbeta(solution, submission, beta=beta, cfg=cfg)\n\n# ========================\n# Build solution\n# ========================\n\ndef create_solution_df(dataset: pl.DataFrame, cfg: Optional[Config] = None) -> pl.DataFrame:\n    cfg = cfg or Config()\n    records: List[pl.DataFrame] = []\n    for row in tqdm(dataset.to_dicts(), total=len(dataset), desc=\"Building solution\"):\n        lab_id: str = row[\"lab_id\"]\n        if lab_id.startswith(\"MABe22\"): continue\n        video_id: int = row[\"video_id\"]\n        annot_path = cfg.train_annot_dir / lab_id / f\"{video_id}.parquet\"\n        if not annot_path.exists():\n            continue\n        try:\n            annot = pl.read_parquet(annot_path).with_columns(\n                [\n                    pl.lit(lab_id).alias(\"lab_id\"),\n                    pl.lit(video_id).alias(\"video_id\"),\n                    pl.lit(row[\"behaviors_labeled\"]).alias(\"behaviors_labeled\"),\n                    pl.concat_str([pl.lit(\"mouse\"), pl.col(\"agent_id\").cast(pl.Utf8)]).alias(\"agent_id\"),\n                    pl.concat_str([pl.lit(\"mouse\"), pl.col(\"target_id\").cast(pl.Utf8)]).alias(\"target_id\"),\n                ]\n            )\n            # cast numeric cols to Int64\n            annot = annot.with_columns([\n                pl.col(\"video_id\").cast(pl.Int64),\n                pl.col(\"start_frame\").cast(pl.Int64),\n                pl.col(\"stop_frame\").cast(pl.Int64),\n            ])\n            annot = annot.select([c for c in cfg.solution_schema.keys() if c in annot.columns])\n            records.append(annot)\n        except Exception as e:\n            print(f\"[WARN] Failed to load {annot_path}: {e}\")\n            continue\n    if not records: raise ValueError(\"No annotation files loaded.\")\n    solution = pl.concat(records, how=\"vertical\")\n    solution = validate_schema(solution, cfg.solution_schema, \"Solution\")\n    return solution\n\n# ========================\n# Main\n# ========================\n\ndef run_submit(cfg: Config, eps: float, prior_scope: str) -> None:\n    print(\"[INFO] Loading train (priors) and test...\")\n    train = pl.read_csv(cfg.train_csv)\n    train_subset = train.filter(~pl.col(\"lab_id\").str.starts_with(\"MABe22\"))\n    solution = create_solution_df(train_subset, cfg)\n\n    test = pl.read_csv(cfg.test_csv)\n    # simple dummy predictions (all background)\n    submission = test.select([\"video_id\"]).with_columns([\n        pl.lit(\"mouse1\").alias(\"agent_id\"),\n        pl.lit(\"mouse2\").alias(\"target_id\"),\n        pl.lit(\"background\").alias(\"action\"),\n        pl.lit(0).alias(\"start_frame\"),\n        pl.lit(10).alias(\"stop_frame\"),\n    ])\n    submission = submission.select(list(cfg.submission_schema.keys())).with_row_index(cfg.row_id_col)\n    submission.write_csv(cfg.submission_file)\n    print(f\"[INFO] Saved submission to {cfg.submission_file}\")\n    print(submission.head())\n\ndef main() -> None:\n    parser = argparse.ArgumentParser(description=\"MABe Mouse Baseline (fixed)\")\n    parser.add_argument(\"--mode\", choices=[\"validate\", \"submit\", \"all\"], default=\"submit\")\n    parser.add_argument(\"--prior-scope\", choices=[\"lab\", \"global\", \"mixed\"], default=\"mixed\")\n    parser.add_argument(\"--eps\", type=float, default=1.0)\n    args = parser.parse_args()\n\n    cfg = Config()\n\n    if args.mode in (\"submit\", \"all\"):\n        run_submit(cfg, args.eps, args.prior_scope)\n\nif __name__ == \"__main__\":\n    warnings.filterwarnings(\"ignore\")\n    main()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-09-21T02:59:43.590748Z","iopub.execute_input":"2025-09-21T02:59:43.591085Z","iopub.status.idle":"2025-09-21T02:59:43.603281Z","shell.execute_reply.started":"2025-09-21T02:59:43.591055Z","shell.execute_reply":"2025-09-21T02:59:43.602421Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python baseline.py --mode submit --prior-scope mixed --eps 1.0\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-21T02:59:43.604031Z","iopub.execute_input":"2025-09-21T02:59:43.604285Z","iopub.status.idle":"2025-09-21T02:59:45.901194Z","shell.execute_reply.started":"2025-09-21T02:59:43.604258Z","shell.execute_reply":"2025-09-21T02:59:45.900115Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}