{"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":13874099,"sourceType":"competition"},{"sourceId":14155623,"sourceType":"datasetVersion","datasetId":8887818}],"dockerImageVersionId":31090,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## 📘 MABe: 2.5D MaxVit + 3D Stem with OOF Optimization\n\n**Goal:** Generate `submission.csv` for the MABe Challenge (Mouse Social Action Recognition).\n\n**Context:**\nThis notebook is an updated version of my previous approach: [2.5D CNN for Mouse Social Action Detection](https://www.kaggle.com/code/kmatsu01/2-5d-cnn-for-mouse-social-action-detection).\nWhile the previous version used EfficientNet-B0 with simple 2.5D stacking, this updated version introduces several improvements\n    \n1. **Backbone:** Upgraded to **MaxVit Tiny** (via `timm`) for better spatial feature extraction.\n2. **3D Stem:** Instead of simple channel stacking, we use a lightweight **3D Convolutional stem** to process temporal inputs before feeding them into the 2D backbone.\n3. **FPS-Adaptive Offsets:** Temporal offsets are calculated in seconds rather than fixed frames, making the model robust to varying frame rates.\n4. **OOF Optimization:** Post-processing parameters (prediction thresholds, minimum bout lengths, and gap filling) are optimized per action class using Out-Of-Fold (OOF) data.\n\n### **What this notebook does**\n\n1. **Preprocessing:**\n* Rasterizes mouse keypoints into grayscale images.\n* Uses a \"Dot Size\" > 1 to draw thicker keypoints for better visibility.\n\n\n2. **Model Inference:**\n* Input: A stack of frames (2.5D) centered on the target frame.\n* Architecture: `Input -> 3D Conv Layers -> MaxVit Tiny 2D Backbone -> Classifier`.\n\n\n3. **Advanced Post-processing:**\n* **Class-specific Thresholds:** Loads optimized probability thresholds for each action.\n* **Bout Refinement:** Merges fragmented bouts (Gap Filling) and removes short noise (Min Length) based on class-specific parameters.\n* **Whitelist Enforcement:** Ensures only valid (Agent, Target, Action) triplets defined in `test.csv` are predicted.\n\n\n**Train Notebook:** [MABe: 2.5D MaxVit + 3D Stem [Train & OOF Opt]](https://www.kaggle.com/code/kmatsu01/mabe-2-5d-maxvit-3d-stem-train-oof-opt)","metadata":{}},{"cell_type":"code","source":"# ===================================================================\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\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms as T\nimport timm\nimport psutil\n\nwarnings.filterwarnings('ignore')\n\n# ===================================================================\n# 1. Configuration\n# ===================================================================\nclass CFG:\n    # Base dataset path\n    KAGGLE_INPUT_DIR = Path(\"/kaggle/input/MABe-mouse-behavior-detection\")\n\n    # =================================================================\n    # PATHS: Directory containing model weights and config files\n    # =================================================================\n    FOLD_ARTIFACT_DIR = Path(\"/kaggle/input/0026-v2-idswap0-5-ep12-lab5folds-maxvittiny224\")\n\n    # Folds to use for ensemble\n    FOLDS_TO_USE = [1, 2, 3, 4]\n\n    # Class map file\n    CLASS_MAP_PATH = FOLD_ARTIFACT_DIR / \"class_to_idx.json\"\n\n    # Model Architecture Settings\n    MODEL_NAME = \"maxvit_tiny_tf_224\" # MaxVit Tiny\n    IMG_SIZE = 224                   # MaxVit input size\n    \n    # =================================================================\n    # PERFORMANCE & HARDWARE SETTINGS\n    # =================================================================\n    # Batch size: With FP16 and Inference Mode, 128 should fit on T4.\n    # If OOM occurs, reduce to 64.\n    BATCH_SIZE = 128\n    \n    # Enable Half Precision (FP16)\n    USE_FP16 = True\n\n    # Multi-GPU Strategy\n    # Mapping folds to specific devices\n    # Fold 0, 1 -> cuda:0\n    # Fold 2, 3 -> cuda:1\n    DEVICE_MAP = {\n        1: 'cuda:0',\n        2: 'cuda:0',\n        3: 'cuda:1',\n        4: 'cuda:1'\n    }\n\n    # Logging\n    ENABLE_MEMORY_LOG = False        # Set False for submission\n    \n    # Default device (fallback)\n    DEFAULT_DEVICE = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n\n    # --- Preprocessing Settings (Must match training) ---\n    DOT_SIZE = 3 \n\n    # ===== Post-processing Defaults =====\n    PREDICTION_THRESHOLD_DEFAULT = 0.34\n    DEFAULT_MIN_BOUT_LENGTH = 4\n    DEFAULT_MAX_GAP         = 0\n\n    # Frame Subsampling (Inference Speedup)\n    INFER_FRAME_STEP = 6\n\n    # ===== Temporal offsets (2.5D channels) =====\n    USE_OFFSETS_SEC = True\n    OFFSETS_SEC = [\n        -25/30.0, -15/30.0, -10/30.0, -5/30.0, -2/30.0,\n         0.0,\n         2/30.0,   5/30.0,   10/30.0,  15/30.0,  25/30.0\n    ]\n    OFFSETS = [-25, -15, -10, -5, -2, 0, 2, 5, 10, 15, 25]\n\n    # ===== Threshold blending / clipping =====\n    THR_BLEND_LAMBDA = 1.0\n    THR_MIN = 0.05\n    THR_MAX = 0.95\n\n# === Memory Logging Utility ===\ndef log_memory_usage(prefix=\"\"):\n    if not CFG.ENABLE_MEMORY_LOG:\n        return\n        \n    # RAM usage\n    mem = psutil.virtual_memory()\n    ram_gb = mem.used / (1024**3)\n    ram_total_gb = mem.total / (1024**3)\n    \n    # VRAM usage\n    gpu_msg = \"\"\n    if torch.cuda.is_available():\n        for i in range(torch.cuda.device_count()):\n            vram_allocated = torch.cuda.memory_allocated(i) / (1024**3)\n            vram_max = torch.cuda.max_memory_allocated(i) / (1024**3)\n            gpu_msg += f\"[GPU{i}: {vram_allocated:.2f}GB / Max {vram_max:.2f}GB] \"\n    else:\n        gpu_msg = \"VRAM: N/A\"\n        \n    print(f\"[{prefix}] RAM: {ram_gb:.2f}/{ram_total_gb:.2f}GB | {gpu_msg}\")\n\n# Check GPU availability\nprint(f\"Available GPUs: {torch.cuda.device_count()}\")\nfor i in range(torch.cuda.device_count()):\n    print(f\"  GPU {i}: {torch.cuda.get_device_name(i)}\")\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 suffixes.\"\"\"\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\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    \"\"\"Calculate grayscale intensity for body parts.\"\"\"\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 rasterization using PIL.\"\"\"\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\n        dot = max(1, int(CFG.DOT_SIZE))\n        if dot == 1:\n            img[y, x] = vals\n        else:\n            h, w = img.shape\n            for xi, yi, gi in zip(x, y, vals):\n                x2 = min(xi + dot, w)\n                y2 = min(yi + dot, h)\n                img[yi:y2, xi:x2] = gi\n    return Image.fromarray(img, mode='L')\n\ndef _precompute_tracking_buffers(df: pd.DataFrame):\n    \"\"\"Precompute arrays for fast rasterization.\"\"\"\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    \"\"\"Fast rasterization to Tensor.\"\"\"\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        g = buf['gray'][idx]\n\n        dot = max(1, int(CFG.DOT_SIZE))\n        if dot == 1:\n            img[y, x] = g\n        else:\n            h, w = img.shape\n            for xi, yi, gi in zip(x, y, g):\n                x2 = min(xi + dot, w)\n                y2 = min(yi + dot, h)\n                img[yi:y2, xi:x2] = gi\n\n    return torch.from_numpy(img).unsqueeze(0).to(dtype=torch.float32).div_(255.0)\n\n# --- FPS and Offset Helpers ---\ndef _fps_or_default(v, default=30.0):\n    try:\n        v = float(v)\n        if not np.isfinite(v) or v <= 0:\n            return float(default)\n        return v\n    except Exception:\n        return float(default)\n\ndef _frame_offsets_for_fps(fps: float) -> list[int]:\n    if CFG.USE_OFFSETS_SEC:\n        return [int(round(sec * fps)) for sec in CFG.OFFSETS_SEC]\n    else:\n        return list(CFG.OFFSETS)\n\ndef build_2p5d_clip(tracking_df_indexed, center_frame_idx: int, width: int, height: int, img_size: int, offsets) -> torch.Tensor:\n    imgs = []\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)\n            imgs.append(t)\n        return torch.cat(imgs, dim=0)\n    # fallback\n    for off in offsets:\n        f = center_frame_idx + off\n        if hasattr(tracking_df_indexed, \"index\") and (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\n        pil = keypoints_to_image_numpy(frame_df, width, height, img_size)\n        t = T.ToTensor()(pil)\n        imgs.append(t)\n    return torch.cat(imgs, dim=0)\n\n# ===================================================================\n# 3. Model: 3D Stem + 2D Backbone (MABeEfficientNet)\n# ===================================================================\nclass MABeEfficientNet(nn.Module):\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_SEC) if CFG.USE_OFFSETS_SEC else len(CFG.OFFSETS))\n        \n        # 3D Stem\n        self.spatial_conv = nn.Conv3d(\n            in_channels=1, out_channels=in_chans,\n            kernel_size=(1, 3, 3), padding=(0, 1, 1), bias=False\n        )\n        self.spatial_bn = nn.BatchNorm3d(in_chans)\n        self.spatial_act = nn.SiLU(inplace=True)\n        \n        self.temporal_conv = nn.Conv3d(\n            in_channels=in_chans, out_channels=in_chans,\n            kernel_size=(3, 1, 1), padding=(1, 0, 0), groups=in_chans, bias=False\n        )\n        self.temporal_bn = nn.BatchNorm3d(in_chans)\n        self.temporal_act = nn.SiLU(inplace=True)\n\n        # 2D Backbone (MaxVit)\n        self.backbone = timm.create_model(\n            model_name,\n            pretrained=pretrained,\n            in_chans=in_chans,\n            num_classes=n_classes,\n        )    \n    \n    def forward(self, x):\n        # x: [B, C, H, W]\n        x = x.unsqueeze(1)                              # [B, 1, C, H, W]\n        x = self.spatial_act(self.spatial_bn(self.spatial_conv(x)))\n        x = self.temporal_act(self.temporal_bn(self.temporal_conv(x)))\n        x = x.mean(dim=2)                               # [B, C, H, W]\n        x = self.backbone(x)                            # [B, n_classes]\n        return 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    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    allow = set()\n    try:\n        items = json.loads(behaviors_labeled_str)\n    except Exception:\n        s = str(behaviors_labeled_str or \"\").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 or []:\n        it = str(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\ndef build_video_class_mask(idx_to_class: dict[int, str], whitelist: set) -> np.ndarray:\n    n_classes = len(idx_to_class)\n    mask = np.zeros(n_classes, dtype=np.float32)\n    for ci, label in idx_to_class.items():\n        decoded = decode_class_label(label)\n        if decoded is None:\n            continue\n        a, t, act = decoded\n        if (a.lower(), t.lower(), act.lower()) in whitelist:\n            mask[ci] = 1.0\n    if mask.sum() == 0:\n        mask[:] = 1.0\n    return mask\n\n# ===================================================================\n# 4.5 Aggregating Metadata (Thresholds/Bouts) across Folds\n# ===================================================================\ndef _load_thr_json_for_fold(artifact_dir: Path, fold: int):\n    p = artifact_dir / f\"thr_per_action_oof_fold{fold}.json\"\n    if p.exists():\n        with open(p, \"r\") as f:\n            return json.load(f)\n    return None\n\ndef _load_bout_csv_for_fold(artifact_dir: Path, fold: int):\n    p = artifact_dir / f\"per_action_bout_params_fold{fold}.csv\"\n    if p.exists():\n        try:\n            return pd.read_csv(p)\n        except Exception:\n            try:\n                return pd.read_parquet(p)\n            except Exception:\n                return None\n    return None\n\ndef _aggregate_thresholds(idx_to_class: dict[int, str], thr_json_list: list[dict]) -> np.ndarray:\n    \"\"\"Load thresholds from multiple folds and aggregate (median/mean).\"\"\"\n    C = len(idx_to_class)\n    percls_vals = {ci: [] for ci in range(C)}\n    global_vals = []\n    \n    for js in thr_json_list:\n        if not js: continue\n        per = js.get(\"per_action_thresholds\") or js.get(\"thresholds\") or {}\n        g = js.get(\"global_best_threshold\", None)\n        if isinstance(g, (int, float)):\n            global_vals.append(float(g))\n            \n        for ci, label in idx_to_class.items():\n            key = str(label).strip().lower()\n            v = per.get(key, None)\n            if isinstance(v, (int, float)):\n                percls_vals[ci].append(float(v))\n                \n    global_med = np.median(global_vals) if len(global_vals)>0 else float(CFG.PREDICTION_THRESHOLD_DEFAULT)\n    out = np.full(C, float(global_med), dtype=np.float32)\n    \n    for ci in range(C):\n        if len(percls_vals[ci]) > 0:\n            med = float(np.median(percls_vals[ci]))\n        else:\n            med = float(global_med)\n            \n        t_final = CFG.THR_BLEND_LAMBDA * med + (1.0 - CFG.THR_BLEND_LAMBDA) * float(global_med)\n        out[ci] = float(np.clip(t_final, CFG.THR_MIN, CFG.THR_MAX))\n    return out\n\ndef _pick_cols(df: pd.DataFrame, candidates: list[str]):\n    cols = [c for c in candidates if c in df.columns]\n    return cols[0] if cols else None\n\ndef _normalize_bout_df(df_raw: pd.DataFrame) -> dict[str, tuple]:\n    if df_raw is None or df_raw.empty:\n        return {}\n    df = df_raw.copy()\n    df.columns = [str(c).strip().lower() for c in df.columns]\n    \n    col_cls = _pick_cols(df, ['action_class','class','label','name'])\n    col_min = _pick_cols(df, ['min_bout_len','min_len','min_bout','minlength','min_bout_sec'])\n    col_gap = _pick_cols(df, ['max_gap','gap','maxgap','max_gap_len','max_gap_sec'])\n    \n    if col_cls is None:\n        return {}\n    \n    out = {}\n    for r in df.itertuples(index=False):\n        name = str(getattr(r, col_cls)).strip().lower()\n        try: minv = getattr(r, col_min) if col_min else CFG.DEFAULT_MIN_BOUT_LENGTH\n        except: minv = CFG.DEFAULT_MIN_BOUT_LENGTH\n        try: gapv = getattr(r, col_gap) if col_gap else CFG.DEFAULT_MAX_GAP\n        except: gapv = CFG.DEFAULT_MAX_GAP\n        \n        try: minv = int(round(float(minv)))\n        except: minv = int(CFG.DEFAULT_MIN_BOUT_LENGTH)\n        try: gapv = int(round(float(gapv)))\n        except: gapv = int(CFG.DEFAULT_MAX_GAP)\n        out[name] = (minv, gapv)\n    return out\n\ndef _aggregate_bout_params(idx_to_class: dict[int, str], bout_df_list: list[pd.DataFrame]):\n    percls_min = {k: [] for k in idx_to_class.values()}\n    percls_gap = {k: [] for k in idx_to_class.values()}\n    \n    for df in bout_df_list:\n        mapping = _normalize_bout_df(df)\n        for cls_name in percls_min.keys():\n            key = str(cls_name).strip().lower()\n            if key in mapping:\n                mn, gp = mapping[key]\n                percls_min[cls_name].append(mn)\n                percls_gap[cls_name].append(gp)\n                \n    C = len(idx_to_class)\n    min_len_arr = np.full(C, int(CFG.DEFAULT_MIN_BOUT_LENGTH), dtype=np.int32)\n    max_gap_arr = np.full(C, int(CFG.DEFAULT_MAX_GAP), dtype=np.int32)\n    \n    for ci, label in idx_to_class.items():\n        mins = percls_min[label]\n        gaps = percls_gap[label]\n        if len(mins) > 0:\n            min_len_arr[ci] = int(round(np.median(mins)))\n        if len(gaps) > 0:\n            max_gap_arr[ci] = int(round(np.median(gaps)))\n            \n    return min_len_arr, max_gap_arr\n\n# ===================================================================\n# 5. Post-processing: Bout Generation\n# ===================================================================\ndef post_process_and_submit(\n    predictions: np.ndarray,         \n    video_id: int,\n    idx_to_class: dict,\n    thresholds_by_class: np.ndarray, \n    min_len_by_class: np.ndarray,    \n    max_gap_by_class: np.ndarray,    \n    whitelist: set,\n    present_mice_nums: list,\n    video_max_frame: int,\n    frame_ids: np.ndarray            ) -> pd.DataFrame:\n\n    def mouse_exists(agent_id: str) -> bool:\n        if not agent_id.startswith('mouse'): return False\n        try: m = int(agent_id[5:])\n        except: 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: continue\n        agent_id, target_id, action = decoded\n\n        if not mouse_exists(agent_id): continue\n        if target_id != 'self' and not mouse_exists(target_id): continue\n        if (agent_id.lower(), target_id.lower(), action.lower()) not in whitelist: continue\n\n        thr = float(thresholds_by_class[ci])\n        prob = predictions[:, ci]\n        bin_ = (prob > thr).astype(np.int8)\n\n        if bin_.sum() == 0: 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        if len(starts) == 0: continue\n\n        g = int(max_gap_by_class[ci])\n        merged = []\n        cur_s_idx = int(starts[0])\n        cur_e_idx = int(stops[0])\n        for k in range(1, len(starts)):\n            prev_stop_frame = int(frame_ids[cur_e_idx])\n            next_start_frame = int(frame_ids[int(starts[k])])\n            gap_len_raw = next_start_frame - prev_stop_frame - 1\n            if gap_len_raw <= g:\n                cur_e_idx = int(stops[k])\n            else:\n                merged.append((cur_s_idx, cur_e_idx))\n                cur_s_idx, cur_e_idx = int(starts[k]), int(stops[k])\n        merged.append((cur_s_idx, cur_e_idx))\n\n        min_len_raw = int(min_len_by_class[ci])\n        for s_idx, e_idx in merged:\n            s_frame = int(frame_ids[s_idx])\n            e_frame = int(frame_ids[e_idx])\n            raw_len = e_frame - s_frame + 1\n            if raw_len >= min_len_raw:\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    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), gdf in df.groupby(['video_id','agent_id','target_id'], sort=False):\n        gdf = gdf.sort_values(['start_frame','stop_frame']).reset_index(drop=True)\n        keep_rows = []\n        last_stop = -1\n        for r in gdf.itertuples(index=False):\n            if r.start_frame >= last_stop:\n                keep_rows.append(r)\n                last_stop = r.stop_frame\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 using {torch.cuda.device_count()} GPU(s)...\")\n\ntest_meta_df = pd.read_csv(CFG.KAGGLE_INPUT_DIR / \"test.csv\")\n\ndef present_mice_nums_from_row(row: pd.Series):\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\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# === Load Models for Folds 0,1,2,3 with Manual Device Placement ===\nEFFECTIVE_IN_CHANS = (len(CFG.OFFSETS_SEC) if CFG.USE_OFFSETS_SEC else len(CFG.OFFSETS))\n\n# Stores models categorized by device: {'cuda:0': [model_fold0, model_fold1], 'cuda:1': [...]}\nmodels_by_device = {f\"cuda:{i}\": [] for i in range(torch.cuda.device_count())}\n\n# Keep track of which fold was loaded\nloaded_folds = []\n\nprint(f\"Loading models from: {CFG.FOLD_ARTIFACT_DIR}\")\n\n# Safety check for single GPU environment\navailable_devices = [f\"cuda:{i}\" for i in range(torch.cuda.device_count())]\n\nfor f in CFG.FOLDS_TO_USE:\n    mp = CFG.FOLD_ARTIFACT_DIR / f\"model_fold{f}_best.pth\"\n    if mp.exists():\n        # Determine target device\n        if f in CFG.DEVICE_MAP and CFG.DEVICE_MAP[f] in available_devices:\n            target_device = CFG.DEVICE_MAP[f]\n        else:\n            target_device = CFG.DEFAULT_DEVICE\n            \n        print(f\"  Loading Fold {f} onto {target_device}...\")\n        \n        m = MABeEfficientNet(CFG.MODEL_NAME, n_classes=n_classes, pretrained=False, in_chans=EFFECTIVE_IN_CHANS)\n        m.to(target_device)\n        \n        state = torch.load(mp, map_location=target_device)\n        m.load_state_dict(state, strict=True)\n        m.eval()\n        \n        if CFG.USE_FP16:\n            m.half() # Convert to FP16\n            \n        # Register model to the dict\n        if target_device not in models_by_device:\n            models_by_device[target_device] = []\n        models_by_device[target_device].append(m)\n        \n        loaded_folds.append(f)\n    else:\n        print(f\"[WARN] Model not found for fold {f}: {mp}\")\n\ntotal_models = sum(len(v) for v in models_by_device.values())\nif total_models == 0:\n    raise RuntimeError(\"No fold models found. Please check paths.\")\n\nprint(f\"[Models] Successfully loaded {total_models} models across devices.\")\nfor dev, m_list in models_by_device.items():\n    print(f\"  {dev}: {len(m_list)} models\")\n\n# === Load & Aggregate Thresholds/Bout Params ===\nthr_json_list = []\nbout_df_list  = []\nfor f in loaded_folds:\n    thr_js  = _load_thr_json_for_fold(CFG.FOLD_ARTIFACT_DIR, f)\n    bout_df = _load_bout_csv_for_fold(CFG.FOLD_ARTIFACT_DIR, f)\n    if thr_js:  thr_json_list.append(thr_js)\n    if bout_df is not None: bout_df_list.append(bout_df)\n\nthresholds_by_class = _aggregate_thresholds(idx_to_class, thr_json_list)\nmin_len_by_class, max_gap_by_class = _aggregate_bout_params(idx_to_class, bout_df_list)\n\nprint(f\"[Thresholds] Aggregated for {n_classes} classes.\")\nprint(f\"[Bout Params] Aggregated for {n_classes} classes.\")\n\nsubmissions = []\n\n# Log memory once\nlog_memory_usage(\"Start Loop\")\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    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        tracking_buf = _precompute_tracking_buffers(tracking_df)\n\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        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        video_max_frame = int(all_frames_full.max())\n        fps = _fps_or_default(row.get('frames_per_second', 30.0))\n        offsets = _frame_offsets_for_fps(fps)\n\n        whitelist = parse_behaviors_whitelist(row['behaviors_labeled'])\n        video_mask = build_video_class_mask(idx_to_class, whitelist)\n        present_mice_nums = present_mice_nums_from_row(row)\n\n        preds_list = []\n        with torch.no_grad():\n            for i in range(0, len(frame_ids), CFG.BATCH_SIZE):\n                if i == 0:\n                    log_memory_usage(f\"Vid {video_id} start\")\n                    \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=offsets\n                    )\n                    clips.append(clip)\n                if not clips:\n                    continue\n                \n                # --- Multi-GPU Inference ---\n                # 1. Stack raw clips on CPU first (lightweight)\n                x_cpu = torch.stack(clips)\n                \n                sum_probs_cpu = None\n                \n                # 2. Iterate over devices (cuda:0, cuda:1)\n                for device_name, m_list in models_by_device.items():\n                    if len(m_list) == 0: continue\n                    \n                    # Move batch to target GPU and convert to FP16\n                    x_dev = x_cpu.to(device_name, non_blocking=True)\n                    if CFG.USE_FP16:\n                        x_dev = x_dev.half()\n                    \n                    # Run models on this device\n                    probs_sum_dev = None\n                    for m in m_list:\n                        logits = m(x_dev)\n                        probs = torch.sigmoid(logits)\n                        if probs_sum_dev is None:\n                            probs_sum_dev = probs\n                        else:\n                            probs_sum_dev += probs\n                    \n                    # Move aggregated sum from this device to CPU\n                    if probs_sum_dev is not None:\n                        # Ensure float32 for CPU accumulation\n                        p_cpu = probs_sum_dev.float().cpu()\n                        if sum_probs_cpu is None:\n                            sum_probs_cpu = p_cpu\n                        else:\n                            sum_probs_cpu += p_cpu\n\n                # 3. Average across total models\n                if sum_probs_cpu is not None:\n                    probs_avg = (sum_probs_cpu / float(total_models)).numpy()\n                    # Apply Video Mask\n                    probs_avg = probs_avg * video_mask[None, :]\n                    preds_list.append(probs_avg)\n\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)\n            video_df = post_process_and_submit(\n                preds, video_id, idx_to_class,\n                thresholds_by_class,\n                min_len_by_class, max_gap_by_class,\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        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\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\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)\nsubmission_df = submission_df[submission_df['stop_frame'] > submission_df['start_frame']].reset_index(drop=True)\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-12-19T12:39:07.147668Z","iopub.execute_input":"2025-12-19T12:39:07.148207Z","iopub.status.idle":"2025-12-19T12:41:21.053015Z","shell.execute_reply.started":"2025-12-19T12:39:07.148181Z","shell.execute_reply":"2025-12-19T12:41:21.052385Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}