{"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":13234212,"sourceType":"datasetVersion","datasetId":8387028}],"dockerImageVersionId":31090,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## 📘 MABe: 2.5D Pose-based Social Action Detection — Inference Notebook\n\n**Goal.** Generate `submission.csv` for the **MABe Challenge — Social Action Recognition in Mice** using a trained 2.5D EfficientNet model that consumes top-down 2D keypoints (pose) and predicts frame-wise probabilities for action classes, then converts them into action bouts.\n\n**What this notebook does**\n\n1. **Preprocess pose**\n\n   * Drop head-mounted parts (e.g., `headpiece_*`).\n   * Map body parts to grayscale intensities (mouse-ID × body-part position) and rasterize per frame.\n   * Build **2.5D clips** by stacking frames at fixed temporal **offsets** around a center frame.\n2. **Model inference**\n\n   * Uses a trained **EfficientNet (timm)** with custom `in_chans = len(OFFSETS)`; loads weights and a `class_to_idx` map.\n   * Iterates test frames (optionally **subsampled** by `INFER_FRAME_STEP`), batches clips, and produces per-class probabilities.\n3. **Post-process**\n\n   * Threshold probabilities → binary sequences → merge consecutive frames into bouts.\n   * Enforce **whitelists** from `behaviors_labeled` (only allowed agent/target/action triplets per video).\n   * Drop invalid bouts and resolve overlaps consistently.\n\n> 🧭 Why this approach?  \n> Many baselines convert pose into **tabular/hand-crafted features** for sequence models. Here we take the opposite route: **turn pose into images**, then apply a **2D CNN** over a **2.5D clip** (stack of offset frames). This keeps preprocessing simple, leverages CNN priors.\n\n## Related tools & notebooks\n\n- **Visualizing Mouse Pose Track** — quick look at keypoint tracks and how they rasterize  \n  https://www.kaggle.com/code/kmatsu01/visualizing-mouse-pose-track","metadata":{}},{"cell_type":"code","source":"\n# ===================================================================\n# 0. Libraries\n# ===================================================================\nimport os, gc, re, json, warnings\nfrom pathlib import Path\nimport pandas as pd\nimport numpy as np\nfrom PIL import Image\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms as T\nimport timm\n\nwarnings.filterwarnings('ignore')\n\n# ===================================================================\n# 1. Configuration\n# ===================================================================\nclass CFG:\n    # Base dataset path (Kaggle competition input)\n    KAGGLE_INPUT_DIR = Path(\"/kaggle/input/MABe-mouse-behavior-detection\")\n\n    # <<< Replace with your uploaded weights and class map >>>\n    MODEL_PATH = \"/kaggle/input/2-5d-cnn-for-mabe/model_fold0_best.pth\"\n    CLASS_MAP_PATH = \"/kaggle/input/2-5d-cnn-for-mabe/class_to_idx.json\"\n\n    # Backbone and input settings (must match training)\n    MODEL_NAME = \"efficientnet_b0\"\n    IMG_SIZE = 128\n    BATCH_SIZE = 256\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # Post-processing\n    # PREDICTION_THRESHOLD = 0.3\n    PREDICTION_THRESHOLD = 0.05\n    MIN_BOUT_LENGTH = 3\n\n    # Frame subsampling for speed (e.g., use every 5th frame)\n    # INFER_FRAME_STEP = 5\n    INFER_FRAME_STEP = 2\n\n    # Temporal offsets (2.5D channels); \n    OFFSETS = [-25, -15, -10, -5, -2, 0, 2, 5, 10, 15, 25]\n\nprint(f\"Using device: {CFG.DEVICE}\")\nif torch.cuda.is_available():\n    print(f\"CUDA device: {torch.cuda.get_device_name(0)}\")\n\n# ===================================================================\n# 2. Utilities: pose cleaning and rasterization\n# ===================================================================\ndef drop_body_parts(df: pd.DataFrame) -> pd.DataFrame:\n    \"\"\"Remove head-mounted parts and unify left/right suffix for bodypart.\"\"\"\n    df = df[~df['bodypart'].str.contains('headpiece')]\n    df['bodypart'] = (\n        df['bodypart']\n        .str.replace('_left',  '', regex=False)\n        .str.replace('_right', '', regex=False)\n    )\n    return df\n\n# Ordered coarse body parts (for a simple grayscale mapping)\nBODYPART_ORDER = ['nose','ear','neck','lateral','body_center','hip','tail_base','tail_midpoint','tail_tip']\n_bodypart_pos_map = {p: i for i, p in enumerate(BODYPART_ORDER)}\n\ndef _gray(mouse_id: int, bodypart: str) -> tuple:\n    \"\"\"Legacy grayscale mapper used by the fallback PIL path.\"\"\"\n    base = {1: 0.9, 2: 0.7, 3: 0.5, 4: 0.3}.get(mouse_id, 0.1)\n    rng = 0.1\n    pos = next((_bodypart_pos_map[p] for p in BODYPART_ORDER if bodypart.startswith(p)), -1)\n    ratio = 0.5 if pos == -1 else pos / (len(BODYPART_ORDER) - 1)\n    v = np.clip(base - ratio * rng, 0.0, 1.0)\n    return (v, v, v)\n\ndef keypoints_to_image_numpy(frame_df: pd.DataFrame, width: int, height: int, img_size: int) -> Image.Image:\n    \"\"\"Fallback (slower) single-frame rasterization via PIL; returns 1-channel image.\"\"\"\n    img = np.zeros((img_size, img_size), dtype=np.uint8)\n    if frame_df is not None and len(frame_df) > 0:\n        width  = width  if (width  is not None and width  > 0) else 1\n        height = height if (height is not None and height > 0) else 1\n        x = (frame_df['x'].values * (img_size / width)).astype(int)\n        y = (frame_df['y'].values * (img_size / height)).astype(int)\n        x = np.clip(x, 0, img_size - 1)\n        y = np.clip(y, 0, img_size - 1)\n        vals = [int(_gray(int(r['mouse_id']), r['bodypart'])[0] * 255) for _, r in frame_df.iterrows()]\n        img[y, x] = vals\n    return Image.fromarray(img, mode='L')\n\n# ===== Fast path used in inference (matches training approach) =====\ndef _precompute_tracking_buffers(df: pd.DataFrame):\n    \"\"\"\n    Precompute arrays for fast per-frame rasterization from a cleaned tracking DataFrame:\n      - grayscale per row (uint8) from mouse-id and body-part position,\n      - a mapping: video_frame -> row indices for that frame.\n    \"\"\"\n    key = df['bodypart'].str.extract(\n        r'^(nose|ear|neck|lateral|body_center|hip|tail_base|tail_midpoint|tail_tip)'\n    )[0]\n    pos = key.map(_bodypart_pos_map).fillna(-1).astype(np.int16).to_numpy()\n\n    base = df['mouse_id'].map({1: 0.9, 2: 0.7, 3: 0.5, 4: 0.3}).fillna(0.1).astype(np.float32).to_numpy()\n    ratio = np.where(pos >= 0, pos / (len(BODYPART_ORDER) - 1), 0.5).astype(np.float32)\n    gray = np.clip(base - 0.1 * ratio, 0.0, 1.0)\n    gray_u8 = (gray * 255.0).astype(np.uint8)\n\n    groups = df.groupby('video_frame', sort=True).indices\n    frame_to_rows = {int(k): v for k, v in groups.items()}\n\n    return {\n        'x': df['x'].to_numpy(np.float32),\n        'y': df['y'].to_numpy(np.float32),\n        'gray': gray_u8,\n        'frame_to_rows': frame_to_rows\n    }\n\ndef _rasterize_frame_tensor(buf: dict, frame: int, width: int, height: int, img_size: int) -> torch.Tensor:\n    \"\"\"\n    Rasterize a single frame directly into a torch tensor (1, H, W) in [0, 1].\n    Avoids PIL for speed; uses precomputed buffers.\n    \"\"\"\n    img = np.zeros((img_size, img_size), dtype=np.uint8)\n    idx = buf['frame_to_rows'].get(int(frame), None)\n    if idx is not None and len(idx) > 0:\n        sx = img_size / max(int(width), 1)\n        sy = img_size / max(int(height), 1)\n        x = (buf['x'][idx] * sx).astype(np.int32)\n        y = (buf['y'][idx] * sy).astype(np.int32)\n        x = np.clip(x, 0, img_size - 1)\n        y = np.clip(y, 0, img_size - 1)\n        img[y, x] = buf['gray'][idx]\n    return torch.from_numpy(img).unsqueeze(0).to(dtype=torch.float32).div_(255.0)\n\ndef build_2p5d_clip(\n    tracking_df_indexed,\n    center_frame_idx: int,\n    width: int,\n    height: int,\n    img_size: int,\n    offsets = CFG.OFFSETS\n) -> torch.Tensor:\n    \"\"\"\n    Build a 2.5D clip by stacking frames at specified temporal offsets around a center frame.\n    Accepts either:\n      - dict buffer from _precompute_tracking_buffers (fast path), or\n      - a DataFrame indexed by 'video_frame' (fallback path).\n    Returns: torch.FloatTensor [C=len(offsets), H, W] with values in [0, 1].\n    \"\"\"\n    imgs = []\n\n    # Fast path (dict buffer)\n    if isinstance(tracking_df_indexed, dict) and 'frame_to_rows' in tracking_df_indexed:\n        buf = tracking_df_indexed\n        for off in offsets:\n            f = center_frame_idx + off\n            t = _rasterize_frame_tensor(buf, f, width, height, img_size)  # (1, H, W)\n            imgs.append(t)\n        return torch.cat(imgs, dim=0)\n\n    # Fallback path (DataFrame + PIL)\n    for off in offsets:\n        f = center_frame_idx + off\n        if f in tracking_df_indexed.index:\n            frame_df = tracking_df_indexed.loc[f]\n            if isinstance(frame_df, pd.Series):\n                frame_df = frame_df.to_frame().T\n        else:\n            frame_df = None  # empty frame\n        pil = keypoints_to_image_numpy(frame_df, width, height, img_size)\n        t = T.ToTensor()(pil)  # (1, H, W) in [0, 1]\n        imgs.append(t)\n    return torch.cat(imgs, dim=0)\n\n# ===================================================================\n# 3. Model: 2.5D EfficientNet\n# ===================================================================\nclass MABeEfficientNet2p5D(nn.Module):\n    \"\"\"\n    EfficientNet backbone with a custom input channel count (= len(OFFSETS)).\n    The classifier head is replaced to match the number of classes.\n    \"\"\"\n    def __init__(self, model_name, n_classes, pretrained=False, in_chans: int = None):\n        super().__init__()\n        if in_chans is None:\n            in_chans = len(CFG.OFFSETS)\n        m = timm.create_model(model_name, pretrained=pretrained, in_chans=in_chans)\n        in_features = m.classifier.in_features\n        m.classifier = nn.Linear(in_features, n_classes)\n        self.model = m\n\n    def forward(self, x):\n        return self.model(x)\n\n# ===================================================================\n# 4. Label decoding and whitelist helpers\n# ===================================================================\n_PATTERNS = [\n    re.compile(r'^(\\d+)[_\\-\\.](\\d+)[_\\-\\.](.+)$'),\n    re.compile(r'^(\\d+)[_\\-\\.]self[_\\-\\.](.+)$'),\n    re.compile(r'^mouse(\\d+)[_\\-\\.]mouse(\\d+)[_\\-\\.](.+)$'),\n    re.compile(r'^mouse(\\d+)[_\\-\\.]self[_\\-\\.](.+)$'),\n]\n\ndef decode_class_label(label: str):\n    \"\"\"\n    Convert a class label into (agent_id, target_id, action) triples in 'mouse{n}' form.\n    Accepts a few common label patterns used in training.\n    \"\"\"\n    s = label.strip().lower()\n    for pat in _PATTERNS:\n        m = pat.match(s)\n        if m:\n            if len(m.groups()) == 3:\n                a, t, act = m.groups()\n                return f\"mouse{int(a)}\", f\"mouse{int(t)}\", act\n            elif len(m.groups()) == 2:\n                a, act = m.groups()\n                return f\"mouse{int(a)}\", \"self\", act\n    return None\n\ndef parse_behaviors_whitelist(behaviors_labeled_str: str):\n    \"\"\"\n    Parse test.csv 'behaviors_labeled' into a set of lowercase (agent, target, action) triplets\n    that are allowed for the video.\n    \"\"\"\n    allow = set()\n    try:\n        items = json.loads(behaviors_labeled_str)\n    except Exception:\n        s = behaviors_labeled_str.strip()\n        if s.startswith('[') and s.endswith(']'):\n            items = json.loads(s)\n        else:\n            items = [x.strip() for x in s.split(',') if x.strip()]\n    for it in items:\n        it = it.replace(\"'\", \"\").strip().lower()\n        parts = [p.strip() for p in it.split(',')]\n        if len(parts) == 3:\n            agent, target, action = parts\n            allow.add((agent, target, action))\n    return allow\n\n# ===================================================================\n# 5. Post-processing: thresholding, boutization, and de-duplication\n# ===================================================================\ndef post_process_and_submit(\n    predictions: np.ndarray,\n    video_id: int,\n    idx_to_class: dict,\n    threshold: float,\n    min_len: int,\n    whitelist: set,\n    present_mice_nums: list,\n    video_max_frame: int,\n    frame_ids: np.ndarray\n) -> pd.DataFrame:\n    \"\"\"\n    Convert per-frame probabilities into valid action bouts:\n      1) Decode class → (agent, target, action)\n      2) Mask by existing mice and whitelist\n      3) Threshold and merge consecutive 1's\n      4) Resolve overlaps by keeping the earliest interval\n    Returns a DataFrame with (video_id, agent_id, target_id, action, start_frame, stop_frame).\n    \"\"\"\n    def mouse_exists(agent_id: str) -> bool:\n        if not agent_id.startswith('mouse'):\n            return False\n        try:\n            m = int(agent_id[5:])\n        except Exception:\n            return False\n        return m in present_mice_nums\n\n    n_frames, n_classes = predictions.shape\n    bouts = []\n\n    for ci in range(n_classes):\n        label = idx_to_class[ci]\n        decoded = decode_class_label(label)\n        if decoded is None:\n            continue\n        agent_id, target_id, action = decoded\n\n        # Mouse presence checks\n        if not mouse_exists(agent_id):\n            continue\n        if target_id != 'self' and not mouse_exists(target_id):\n            continue\n\n        # Enforce per-video whitelist\n        if (agent_id.lower(), target_id.lower(), action.lower()) not in whitelist:\n            continue\n\n        prob = predictions[:, ci]\n        bin_ = (prob > threshold).astype(np.int8)\n        if bin_.sum() < min_len:\n            continue\n\n        diff = np.diff(np.concatenate(([0], bin_, [0])))\n        starts = np.where(diff == 1)[0]\n        stops  = np.where(diff == -1)[0] - 1\n\n        for s_idx, e_idx in zip(starts, stops):\n            if e_idx - s_idx + 1 >= min_len:\n                s_frame = int(frame_ids[s_idx])\n                e_frame = int(frame_ids[e_idx])\n                s_frame = max(0, s_frame)\n                e_frame = min(int(video_max_frame), e_frame)\n                if e_frame > s_frame:\n                    bouts.append({\n                        'video_id': int(video_id),\n                        'agent_id': agent_id,\n                        'target_id': target_id,\n                        'action': action,\n                        'start_frame': s_frame,\n                        'stop_frame': e_frame\n                    })\n\n    if not bouts:\n        return pd.DataFrame(columns=['video_id','agent_id','target_id','action','start_frame','stop_frame'])\n\n    # Within (agent, target): keep earliest interval if overlaps occur\n    df = pd.DataFrame(bouts).sort_values(\n        ['video_id','agent_id','target_id','start_frame','stop_frame','action']\n    )\n    cleaned_groups = []\n    for (vid, ag, tg), g in df.groupby(['video_id','agent_id','target_id'], sort=False):\n        g = g.sort_values(['start_frame','stop_frame']).reset_index(drop=True)\n        keep_rows = []\n        last_stop = -1\n        for r in g.itertuples(index=False):\n            if r.start_frame >= last_stop:\n                keep_rows.append(r)\n                last_stop = r.stop_frame\n            else:\n                # If overlapped, drop the later one (mirrors sample behavior)\n                pass\n        if keep_rows:\n            cleaned_groups.append(pd.DataFrame(keep_rows))\n    if cleaned_groups:\n        out = pd.concat(cleaned_groups, ignore_index=True)\n        out.columns = ['video_id','agent_id','target_id','action','start_frame','stop_frame']\n    else:\n        out = pd.DataFrame(columns=['video_id','agent_id','target_id','action','start_frame','stop_frame'])\n\n    out = out.dropna()\n    out = out[(out['stop_frame'] > out['start_frame'])].copy()\n    return out\n\n# ===================================================================\n# 6. Inference loop\n# ===================================================================\nprint(f\"Starting inference on {CFG.DEVICE}...\")\n\ntest_meta_df = pd.read_csv(CFG.KAGGLE_INPUT_DIR / \"test.csv\")\n\ndef present_mice_nums_from_row(row: pd.Series):\n    \"\"\"Infer which mouse indices (1..4) exist in the video from metadata.\"\"\"\n    nums = []\n    for i in [1, 2, 3, 4]:\n        if not pd.isna(row.get(f'mouse{i}_strain', np.nan)):\n            nums.append(i)\n    return nums\n\n# Load class map (class -> index), then invert\nwith open(CFG.CLASS_MAP_PATH, 'r') as f:\n    class_to_idx = json.load(f)\nidx_to_class = {v: k for k, v in class_to_idx.items()}\nn_classes = len(idx_to_class)\n\n# Build model and load weights\nmodel = MABeEfficientNet2p5D(\n    CFG.MODEL_NAME, n_classes=n_classes, pretrained=False, in_chans=len(CFG.OFFSETS)\n).to(CFG.DEVICE)\nmodel.load_state_dict(torch.load(CFG.MODEL_PATH, map_location=CFG.DEVICE))\nmodel.eval()\n\n# Normalization (must match training)\nnorm_transform = T.Normalize(mean=[0.5] * len(CFG.OFFSETS), std=[0.5] * len(CFG.OFFSETS))\n\nsubmissions = []\n\nfor _, row in tqdm(test_meta_df.iterrows(), total=len(test_meta_df), desc=\"Inferencing test videos\"):\n    video_id = int(row['video_id'])\n    lab_id   = str(row['lab_id'])\n\n    # Tracking file path (some datasets may omit the lab subfolder)\n    fpath = CFG.KAGGLE_INPUT_DIR / \"test_tracking\" / lab_id / f\"{video_id}.parquet\"\n    if not fpath.exists():\n        fpath = CFG.KAGGLE_INPUT_DIR / \"test_tracking\" / f\"{video_id}.parquet\"\n    if not fpath.exists():\n        print(f\"[WARN] file not found: {video_id}\")\n        continue\n\n    try:\n        tracking_df = pd.read_parquet(fpath)\n        tracking_df = drop_body_parts(tracking_df)\n\n        # Precompute fast buffers for rasterization\n        tracking_buf = _precompute_tracking_buffers(tracking_df)\n\n        # Image size in pixels (fallback to 1 to avoid division by zero)\n        width_val  = row.get('video_width_pix',  np.nan)\n        height_val = row.get('video_height_pix', np.nan)\n        width  = int(width_val)  if pd.notna(width_val)  and width_val  else 1\n        height = int(height_val) if pd.notna(height_val) and height_val else 1\n\n        # Frame IDs present in the file, optionally subsampled for speed\n        all_frames_full = np.array(sorted(tracking_buf['frame_to_rows'].keys()))\n        if len(all_frames_full) == 0:\n            continue\n        frame_ids = all_frames_full[::CFG.INFER_FRAME_STEP] if CFG.INFER_FRAME_STEP > 1 else all_frames_full\n        if len(frame_ids) == 0:\n            frame_ids = all_frames_full[[-1]]\n\n        # For clipping bounds and final validation\n        video_max_frame = int(all_frames_full.max())\n\n        whitelist = parse_behaviors_whitelist(row['behaviors_labeled'])\n        present_mice_nums = present_mice_nums_from_row(row)\n\n        # Run model\n        preds_list = []\n        with torch.no_grad():\n            for i in range(0, len(frame_ids), CFG.BATCH_SIZE):\n                frames_batch = frame_ids[i:i + CFG.BATCH_SIZE]\n                clips = []\n                for fr in frames_batch:\n                    clip = build_2p5d_clip(\n                        tracking_df_indexed=tracking_buf,\n                        center_frame_idx=int(fr),\n                        width=width, height=height,\n                        img_size=CFG.IMG_SIZE,\n                        offsets=CFG.OFFSETS\n                    )  # (C, H, W) in [0, 1]\n                    clips.append(clip)\n                if not clips:\n                    continue\n                x = torch.stack(clips).to(CFG.DEVICE, non_blocking=True)  # (B, C, H, W)\n                x = norm_transform(x)\n                logits = model(x)\n                probs = torch.sigmoid(logits).cpu().numpy()\n                preds_list.append(probs)\n\n        # Build per-video submission rows\n        if len(preds_list) == 0:\n            video_df = pd.DataFrame(columns=['video_id','agent_id','target_id','action','start_frame','stop_frame'])\n        else:\n            preds = np.concatenate(preds_list, axis=0)  # (len(frame_ids), n_classes)\n            video_df = post_process_and_submit(\n                preds, video_id, idx_to_class,\n                CFG.PREDICTION_THRESHOLD, CFG.MIN_BOUT_LENGTH,\n                whitelist, present_mice_nums, video_max_frame,\n                frame_ids=frame_ids\n            )\n\n        if not video_df.empty:\n            submissions.append(video_df)\n\n        # Cleanup\n        del tracking_df, tracking_buf, preds_list\n        gc.collect()\n\n    except Exception as e:\n        print(f\"[ERROR] {video_id} :: {e}\")\n\n# ===================================================================\n# 7. Save submission.csv (schema strictly matches sample)\n# ===================================================================\nif len(submissions) > 0:\n    submission_df = pd.concat(submissions, ignore_index=True)\nelse:\n    submission_df = pd.DataFrame(columns=['video_id','agent_id','target_id','action','start_frame','stop_frame'])\n\n# Column order and dtypes\nsubmission_df = submission_df[['video_id','agent_id','target_id','action','start_frame','stop_frame']].copy()\nsubmission_df['video_id'] = submission_df['video_id'].astype(int)\nsubmission_df['agent_id'] = submission_df['agent_id'].astype(str)\nsubmission_df['target_id'] = submission_df['target_id'].astype(str)\nsubmission_df['action'] = submission_df['action'].astype(str)\nsubmission_df['start_frame'] = submission_df['start_frame'].astype(int)\nsubmission_df['stop_frame']  = submission_df['stop_frame'].astype(int)\n\n# Keep only valid intervals (start < stop)\nsubmission_df = submission_df[submission_df['stop_frame'] > submission_df['start_frame']].reset_index(drop=True)\n\n# Write with index named 'row_id' (as in the sample)\nsubmission_df.index.name = 'row_id'\nsubmission_df.to_csv('submission.csv')\n\nprint(\"\\nsubmission.csv created successfully\")\nprint(\"Shape (rows, cols):\", submission_df.shape)\nprint(submission_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-02T01:28:40.147133Z","iopub.execute_input":"2025-10-02T01:28:40.147404Z","iopub.status.idle":"2025-10-02T01:29:14.966003Z","shell.execute_reply.started":"2025-10-02T01:28:40.147375Z","shell.execute_reply":"2025-10-02T01:29:14.965321Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}