{"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"}],"dockerImageVersionId":31090,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### Notebook Description (Introduction)\n\n**Goal:** Train a 2.5D MaxVit model with a 3D Stem for the MABe Challenge (Mouse Social Action Recognition) and generate optimized thresholds/parameters.\n\n**Context:**\nThis is the **training notebook** for the solution. The trained weights and optimized parameters (thresholds, minimum bout lengths, gap durations) generated here are used in the inference notebook.\n\n* **Inference Notebook:** [MABe: 2.5D MaxVit + 3D Stem [Ensemble & FP16]](https://www.kaggle.com/code/kmatsu01/mabe-2-5d-maxvit-3d-stem)\n\n**Key Features:**\n\n1. **Architecture:**\n* **Input:** Stacked grayscale frames (2.5D) centered on the target frame.\n* **3D Stem:** Processes temporal information using `Conv3d` layers before the 2D backbone.\n* **Backbone:** `MaxVit Tiny` (via `timm`) for robust feature extraction.\n\n\n2. **OOF Optimization:**\n* Calculates Out-Of-Fold (OOF) predictions.\n* Optimizes probability thresholds for each action class to maximize F1 score.\n* Optimizes post-processing parameters (min_bout_len, max_gap) per action.\n\n\n3. **Pseudo-Labeling (Optional):**\n* Includes logic to generate annotations for unlabeled labs if all 5 folds are trained.","metadata":{}},{"cell_type":"code","source":"# ===================================================================\n# 0. Import Libraries\n# ===================================================================\nimport os\nimport gc\nimport re\nimport json\nimport argparse\nimport warnings\nfrom pathlib import Path\nfrom datetime import datetime\nimport time\nimport ast\nimport random\nfrom collections import defaultdict\n\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\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\n\nimport torchvision.transforms as T\nimport timm\nimport albumentations as A\nimport cv2\n\n#For plotting learning curves\nimport matplotlib.pyplot as plt\n\nwarnings.filterwarnings(\"ignore\")\n\n\n# ===================================================================\n# 1. Configuration\n# ===================================================================\nclass CFG:\n    # Random Seed\n    SEED = 42\n\n    # Keypoint dot size (pixels)\n    DOT_SIZE = 3\n\n    # Augmentation settings\n    AugAngle = 10\n    AugShift = 0.20\n    AugScale = 0.15\n    \n    N_SPLITS = 5\n    # FOLDS_TO_RUN = [0]       # Example: Run single fold\n    FOLDS_TO_RUN = [0, 1, 2, 3, 4]  # Run all 5 folds to enable re-annotation\n\n    BASE_PATH = Path(\"../kaggle/input/MABe-mouse-behavior-detection\")\n\n    # Training Hyperparameters\n    MODEL_NAME = \"maxvit_tiny_tf_224\"\n    IMG_SIZE = 224 # MaxVit uses 224\n    BATCH_SIZE = 64\n\n    EPOCHS = 12\n    LR = 2e-4\n    ETA_MIN = 1e-6\n    WEIGHT_DECAY = 0.01\n\n    # Regularization\n    DROP_3D_STEM   = 0.0\n    DROP_2D        = 0.0\n    DROP_DENSE     = 0.0\n    DROP_PATH_RATE = 0.5\n\n    # F1 score monitoring (Used for logging, not for saving best model)\n    F1_GLOBAL_THRESHOLD = 0.25\n    F1_MONITOR_AVERAGE  = \"micro\"\n\n    # ID Swap Augmentation Probability\n    ID_SWAP_PROB = 0.5\n\n    TITLE = f\"0026_v2_idswap{ID_SWAP_PROB}_ep{EPOCHS}_Lab5folds_maxvittiny224\"\n    OUTPUT_DIR = Path(f\"../output/{datetime.now().strftime('%Y%m%d_%H%M%S')}_{TITLE}\")\n\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # Frame subsampling\n    FRAME_SKIP = 11\n\n    # Save visualization of augmented images (First 100 only)\n    SAVE_AUG_VIS = True\n    AUG_VIS_MAX = 100\n    AUG_VIS_PRE_DIR = \"aug_vis\"\n    SAVE_AUG_VIS_DURING_TRAIN = False\n\n    # DataLoader optimizations\n    NUM_WORKERS = min(8, (os.cpu_count() or 2))\n    PIN_MEMORY = torch.cuda.is_available()\n\n    # 2.5D Clip Configuration (Center frame ± offsets)\n    OFFSETS = [-25, -15, -10, -5, -2, 0, 2, 5, 10, 15, 25]\n\n    # FPS Correction\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\n    # Grid for OOF Optimization (Seconds)\n    MIN_LEN_GRID_SEC = (0.05, 0.10, 0.20, 0.33, 0.50, 0.75, 1.00, 1.25, 1.5)\n    MAX_GAP_GRID_SEC = (0.00, 0.03, 0.05, 0.10, 0.20, 0.40, 0.80)\n    DEFAULT_MIN_LEN_SEC = 0.20\n    DEFAULT_MAX_GAP_SEC = 0.05\n    MIN_SUPPORT_BOUTS = 3\n\n    # ====== Re-annotation Settings (Step 2) ======\n    RELABEL_LABELED_VIDEOS = False\n    ENSEMBLE_FOLDS = [0, 1, 2, 3, 4]\n    ALLANN_ROOT_NAME = \"train_annotation_Allannotation\"\n    UNLABELED_LABS_HINT = [\"MABe22_keypoints\", \"MABe22_movies\"]\n\n\n# ===================================================================\n# 1.5 Utility to fix random seeds\n# ===================================================================\ndef seed_everything(seed: int = 42):\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    os.environ.setdefault(\"CUBLAS_WORKSPACE_CONFIG\", \":4096:8\")\n    random.seed(seed); np.random.seed(seed)\n    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    try:\n        torch.use_deterministic_algorithms(True, warn_only=True)\n    except Exception:\n        pass\n\ndef _seed_worker(worker_id: int):\n    worker_seed = CFG.SEED + worker_id\n    np.random.seed(worker_seed); random.seed(worker_seed)\n\ndef _make_loader_common():\n    g = torch.Generator(); g.manual_seed(CFG.SEED)\n    kwargs = dict(\n        batch_size=CFG.BATCH_SIZE,\n        num_workers=CFG.NUM_WORKERS,\n        pin_memory=CFG.PIN_MEMORY,\n        worker_init_fn=_seed_worker,\n        generator=g,\n    )\n    if CFG.NUM_WORKERS > 0:\n        kwargs.update(dict(persistent_workers=True, prefetch_factor=2))\n    return kwargs\n\n\n# ===================================================================\n# 2. Preprocessing (Pose -> Image, etc.)\n# ===================================================================\ndef drop_body_parts(df: pd.DataFrame) -> pd.DataFrame:\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 = [\n    'nose','ear','neck','lateral','body_center','hip','tail_base','tail_midpoint','tail_tip'\n]\n_bodypart_pos_map = {p: i for i, p in enumerate(BODYPART_ORDER)}\n\ndef _precompute_tracking_buffers(df: pd.DataFrame):\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    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    groups = df.groupby('video_frame', sort=True).indices\n    frame_to_rows = {int(k): v for k, v in groups.items()}\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    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        xs = (buf['x'][idx] * sx).astype(np.int32)\n        ys = (buf['y'][idx] * sy).astype(np.int32)\n        xs = np.clip(xs, 0, img_size - 1)\n        ys = np.clip(ys, 0, img_size - 1)\n        gr = buf['gray'][idx]\n        dot = max(1, int(CFG.DOT_SIZE))\n        if dot == 1:\n            img[ys, xs] = gr\n        else:\n            h, w = img.shape\n            for x, y, g in zip(xs, ys, gr):\n                x2 = min(x + dot, w)\n                y2 = min(y + dot, h)\n                img[y:y2, x:x2] = g\n    return torch.from_numpy(img).unsqueeze(0).to(dtype=torch.float32).div_(255.0)\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    else:\n        raise RuntimeError(\"Expected precomputed buffer dict.\")\n\ndef _rasterize_points_tensor(points, gray_values, img_size: int) -> torch.Tensor:\n    img = np.zeros((img_size, img_size), dtype=np.uint8)\n    if points is not None and len(points) > 0:\n        x = points[:, 0].astype(np.int32)\n        y = points[:, 1].astype(np.int32)\n        valid_indices = (x >= 0) & (x < img_size) & (y >= 0) & (y < img_size)\n        x, y = x[valid_indices], y[valid_indices]\n        valid_gray = gray_values[valid_indices]\n        dot = max(1, int(CFG.DOT_SIZE))\n        if dot == 1:\n            img[y, x] = valid_gray\n        else:\n            h, w = img.shape\n            for xi, yi, gi in zip(x, y, valid_gray):\n                x2 = min(xi + dot, w)\n                y2 = min(yi + dot, h)\n                img[yi:y2, xi:x2] = gi\n    return torch.from_numpy(img).unsqueeze(0).to(dtype=torch.float32).div_(255.0)\n\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\n\n# ===================================================================\n# 3. Behaviors_labeled -> Allowed Class Mask\n# ===================================================================\ndef _normalize_beh_label_to_classname(s: str) -> str | None:\n    if not isinstance(s, str) or not s:\n        return None\n    parts = [p.strip() for p in s.split(',')]\n    if len(parts) != 3:\n        return None\n    agent, target, action = parts[0], parts[1], parts[2].lower()\n    agent_num = agent.replace('mouse', '')\n    if target.lower() == 'self':\n        target_part = 'self'\n    else:\n        target_part = target.replace('mouse', '')\n    if not agent_num or not target_part or not action:\n        return None\n    return f\"{agent_num}_{target_part}_{action}\"\n\ndef build_behaviors_masks(meta_df: pd.DataFrame, class_to_idx: dict[int, int] | dict[str, int]) -> dict[str, torch.Tensor]:\n    n_classes = len(class_to_idx)\n    masks = {}\n    for _, row in meta_df.iterrows():\n        vid = str(row['video_id'])\n        bl = row.get('behaviors_labeled', None)\n        mask = np.zeros(n_classes, dtype=np.float32)\n        items = []\n        if isinstance(bl, str) and bl.strip():\n            try:\n                items = json.loads(bl)\n            except Exception:\n                try:\n                    items = ast.literal_eval(bl)\n                except Exception:\n                    items = []\n        for s in items or []:\n            cname = _normalize_beh_label_to_classname(str(s))\n            if cname is None:\n                continue\n            idx = class_to_idx.get(cname, None)\n            if idx is not None:\n                mask[idx] = 1.0\n        if mask.sum() == 0:\n            mask[:] = 1.0\n        masks[vid] = torch.from_numpy(mask)\n    return masks\n\n\n# ===================================================================\n# 4. Dataset\n# ===================================================================\nclass MABeFrameDataset(Dataset):\n    \"\"\"\n    2.5D: Stack frames (center +/- offsets) along the channel dimension.\n    Returns: (clip[C,H,W] (0..1), label[n_classes] (multi-hot), mask[n_classes] (0/1))\n    \"\"\"\n    def __init__(self, video_ids, meta_df, frame_labels, all_tracking_data,\n                 behaviors_masks: dict[str, torch.Tensor],\n                 augmentor=None, class_to_idx: dict = None, id_swap_prob: float = None):\n        self.video_ids = video_ids\n        self.meta_df = meta_df.set_index('video_id')\n        self.frame_labels = frame_labels\n        self.all_tracking_data = all_tracking_data\n        self.behaviors_masks = behaviors_masks\n        self.augmentor = augmentor\n        self.class_to_idx = class_to_idx\n        self.idx_to_class = {v: k for k, v in class_to_idx.items()} if class_to_idx is not None else None\n        self.id_swap_prob = CFG.ID_SWAP_PROB if id_swap_prob is None else float(id_swap_prob)\n\n        self._allones_mask = torch.ones(len(class_to_idx), dtype=torch.float32) if class_to_idx is not None else None\n\n        self.index_map = []\n        for video_id in self.video_ids:\n            total_frames = self.frame_labels[video_id].shape[0]\n            frames = list(range(0, total_frames, max(1, CFG.FRAME_SKIP)))\n            for fr in frames:\n                self.index_map.append((video_id, fr))\n\n        self._aug_vis_enabled = bool(CFG.SAVE_AUG_VIS_DURING_TRAIN)\n        self._aug_vis_saved = 0\n        self._epoch = 0\n        self._aug_vis_dir = (CFG.OUTPUT_DIR / \"aug_vis_epoch1\") if self._aug_vis_enabled else None\n        if self._aug_vis_enabled:\n            try:\n                self._aug_vis_dir.mkdir(parents=True, exist_ok=True)\n            except Exception:\n                pass\n\n    def set_epoch(self, epoch: int):\n        self._epoch = int(epoch)\n        if self._aug_vis_enabled and self._epoch == 0:\n            self._aug_vis_saved = 0\n\n    @staticmethod\n    def _sample_perm():\n        arr = np.random.permutation([1,2,3,4]).tolist()\n        return {i+1: arr[i] for i in range(4)}\n\n    @staticmethod\n    def _base_u8_for_id(mouse_id: int) -> int:\n        base01 = {1: 0.9, 2: 0.7, 3: 0.5, 4: 0.3}.get(int(mouse_id), 0.1)\n        return int(round(base01 * 255.0))\n\n    def _id_swap_clip(self, clip: torch.Tensor, perm: dict) -> torch.Tensor:\n        arr = (clip.cpu().numpy() * 255.0).astype(np.int16)\n        rng = int(round(0.1 * 255.0))\n        out = arr.copy()\n        for src_id in (1,2,3,4):\n            dst_id = perm.get(src_id, src_id)\n            src_base = self._base_u8_for_id(src_id)\n            dst_base = self._base_u8_for_id(dst_id)\n            mask = (arr <= src_base) & (arr >= (src_base - rng))\n            if not mask.any():\n                continue\n            delta = (src_base - arr).clip(0, rng)\n            new_vals = (dst_base - delta).clip(0, 255)\n            out = np.where(mask, new_vals, out)\n        out = np.clip(out, 0, 255).astype(np.uint8)\n        out = out.astype(np.float32) / 255.0\n        return torch.from_numpy(out)\n\n    def _remap_label(self, label_vec: torch.Tensor, perm: dict) -> torch.Tensor:\n        if self.class_to_idx is None or self.idx_to_class is None:\n            return label_vec\n        vec = label_vec.detach().cpu().numpy()\n        if vec.ndim != 1:\n            vec = vec.reshape(-1)\n        active = np.where(vec > 0.5)[0].tolist()\n        if not active:\n            return label_vec\n        new_vec = np.zeros_like(vec, dtype=np.float32)\n        for idx in active:\n            name = self.idx_to_class.get(int(idx))\n            if not name:\n                continue\n            try:\n                parts = name.split('_', 2)\n                a = int(parts[0]); t = parts[1]; action = parts[2]\n                a2 = perm.get(a, a)\n                if t == 'self':\n                    t2s = 'self'\n                else:\n                    t2 = perm.get(int(t), int(t))\n                    t2s = str(t2)\n                new_name = f\"{a2}_{t2s}_{action}\"\n                new_idx = self.class_to_idx.get(new_name, int(idx))\n                new_vec[new_idx] = 1.0\n            except Exception:\n                new_vec[idx] = 1.0\n        return torch.from_numpy(new_vec)\n\n    def __len__(self):\n        return len(self.index_map)\n\n    def __getitem__(self, idx):\n        video_id, frame_idx = self.index_map[idx]\n        tracking_buf = self.all_tracking_data[video_id]\n        video_meta = self.meta_df.loc[video_id]\n        width, height = int(video_meta['video_width_pix']), int(video_meta['video_height_pix'])\n        fps = _fps_or_default(video_meta.get('frames_per_second', 30.0))\n        offsets = _frame_offsets_for_fps(fps)\n\n        label = self.frame_labels[video_id][frame_idx]\n        beh_mask = self.behaviors_masks.get(video_id, self._allones_mask)\n\n        transform = None\n        if self.augmentor is not None:\n            transform = A.ReplayCompose([\n                A.OneOf([\n                    A.HorizontalFlip(p=1.0),\n                    A.VerticalFlip(p=1.0),\n                ], p=0.5),\n                A.RandomRotate90(p=0.5),\n                A.ShiftScaleRotate(\n                    shift_limit=CFG.AugShift,\n                    scale_limit=CFG.AugScale,\n                    rotate_limit=CFG.AugAngle,\n                    border_mode=cv2.BORDER_CONSTANT,\n                    value=0,\n                    p=0.5,\n                ),\n            ], keypoint_params=A.KeypointParams(format='xy', remove_invisible=False))\n\n        imgs = []\n        replay_data = None\n        sx = CFG.IMG_SIZE / max(int(width), 1)\n        sy = CFG.IMG_SIZE / max(int(height), 1)\n\n        for off in offsets:\n            f = frame_idx + off\n            row_indices = tracking_buf['frame_to_rows'].get(int(f), None)\n            if row_indices is not None and len(row_indices) > 0:\n                coords = np.vstack([\n                    tracking_buf['x'][row_indices] * sx,\n                    tracking_buf['y'][row_indices] * sy\n                ]).T\n                gray_vals = tracking_buf['gray'][row_indices]\n                if transform is not None:\n                    if replay_data is None:\n                        data = transform(image=np.zeros((CFG.IMG_SIZE, CFG.IMG_SIZE)), keypoints=coords)\n                        replay_data = data['replay']\n                        augmented_coords = np.array(data['keypoints'])\n                    else:\n                        data = A.ReplayCompose.replay(replay_data, image=np.zeros((CFG.IMG_SIZE, CFG.IMG_SIZE)), keypoints=coords)\n                        augmented_coords = np.array(data['keypoints'])\n                    t = _rasterize_points_tensor(augmented_coords, gray_vals, CFG.IMG_SIZE)\n                else:\n                    t = _rasterize_points_tensor(coords, gray_vals, CFG.IMG_SIZE)\n            else:\n                t = torch.zeros((1, CFG.IMG_SIZE, CFG.IMG_SIZE), dtype=torch.float32)\n            imgs.append(t)\n        clip = torch.cat(imgs, dim=0)\n\n        if self.augmentor is not None and np.random.rand() < self.id_swap_prob:\n            perm_used = self._sample_perm()\n            clip = self._id_swap_clip(clip, perm_used)\n            label = self._remap_label(label, perm_used)\n        return clip, label, beh_mask\n\n\nclass MABeValAllFramesDataset(Dataset):\n    \"\"\"\n    Returns: (clip, (video_id, frame_idx), beh_mask[n_classes])\n    \"\"\"\n    def __init__(self, video_ids, meta_df, frame_labels, all_tracking_data, behaviors_masks):\n        self.video_ids = video_ids\n        self.meta_df = meta_df.set_index('video_id')\n        self.frame_labels = frame_labels\n        self.all_tracking_data = all_tracking_data\n        self.behaviors_masks = behaviors_masks\n\n        any_vid = video_ids[0] if len(video_ids) else None\n        n_classes = frame_labels[any_vid].shape[1] if any_vid is not None else 1\n        self._allones_mask = torch.ones(n_classes, dtype=torch.float32)\n\n        self.index_map = []\n        for vid in self.video_ids:\n            total_frames = self.frame_labels[vid].shape[0]\n            for fr in range(total_frames):\n                self.index_map.append((vid, fr))\n\n    def __len__(self):\n        return len(self.index_map)\n\n    def __getitem__(self, idx):\n        video_id, frame_idx = self.index_map[idx]\n        tracking_df_full = self.all_tracking_data[video_id]\n        meta_row = self.meta_df.loc[video_id]\n        width, height = int(meta_row['video_width_pix']), int(meta_row['video_height_pix'])\n        fps = _fps_or_default(meta_row.get('frames_per_second', 30.0))\n        offsets = _frame_offsets_for_fps(fps)\n        clip = build_2p5d_clip(\n            tracking_df_full, center_frame_idx=frame_idx,\n            width=width, height=height, img_size=CFG.IMG_SIZE,\n            offsets=offsets\n        )\n        beh_mask = self.behaviors_masks.get(video_id, self._allones_mask)\n        return clip, (video_id, frame_idx), beh_mask\n\n\n# ===================================================================\n# 5. Model (3D Stem + 2D EfficientNet/MaxVit)\n# ===================================================================\nclass MABeEfficientNet(nn.Module):\n    def __init__(self, model_name, n_classes, pretrained=True, 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        self.in_chans = in_chans\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        self.drop3d_1 = nn.Dropout3d(p=CFG.DROP_3D_STEM) if CFG.DROP_3D_STEM > 0 else nn.Identity()\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),\n            groups=in_chans, bias=False\n        )\n        self.temporal_bn = nn.BatchNorm3d(in_chans)\n        self.temporal_act = nn.SiLU(inplace=True)\n        self.drop3d_2 = nn.Dropout3d(p=CFG.DROP_3D_STEM) if CFG.DROP_3D_STEM > 0 else nn.Identity()\n        self.drop2d = nn.Dropout2d(p=CFG.DROP_2D) if CFG.DROP_2D > 0 else nn.Identity()\n        self.backbone = timm.create_model(\n            model_name,\n            pretrained=pretrained,\n            in_chans=in_chans,\n            num_classes=n_classes,\n            drop_rate=CFG.DROP_DENSE,\n            drop_path_rate=CFG.DROP_PATH_RATE\n        )\n\n    def forward(self, x):\n        x = x.unsqueeze(1)  # [B, 1, C, H, W]\n        x = self.spatial_conv(x)\n        x = self.spatial_bn(x)\n        x = self.spatial_act(x)\n        x = self.drop3d_1(x)\n        x = self.temporal_conv(x)\n        x = self.temporal_bn(x)\n        x = self.temporal_act(x)\n        x = self.drop3d_2(x)\n        x = x.mean(dim=2)   # [B, C, H, W]\n        x = self.drop2d(x)\n        x = self.backbone(x)  # [B, n_classes]\n        return x\n\n\n# ===================================================================\n# 6. Prepare Labels\n# ===================================================================\ndef prepare_labels_and_classes(meta_df, annotation_path):\n    all_actions = set()\n    annotation_files = list(annotation_path.glob(\"*/*.parquet\"))\n    print(\"Collecting all possible actions...\")\n    for file_path in tqdm(annotation_files):\n        df = pd.read_parquet(file_path)\n        for _, row in df.iterrows():\n            if pd.notna(row['target_id']) and row['agent_id'] == row['target_id']:\n                target = 'self'\n            else:\n                target = row['target_id'] if pd.notna(row['target_id']) else 'self'\n            def _normalize_mouse(s):\n                s = str(s)\n                return s.replace('mouse', '') if s.startswith('mouse') else s\n            agent_num = _normalize_mouse(row['agent_id'])\n            target_part = 'self' if target == 'self' else _normalize_mouse(target)\n            action = str(row['action']).strip().lower()\n            all_actions.add(f\"{agent_num}_{target_part}_{action}\")\n    sorted_actions = sorted(list(all_actions))\n    class_to_idx = {cls_name: i for i, cls_name in enumerate(sorted_actions)}\n    actions_df = pd.DataFrame(sorted_actions, columns=['action_class'])\n    (CFG.OUTPUT_DIR / 'all_actions.csv').parent.mkdir(parents=True, exist_ok=True)\n    actions_df.to_csv(CFG.OUTPUT_DIR / 'all_actions.csv', index=False)\n    print(f\"✅ Found {len(sorted_actions)} unique action classes. Saved: {CFG.OUTPUT_DIR / 'all_actions.csv'}\")\n\n    print(\"Creating frame-wise labels...\")\n    frame_labels = {}\n    for file_path in tqdm(annotation_files):\n        video_id = file_path.stem\n        labels_df = pd.read_parquet(file_path)\n        video_meta = meta_df[meta_df['video_id'] == video_id]\n        if video_meta.empty: \n            continue\n        max_frame = int(labels_df[['start_frame','stop_frame']].max().max()) + 1\n        video_label_array = np.zeros((max_frame, len(class_to_idx)), dtype=np.float32)\n        for _, r in labels_df.iterrows():\n            if pd.notna(r['target_id']) and r['agent_id'] == r['target_id']:\n                target = 'self'\n            else:\n                target = r['target_id'] if pd.notna(r['target_id']) else 'self'\n            def _normalize_mouse(s):\n                s = str(s)\n                return s.replace('mouse', '') if s.startswith('mouse') else s\n            agent_num = _normalize_mouse(r['agent_id'])\n            target_part = 'self' if target == 'self' else _normalize_mouse(target)\n            action = str(r['action']).strip().lower()\n            class_name = f\"{agent_num}_{target_part}_{action}\"\n            if class_name in class_to_idx:\n                class_idx = class_to_idx[class_name]\n                s = int(r['start_frame']); e = int(r['stop_frame'])\n                s = max(0, s); e = max(s+1, e)\n                s = min(s, max_frame-1); e = min(e, max_frame)\n                video_label_array[s:e, class_idx] = 1.0\n        frame_labels[video_id] = torch.from_numpy(video_label_array)\n    return class_to_idx, frame_labels\n\n\n# ===================================================================\n# 7. Train & Validation Loop\n# ===================================================================\ndef train_loop(model, train_loader, criterion_none, device):\n    model.train()\n    total_loss = 0.0\n    total_weights = 0.0\n    pbar = tqdm(train_loader, desc=\"Training\", leave=False)\n    for batch in pbar:\n        clips, labels, masks = batch\n        clips = clips.to(device, non_blocking=True)\n        labels = labels.to(device, non_blocking=True)\n        masks = masks.to(device, non_blocking=True)\n        optimizer = train_loop.optimizer\n        optimizer.zero_grad()\n        logits = model(clips)\n        loss_elem = criterion_none(logits, labels)\n        loss_masked_sum = (loss_elem * masks).sum()\n        mask_count = masks.sum().clamp_min(1.0)\n        loss = loss_masked_sum / mask_count\n        loss.backward()\n        optimizer.step()\n        total_loss += float(loss_masked_sum.detach().cpu().item())\n        total_weights += float(mask_count.detach().cpu().item())\n        pbar.set_postfix({'loss': f'{(total_loss/total_weights):.4f}'})\n    return (total_loss / max(1.0, total_weights))\n\n@torch.no_grad()\ndef validate_loop(model, val_loader, criterion_none, device):\n    model.eval()\n    total_loss = 0.0\n    total_weights = 0.0\n    pbar = tqdm(val_loader, desc=\"Validating\", leave=False)\n    for batch in pbar:\n        clips, labels, masks = batch\n        clips = clips.to(device, non_blocking=True)\n        labels = labels.to(device, non_blocking=True)\n        masks = masks.to(device, non_blocking=True)\n        logits = model(clips)\n        loss_elem = criterion_none(logits, labels)\n        loss_masked_sum = (loss_elem * masks).sum()\n        mask_count = masks.sum().clamp_min(1.0)\n        total_loss += float(loss_masked_sum.detach().cpu().item())\n        total_weights += float(mask_count.detach().cpu().item())\n        pbar.set_postfix({'loss': f'{(total_loss/total_weights):.4f}'})\n    return (total_loss / max(1.0, total_weights))\n\n@torch.no_grad()\ndef evaluate_f1_on_loader(model, loader, device, threshold: float = None):\n    if threshold is None:\n        threshold = CFG.F1_GLOBAL_THRESHOLD\n    model.eval()\n    first_batch = True\n    TP_c = FP_c = FN_c = P_c = None\n    for clips, labels, masks in tqdm(loader, desc=\"Eval F1\", leave=False):\n        clips  = clips.to(device, non_blocking=True)\n        labels = labels.to(device, non_blocking=True)\n        masks  = masks.to(device, non_blocking=True)\n        logits = model(clips)\n        probs  = torch.sigmoid(logits)\n        y_true = (labels > 0.5) & (masks > 0.5)\n        y_pred = (probs >= threshold) & (masks > 0.5)\n        TP_b = (y_true & y_pred).sum(dim=0).detach().cpu().numpy()\n        FP_b = ((~y_true) & y_pred).sum(dim=0).detach().cpu().numpy()\n        FN_b = (y_true & (~y_pred)).sum(dim=0).detach().cpu().numpy()\n        P_b  = (y_true).sum(dim=0).detach().cpu().numpy()\n        if first_batch:\n            first_batch = False\n            TP_c = TP_b.astype(np.float64)\n            FP_c = FP_b.astype(np.float64)\n            FN_c = FN_b.astype(np.float64)\n            P_c  = P_b.astype(np.float64)\n        else:\n            TP_c += TP_b; FP_c += FP_b; FN_c += FN_b; P_c += P_b\n    TP = TP_c.sum(); FP = FP_c.sum(); FN = FN_c.sum()\n    f1_micro = (2*TP) / (2*TP + FP + FN + 1e-9)\n    valid = (P_c > 0)\n    f1_per_class = np.zeros_like(P_c, dtype=np.float64)\n    denom = (2*TP_c + FP_c + FN_c + 1e-9)\n    f1_per_class[valid] = (2*TP_c[valid]) / denom[valid]\n    f1_macro = f1_per_class[valid].mean() if valid.sum() > 0 else 0.0\n    return float(f1_micro), float(f1_macro)\n\n\n@torch.no_grad()\ndef pre_epoch_augmentation_visualization(train_dataset: MABeFrameDataset):\n    if not CFG.SAVE_AUG_VIS:\n        return\n    save_dir = CFG.OUTPUT_DIR / CFG.AUG_VIS_PRE_DIR\n    save_dir.mkdir(parents=True, exist_ok=True)\n    n = min(CFG.AUG_VIS_MAX, len(train_dataset.index_map))\n    if n <= 0:\n        return\n    sel_idx = np.random.choice(len(train_dataset.index_map), size=n, replace=False)\n    for k, idx in enumerate(tqdm(sel_idx, desc=\"Saving augmentation samples\")):\n        video_id, frame_idx = train_dataset.index_map[idx]\n        tracking_buf = train_dataset.all_tracking_data[video_id]\n        meta_row = train_dataset.meta_df.loc[video_id]\n        width, height = int(meta_row['video_width_pix']), int(meta_row['video_height_pix'])\n        fps = _fps_or_default(meta_row.get('frames_per_second', 30.0))\n        offsets = _frame_offsets_for_fps(fps)\n        if 0 in offsets:\n            ch = offsets.index(0)\n        else:\n            ch = int(np.argmin([abs(o) for o in offsets]))\n        clip_before = build_2p5d_clip(\n            tracking_buf, center_frame_idx=frame_idx,\n            width=width, height=height, img_size=CFG.IMG_SIZE,\n            offsets=offsets\n        ).clone()\n        clip_after, _, _ = train_dataset[idx]\n        b = (clip_before[ch].cpu().numpy() * 255.0).astype(np.uint8)\n        a = (clip_after[ch].cpu().numpy() * 255.0).astype(np.uint8)\n        comb = np.concatenate([b, a], axis=1)\n        img = Image.fromarray(comb, mode='L')\n        out_path = save_dir / f\"vid{video_id}_f{frame_idx:06d}_{k:03d}.png\"\n        try:\n            img.save(out_path)\n        except Exception:\n            pass\n\n\n# ===================================================================\n# 8. OOF Output (Masked + Unmasked)\n# ===================================================================\n@torch.no_grad()\ndef write_oof_predictions(\n    model, video_ids, meta_df, frame_labels, all_tracking_data,\n    behaviors_masks, fold_idx, out_dir: Path, device: str\n):\n    \"\"\"\n    Save both masked version (legacy) and unmasked version (full class)\n      - oof_probs_fold{fold}.npy            (Masked)\n      - oof_probs_unmasked_fold{fold}.npy   (Unmasked)\n    \"\"\"\n    model.eval()\n    meta_idx = meta_df.set_index(\"video_id\")\n    ds = MABeValAllFramesDataset(video_ids, meta_df, frame_labels, all_tracking_data, behaviors_masks)\n\n    g = torch.Generator(); g.manual_seed(CFG.SEED)\n    dl = DataLoader(\n        ds, batch_size=CFG.BATCH_SIZE, shuffle=False,\n        num_workers=CFG.NUM_WORKERS, pin_memory=CFG.PIN_MEMORY,\n        worker_init_fn=_seed_worker, generator=g,\n        persistent_workers=(CFG.NUM_WORKERS > 0),\n        prefetch_factor=(2 if CFG.NUM_WORKERS > 0 else None),\n    )\n\n    probs_masked_all = []\n    probs_unmasked_all = []\n    all_meta_info = []\n\n    pbar = tqdm(dl, desc=f\"OOF predict (fold {fold_idx})\", leave=False)\n    for clips, meta, masks in pbar:\n        clips = clips.to(device, non_blocking=True)\n        logits = model(clips)\n        probs = torch.sigmoid(logits)  # [B, C]\n\n        # ---- Unmasked (Raw) ----\n        probs_unmasked = probs.detach().cpu().numpy().astype(np.float16)\n        probs_unmasked_all.append(probs_unmasked)\n\n        # ---- Masked (Legacy) ----\n        masks = masks.to(device, non_blocking=True)\n        probs_masked = (probs * masks).detach().cpu().numpy().astype(np.float16)\n        probs_masked_all.append(probs_masked)\n\n        video_ids_batch, frames_batch = meta\n        for i in range(len(video_ids_batch)):\n            vid = str(video_ids_batch[i])\n            frame_idx = int(frames_batch[i].item())\n            lab = str(meta_idx.loc[vid, \"lab_id\"]) if \"lab_id\" in meta_idx.columns else \"\"\n            parquet_name = f\"{vid}.parquet\"\n            all_meta_info.append((lab, parquet_name, fold_idx, frame_idx))\n\n    # Save\n    probs_masked_all = np.concatenate(probs_masked_all, axis=0)\n    probs_unmasked_all = np.concatenate(probs_unmasked_all, axis=0)\n    meta_df_oof = pd.DataFrame(all_meta_info, columns=['lab', 'parquet', 'fold', 'frame'])\n\n    out_dir.mkdir(parents=True, exist_ok=True)\n    np.save(out_dir / f\"oof_probs_fold{fold_idx}.npy\", probs_masked_all)\n    np.save(out_dir / f\"oof_probs_unmasked_fold{fold_idx}.npy\", probs_unmasked_all)\n    meta_df_oof.to_parquet(out_dir / f\"oof_meta_fold{fold_idx}.parquet\", index=False)\n    print(f\"📝 OOF saved: masked= oof_probs_fold{fold_idx}.npy , unmasked= oof_probs_unmasked_fold{fold_idx}.npy\")\n\n\n# ===================================================================\n# 9. Threshold / Bout Optimization\n# ===================================================================\ndef compute_and_save_per_action_thresholds_from_oof(\n    oof_probs_path: Path,\n    oof_meta_path: Path,\n    frame_labels: dict,\n    behaviors_masks: dict,\n    class_to_idx: dict[str, int],\n    out_json_path: Path,\n    threshold_grid: np.ndarray | None = None,\n):\n    if threshold_grid is None:\n        threshold_grid = np.linspace(0.05, 0.95, 91)\n    probs = np.load(oof_probs_path).astype(np.float32)   # [N, C]  <- Input Masked Version\n    meta  = pd.read_parquet(oof_meta_path)               # columns: lab, parquet, fold, frame\n    if len(meta) != len(probs):\n        raise RuntimeError(\"OOF meta and probs length mismatch\")\n    num_classes = probs.shape[1]\n    T = len(threshold_grid)\n    TP = np.zeros((num_classes, T), dtype=np.float64)\n    FP = np.zeros((num_classes, T), dtype=np.float64)\n    P  = np.zeros(num_classes, dtype=np.float64)\n    seen = np.zeros(num_classes, dtype=np.int64)\n    meta = meta.reset_index(drop=True)\n    meta[\"video_id\"] = meta[\"parquet\"].astype(str).str.replace(\".parquet\", \"\", regex=False)\n    idx_to_class = {v: k for k, v in class_to_idx.items()}\n    for vid, idxs in tqdm(meta.groupby(\"video_id\").groups.items(), desc=\"OOF thr-opt per-action\"):\n        idxs = np.asarray(list(idxs), dtype=np.int64)\n        vid_probs = probs[idxs]\n        frames    = meta.loc[idxs, \"frame\"].astype(int).to_numpy()\n        lab_mat   = frame_labels[str(vid)].numpy()\n        mask_vec  = behaviors_masks[str(vid)].numpy().astype(np.float32)\n        valid = (frames >= 0) & (frames < lab_mat.shape[0])\n        if not np.all(valid):\n            frames = frames[valid]; vid_probs = vid_probs[valid]\n            if len(frames) == 0: continue\n        y_true = lab_mat[frames, :] * mask_vec\n        y_true_bool = y_true > 0.5\n        P += y_true.sum(axis=0).astype(np.float64)\n        seen += (mask_vec > 0).astype(np.int64)\n        for t_idx, thr in enumerate(threshold_grid):\n            pred_t = (vid_probs >= float(thr)) & (mask_vec > 0)\n            TP[:, t_idx] += (pred_t & y_true_bool).sum(axis=0).astype(np.float64)\n            FP[:, t_idx] += (pred_t & (~y_true_bool)).sum(axis=0).astype(np.float64)\n    FN = (P[:, None] - TP)\n    denom = (2 * TP + FP + FN + 1e-9)\n    F1 = (2 * TP) / denom\n    best_thr_map = {}\n    for c in range(num_classes):\n        cname = idx_to_class[c]\n        if (seen[c] == 0) or (P[c] == 0):\n            best_thr_map[cname] = float(CFG.F1_GLOBAL_THRESHOLD)\n            continue\n        j = int(np.argmax(F1[c]))\n        best_thr_map[cname] = float(threshold_grid[j])\n    TPg = TP.sum(axis=0); FPg = FP.sum(axis=0); Psum = P.sum()\n    FNg = (Psum - TPg)\n    F1g = (2 * TPg) / (2 * TPg + FPg + FNg + 1e-9)\n    best_global = float(threshold_grid[int(np.argmax(F1g))])\n    payload = {\n        \"per_action_thresholds\": best_thr_map,\n        \"global_best_threshold\": best_global,\n        \"grid\": list(map(float, threshold_grid)),\n        \"note\": \"Optimized on OOF full frames (validation videos).\",\n    }\n    out_json_path.parent.mkdir(parents=True, exist_ok=True)\n    with open(out_json_path, \"w\") as f:\n        json.dump(payload, f, indent=2, ensure_ascii=False)\n    print(f\"✅ OOF per-action thresholds saved to: {out_json_path}\")\n    print(f\"   global_best_threshold (micro-F1): {best_global:.4f}\")\n    return payload\n\n\n# === Bout Optimization (Legacy/Standard) ===\ndef _rle_binary(arr: np.ndarray):\n    if arr.size == 0:\n        return np.array([], dtype=arr.dtype), np.array([], dtype=int)\n    diffs = np.diff(arr.astype(np.int8))\n    idx = np.where(diffs != 0)[0] + 1\n    lengths = np.diff(np.r_[0, idx, arr.size])\n    values  = arr[np.r_[idx - 1, arr.size - 1]]\n    return values, lengths\n\ndef _fill_short_zeros(arr: np.ndarray, max_gap: int) -> np.ndarray:\n    if max_gap <= 0 or arr.sum() == 0:\n        return arr.copy()\n    v, l = _rle_binary(arr)\n    if len(v) == 0:\n        return arr.copy()\n    out = arr.copy()\n    pos = 0\n    for val, ln in zip(v, l):\n        if val == 0 and ln <= max_gap:\n            out[pos:pos+ln] = 1\n        pos += ln\n    return out\n\ndef _remove_short_ones(arr: np.ndarray, min_len: int) -> np.ndarray:\n    if min_len <= 1 or arr.sum() == 0:\n        return arr.copy()\n    v, l = _rle_binary(arr)\n    if len(v) == 0:\n        return arr.copy()\n    out = arr.copy()\n    pos = 0\n    for val, ln in zip(v, l):\n        if val == 1 and ln < min_len:\n            out[pos:pos+ln] = 0\n        pos += ln\n    return out\n\ndef _boutize_from_probs(probs: np.ndarray, thr: float, min_len: int, max_gap: int) -> np.ndarray:\n    b = (probs > thr).astype(np.uint8)\n    b = _fill_short_zeros(b, max_gap)\n    b = _remove_short_ones(b, min_len)\n    return b\n\ndef _sec_to_frames(seconds: float, fps: float, frame_step: int = 1) -> int:\n    return max(1, int(round(seconds * fps / max(1, frame_step))))\n\ndef optimize_bout_params_per_action(\n    oof_probs: np.ndarray,\n    oof_meta: pd.DataFrame,\n    frame_labels: dict,\n    behaviors_masks: dict,\n    class_to_idx: dict,\n    per_action_thresholds: dict,\n    meta_df_for_fps: pd.DataFrame,\n    *,\n    oof_frame_step: int = 1,\n    min_len_grid_sec: tuple = CFG.MIN_LEN_GRID_SEC,\n    max_gap_grid_sec: tuple = CFG.MAX_GAP_GRID_SEC,\n    global_thr: float = 0.30,\n    min_support_bouts: int = CFG.MIN_SUPPORT_BOUTS,\n    default_min_len_sec: float = CFG.DEFAULT_MIN_LEN_SEC,\n    default_max_gap_sec: float = CFG.DEFAULT_MAX_GAP_SEC,\n):\n    idx_to_class = {v: k for k, v in class_to_idx.items()}\n    n_classes = len(idx_to_class)\n    meta_fps = meta_df_for_fps.set_index(\"video_id\")[\"frames_per_second\"].to_dict()\n    oof_meta = oof_meta.reset_index(drop=True)\n    oof_meta[\"video_id\"] = oof_meta[\"parquet\"].astype(str).str.replace(\".parquet\", \"\", regex=False)\n    vid_groups = oof_meta.groupby(\"video_id\").groups\n    best_min_len_sec = {}\n    best_max_gap_sec = {}\n    best_scores      = {}\n    for ci in range(n_classes):\n        cname = idx_to_class[ci]\n        thr_c = float(per_action_thresholds.get(cname, global_thr))\n        support_bouts = 0\n        for vid, idxs in vid_groups.items():\n            y = frame_labels[str(vid)].numpy()[:, ci].astype(np.uint8)\n            v, l = _rle_binary(y)\n            support_bouts += int(((v == 1) & (l > 0)).sum())\n        if support_bouts < min_support_bouts:\n            best_min_len_sec[cname] = default_min_len_sec\n            best_max_gap_sec[cname] = default_max_gap_sec\n            best_scores[cname]      = None\n            continue\n        best_f1 = -1.0\n        best_pair = (default_min_len_sec, default_max_gap_sec)\n        for L_sec in min_len_grid_sec:\n            for G_sec in max_gap_grid_sec:\n                tp = fp = fn = 0\n                for vid, idxs in vid_groups.items():\n                    idxs = np.asarray(list(idxs), dtype=np.int64)\n                    fps_v = float(_fps_or_default(meta_fps.get(float(vid) if isinstance(list(meta_fps.keys())[0], float) else vid, 30.0)))\n                    min_len_frames = _sec_to_frames(L_sec, fps_v, oof_frame_step)\n                    max_gap_frames = _sec_to_frames(G_sec, fps_v, oof_frame_step)\n                    vid_probs = oof_probs[idxs, ci]\n                    frames    = oof_meta.loc[idxs, \"frame\"].astype(int).to_numpy()\n                    yall      = frame_labels[str(vid)].numpy()[:, ci].astype(np.uint8)\n                    mask_c    = int(behaviors_masks[str(vid)].numpy()[ci] > 0.0)\n                    if mask_c == 0:\n                        continue\n                    valid = (frames >= 0) & (frames < yall.shape[0])\n                    if not np.any(valid):\n                        continue\n                    frames    = frames[valid]\n                    vid_probs = vid_probs[valid]\n                    y_true    = yall[frames]\n                    y_pred = _boutize_from_probs(vid_probs, thr_c, min_len_frames, max_gap_frames)\n                    tp += int(((y_true == 1) & (y_pred == 1)).sum())\n                    fp += int(((y_true == 0) & (y_pred == 1)).sum())\n                    fn += int(((y_true == 1) & (y_pred == 0)).sum())\n                denom = (2*tp + fp + fn)\n                f1 = (2*tp / denom) if denom > 0 else 0.0\n                if (f1 > best_f1) or (f1 == best_f1 and (L_sec < best_pair[0] or (L_sec == best_pair[0] and G_sec < best_pair[1]))):\n                    best_f1 = f1; best_pair = (L_sec, G_sec)\n        best_min_len_sec[cname], best_max_gap_sec[cname] = best_pair\n        best_scores[cname] = best_f1\n    return best_min_len_sec, best_max_gap_sec, best_scores\n\ndef save_extended_json_and_csv(\n    out_json_path: Path,\n    base_payload: dict,\n    per_action_min_len_sec: dict,\n    per_action_max_gap_sec: dict,\n    defaults: dict,\n    meta: dict,\n    out_csv_path: Path | None = None\n):\n    payload = dict(base_payload) if base_payload else {}\n    payload[\"per_action_min_bout_sec\"] = per_action_min_len_sec\n    payload[\"per_action_max_gap_sec\"]  = per_action_max_gap_sec\n    payload[\"defaults\"] = defaults\n    payload[\"meta\"] = meta\n    out_json_path.parent.mkdir(parents=True, exist_ok=True)\n    with open(out_json_path, \"w\") as f:\n        json.dump(payload, f, indent=2, ensure_ascii=False)\n    if out_csv_path is not None:\n        rows = []\n        keys = sorted(per_action_min_len_sec.keys())\n        for k in keys:\n            rows.append({\n                \"action_class\": k,\n                \"min_bout_sec\": per_action_min_len_sec.get(k, None),\n                \"max_gap_sec\":  per_action_max_gap_sec.get(k, None),\n                \"threshold\":    payload.get(\"per_action_thresholds\", {}).get(k, None),\n            })\n        pd.DataFrame(rows).to_csv(out_csv_path, index=False)\n    print(f\"✅ Extended JSON saved: {out_json_path}\")\n    if out_csv_path:\n        print(f\"✅ Per-action bout params CSV saved: {out_csv_path}\")\n\n\n# ===================================================================\n# 10. Utility for generating full re-annotation (Pseudo-labeling)\n# ===================================================================\ndef _load_fold_params_and_aggregate(output_dir: Path, folds: list[int]):\n    \"\"\"Read per-action thresholds/bout-seconds for each fold and aggregate by median.\"\"\"\n    thr_list = []\n    minlen_map = defaultdict(list)\n    maxgap_map = defaultdict(list)\n    for f in folds:\n        thr_json = output_dir / f\"thr_per_action_oof_fold{f}.json\"\n        bout_csv = output_dir / f\"per_action_bout_params_fold{f}.csv\"\n        if thr_json.exists():\n            j = json.loads(Path(thr_json).read_text(encoding=\"utf-8\"))\n            thr_list.append(j.get(\"per_action_thresholds\", {}))\n        if bout_csv.exists():\n            df = pd.read_csv(bout_csv)\n            for _, r in df.iterrows():\n                cls = r[\"action_class\"]\n                if pd.notna(r[\"min_bout_sec\"]): minlen_map[cls].append(float(r[\"min_bout_sec\"]))\n                if pd.notna(r[\"max_gap_sec\"]):  maxgap_map[cls].append(float(r[\"max_gap_sec\"]))\n    all_classes = set()\n    for d in thr_list:\n        all_classes |= set(d.keys())\n    thr_agg = {}\n    for cls in all_classes:\n        vals = [d[cls] for d in thr_list if cls in d]\n        if len(vals) == 0:\n            continue\n        thr_agg[cls] = float(np.median(vals))\n    minlen_agg = {k: float(np.median(v)) for k, v in minlen_map.items() if len(v) > 0}\n    maxgap_agg = {k: float(np.median(v)) for k, v in maxgap_map.items() if len(v) > 0}\n    return thr_agg, minlen_agg, maxgap_agg\n\ndef _class_name_to_triplet(cls_name: str):\n    \"\"\"'1_self_sniff' -> ('mouse1','self','sniff') / '2_1_attack' -> ('mouse2','mouse1','attack')\"\"\"\n    a, t, act = cls_name.split('_', 2)\n    agent_id = f\"mouse{a}\"\n    target_id = \"self\" if t == \"self\" else f\"mouse{t}\"\n    return agent_id, target_id, act\n\ndef _segments_from_binary(arr: np.ndarray):\n    \"\"\"Return continuous segments [start, stop) from 0/1 array.\"\"\"\n    if arr.ndim != 1: arr = arr.reshape(-1)\n    n = len(arr)\n    if n == 0: return []\n    out = []\n    v_prev = 0\n    start = None\n    for i, v in enumerate(arr):\n        if v_prev == 0 and v == 1:\n            start = i\n        if v_prev == 1 and v == 0 and start is not None:\n            out.append((start, i))\n            start = None\n        v_prev = v\n    if v_prev == 1 and start is not None:\n        out.append((start, n))\n    return out\n\n@torch.no_grad()\ndef _infer_probs_for_videos_with_model(model, video_ids, meta_df, tracking_map, n_classes):\n    \"\"\"Return probability for specific videos using a single model. dict[video_id] = (frames, probs[C])\"\"\"\n    ds_vids = []\n    for vid in video_ids:\n        if vid in tracking_map:\n            ds_vids.append(vid)\n    meta_idx = meta_df.set_index('video_id')\n    results = {}\n    for vid in ds_vids:\n        width  = int(meta_idx.loc[vid, 'video_width_pix'])\n        height = int(meta_idx.loc[vid, 'video_height_pix'])\n        fps    = _fps_or_default(meta_idx.loc[vid, 'frames_per_second'])\n        offsets = _frame_offsets_for_fps(fps)\n        buf = tracking_map[vid]\n        frames = sorted(buf['frame_to_rows'].keys())\n        all_probs = []\n        B = CFG.BATCH_SIZE\n        for i in range(0, len(frames), B):\n            batch = frames[i:i+B]\n            clips = []\n            for f in batch:\n                clips.append(build_2p5d_clip(buf, f, width, height, CFG.IMG_SIZE, offsets))\n            clips = torch.stack(clips, dim=0).to(CFG.DEVICE)  # [B, C, H, W]\n            logits = model(clips)\n            probs  = torch.sigmoid(logits).detach().cpu().numpy()  # [B, C]\n            all_probs.append(probs)\n        probs_all = np.concatenate(all_probs, axis=0) if all_probs else np.zeros((0, n_classes), dtype=np.float32)\n        results[vid] = (np.array(frames, dtype=np.int32), probs_all.astype(np.float32))\n    return results\n\ndef _ensemble_probs_across_folds(\n    model_paths: list[Path], video_ids: list[str], meta_df: pd.DataFrame,\n    tracking_map: dict, n_classes: int, in_chans: int\n):\n    \"\"\"Load fold models and ensemble average probabilities.\"\"\"\n    sum_probs = {vid: None for vid in video_ids}\n    frame_idx_ref = {}\n    for mp in model_paths:\n        model = MABeEfficientNet(CFG.MODEL_NAME, n_classes=n_classes, pretrained=False, in_chans=in_chans).to(CFG.DEVICE)\n        state = torch.load(mp, map_location=CFG.DEVICE)\n        model.load_state_dict(state, strict=True)\n        model.eval()\n        res = _infer_probs_for_videos_with_model(model, video_ids, meta_df, tracking_map, n_classes)\n        for vid, (frames, probs) in res.items():\n            if sum_probs[vid] is None:\n                sum_probs[vid] = probs.copy()\n                frame_idx_ref[vid] = frames\n            else:\n                if not np.array_equal(frame_idx_ref[vid], frames):\n                    raise RuntimeError(f\"Frame index mismatch in ensemble for video {vid}\")\n                sum_probs[vid] += probs\n        del model; gc.collect()\n        if torch.cuda.is_available(): torch.cuda.empty_cache()\n    avg_probs = {}\n    for vid in video_ids:\n        if sum_probs[vid] is None:\n            continue\n        avg_probs[vid] = (frame_idx_ref[vid], sum_probs[vid] / max(1, len(model_paths)))\n    return avg_probs\n\ndef _write_annotation_parquet(out_path: Path, rows: list[dict]):\n    if len(rows) == 0:\n        df = pd.DataFrame(columns=[\"agent_id\",\"target_id\",\"action\",\"start_frame\",\"stop_frame\"])\n    else:\n        df = pd.DataFrame(rows)\n    out_path.parent.mkdir(parents=True, exist_ok=True)\n    df.to_parquet(out_path, index=False)\n\ndef _copy_existing_annotations_to_allann(src_ann_root: Path, dst_allann_root: Path):\n    files = list(src_ann_root.glob(\"*/*.parquet\"))\n    for p in tqdm(files, desc=\"Copy GT annotations\"):\n        lab = p.parent.name\n        vid = p.name\n        dst = dst_allann_root / lab / vid\n        dst.parent.mkdir(parents=True, exist_ok=True)\n        df = pd.read_parquet(p)\n        df.to_parquet(dst, index=False)\n\ndef _generate_annotations_for_unlabeled_labs(\n    meta_df: pd.DataFrame,\n    class_to_idx: dict,\n    all_tracking_files: list[Path],\n    gt_ann_root: Path,\n    out_allann_root: Path,\n    folds: list[int],\n    model_dir: Path\n):\n    \"\"\"Generate predicted annotations for unlabeled labs (those not in train_annotation).\"\"\"\n    labeled_videos = set([p.stem for p in gt_ann_root.glob(\"*/*.parquet\")])\n    lab_to_videos = defaultdict(list)\n    for fp in all_tracking_files:\n        lab = fp.parent.name\n        vid = fp.stem\n        lab_to_videos[lab].append(vid)\n    unlabeled_labs = [lab for lab, vids in lab_to_videos.items()\n                      if all((v not in labeled_videos) for v in vids)]\n    for h in CFG.UNLABELED_LABS_HINT:\n        if h not in unlabeled_labs and h in lab_to_videos:\n            unlabeled_labs.append(h)\n    if len(unlabeled_labs) == 0:\n        print(\"No unlabeled labs detected. Skipping auto-annotation.\")\n        return\n\n    print(f\"Unlabeled labs to annotate: {unlabeled_labs}\")\n\n    meta_df = meta_df.copy()\n    # Fix: Keep lab_id as string to avoid issues, only convert numeric columns.\n    for col in ['frames_per_second','video_width_pix','video_height_pix']:\n        if col in meta_df.columns:\n            meta_df[col] = pd.to_numeric(meta_df[col], errors='coerce')\n    meta_df['video_id'] = meta_df['video_id'].astype(str)\n    meta_df['lab_id']   = meta_df['lab_id'].astype(str)\n\n    tracking_map = {}\n    for lab in unlabeled_labs:\n        for fp in [p for p in all_tracking_files if p.parent.name == lab]:\n            vid = fp.stem\n            try:\n                df = pd.read_parquet(fp)\n            except Exception:\n                continue\n            df = drop_body_parts(df)\n            tracking_map[vid] = _precompute_tracking_buffers(df)\n\n    in_chans = (len(CFG.OFFSETS_SEC) if CFG.USE_OFFSETS_SEC else len(CFG.OFFSETS))\n    n_classes = len(class_to_idx)\n    model_paths = [model_dir / f\"model_fold{f}_best.pth\" for f in folds if (model_dir / f\"model_fold{f}_best.pth\").exists()]\n    if len(model_paths) == 0:\n        print(\"[WARN] No fold model weights found for ensemble.\")\n        return\n    target_videos = sorted(tracking_map.keys())\n    print(f\"Ensemble inference for {len(target_videos)} videos ...\")\n    ens = _ensemble_probs_across_folds(model_paths, target_videos, meta_df, tracking_map, n_classes, in_chans)\n\n    thr_agg, minlen_agg, maxgap_agg = _load_fold_params_and_aggregate(model_dir, folds)\n    idx_to_class = {v: k for k, v in class_to_idx.items()}\n    meta_idx = meta_df.set_index('video_id')\n\n    for vid, (frames, probs) in tqdm(ens.items(), desc=\"Write predicted annotations\"):\n        try:\n            fps = _fps_or_default(meta_idx.loc[vid, 'frames_per_second'])\n        except Exception:\n            fps = 30.0\n        rows = []\n        for ci in range(n_classes):\n            cname = idx_to_class[ci]\n            thr   = float(thr_agg.get(cname, CFG.F1_GLOBAL_THRESHOLD))\n            L_sec = float(minlen_agg.get(cname, CFG.DEFAULT_MIN_LEN_SEC))\n            G_sec = float(maxgap_agg.get(cname, CFG.DEFAULT_MAX_GAP_SEC))\n            min_len = _sec_to_frames(L_sec, fps, frame_step=1)\n            max_gap = _sec_to_frames(G_sec, fps, frame_step=1)\n            yb = _boutize_from_probs(probs[:, ci], thr, min_len, max_gap)\n            for s, e in _segments_from_binary(yb):\n                agent_id, target_id, action = _class_name_to_triplet(cname)\n                rows.append(dict(\n                    agent_id=agent_id,\n                    target_id=target_id,\n                    action=action,\n                    start_frame=int(frames[s]),\n                    stop_frame=int(frames[e]),\n                ))\n        out_path = out_allann_root / meta_idx.loc[vid, 'lab_id'] / f\"{vid}.parquet\"\n        _write_annotation_parquet(out_path, rows)\n\n\n# ===================================================================\n# 11. Main\n# ===================================================================\ndef main():\n    seed_everything(CFG.SEED)\n    CFG.OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n    print(\"Device:\", CFG.DEVICE)\n    print(\"Output dir:\", CFG.OUTPUT_DIR)\n\n    base_path = CFG.BASE_PATH\n\n    # ==== Fix: Do not convert lab_id to numeric. Only fps/width/height. ====\n    meta_df = pd.read_csv(base_path / \"train.csv\", dtype=str)\n    for col in ['frames_per_second','video_width_pix','video_height_pix']:\n        if col in meta_df.columns:\n            meta_df[col] = pd.to_numeric(meta_df[col], errors='coerce')\n    meta_df['video_id'] = meta_df['video_id'].astype(str)\n    meta_df['lab_id']   = meta_df['lab_id'].astype(str)\n    # =====================================================================\n\n    class_to_idx, frame_labels = prepare_labels_and_classes(meta_df, base_path / \"train_annotation\")\n    n_classes = len(class_to_idx)\n    with open(CFG.OUTPUT_DIR / 'class_to_idx.json', 'w', encoding=\"utf-8\") as f:\n        json.dump(class_to_idx, f, indent=4, ensure_ascii=False)\n    print(f\"✅ class_to_idx.json saved to {CFG.OUTPUT_DIR}\")\n\n    behaviors_masks = build_behaviors_masks(meta_df, class_to_idx)\n\n    labeled_video_ids = set(frame_labels.keys())\n    all_tracking_data = {}\n    all_tracking_files = list((base_path / \"train_tracking\").glob(\"*/*.parquet\"))\n    print(\"\\nLoading tracking data (labeled videos) ...\")\n    for file_path in tqdm(all_tracking_files):\n        vid = file_path.stem\n        if vid in labeled_video_ids:\n            df = pd.read_parquet(file_path)\n            df = drop_body_parts(df)\n            all_tracking_data[vid] = _precompute_tracking_buffers(df)\n\n    # --- 5-Fold split (Split videos evenly per lab) ---\n    lab_groups = meta_df.groupby('lab_id')['video_id'].apply(list).to_dict()\n    folds = [[] for _ in range(CFG.N_SPLITS)]\n    start_fold = 0\n    for lab, vids in sorted(lab_groups.items()):\n        vids_lab = [str(v) for v in vids if str(v) in labeled_video_ids]\n        vids_lab = sorted(set(vids_lab))\n        if not vids_lab:\n            continue\n        for j, vid in enumerate(vids_lab):\n            folds[(start_fold + j) % CFG.N_SPLITS].append(vid)\n        start_fold = (start_fold + len(vids_lab)) % CFG.N_SPLITS\n    total_assigned = sum(len(f) for f in folds)\n    unique_assigned = len(set(v for f in folds for v in f))\n    print(f\"[Split] total videos assigned: {total_assigned}, unique: {unique_assigned}\")\n    print(f\"[Split] fold sizes: {[len(f) for f in folds]}\")\n\n    loader_common = _make_loader_common()\n    best_model_path = None\n    EFFECTIVE_IN_CHANS = (len(CFG.OFFSETS_SEC) if CFG.USE_OFFSETS_SEC else len(CFG.OFFSETS))\n    criterion_none = nn.BCEWithLogitsLoss(reduction='none')\n\n    # ====== Train each fold ======\n    for current_fold in range(CFG.N_SPLITS):\n        if current_fold not in CFG.FOLDS_TO_RUN:\n            continue\n        print(f\"\\n{'='*20} FOLD {current_fold} {'='*20}\")\n        start_time = time.time()\n        val_ids   = folds[current_fold]\n        train_ids = [v for i, fold in enumerate(folds) if i != current_fold for v in fold]\n        print(f\"Train videos: {len(train_ids)}, Val videos: {len(val_ids)}\")\n\n        train_dataset = MABeFrameDataset(\n            train_ids, meta_df, frame_labels, all_tracking_data,\n            behaviors_masks=behaviors_masks,\n            augmentor=True, class_to_idx=class_to_idx, id_swap_prob=CFG.ID_SWAP_PROB\n        )\n        val_dataset = MABeFrameDataset(\n            val_ids, meta_df, frame_labels, all_tracking_data,\n            behaviors_masks=behaviors_masks,\n            augmentor=None, class_to_idx=class_to_idx, id_swap_prob=0.0\n        )\n        pre_epoch_augmentation_visualization(train_dataset)\n\n        train_loader = DataLoader(train_dataset, shuffle=True,  **loader_common)\n        val_loader   = DataLoader(val_dataset,   shuffle=False, **loader_common)\n\n        model = MABeEfficientNet(CFG.MODEL_NAME, n_classes=n_classes, pretrained=True, in_chans=EFFECTIVE_IN_CHANS).to(CFG.DEVICE)\n        optimizer = torch.optim.AdamW(model.parameters(), lr=CFG.LR, weight_decay=CFG.WEIGHT_DECAY)\n        scheduler = CosineAnnealingLR(optimizer, T_max=CFG.EPOCHS, eta_min=CFG.ETA_MIN)\n        train_loop.optimizer = optimizer\n\n        history = []\n        best_val_loss = float('inf')\n        for epoch in range(CFG.EPOCHS):\n            try: train_dataset.set_epoch(epoch)\n            except Exception: pass\n            print(f\"--- Epoch {epoch+1}/{CFG.EPOCHS} ---\")\n            train_loss = train_loop(model, train_loader, criterion_none, CFG.DEVICE)\n            val_loss   = validate_loop(model, val_loader,   criterion_none, CFG.DEVICE)\n            f1_micro, f1_macro = evaluate_f1_on_loader(model, val_loader, CFG.DEVICE, CFG.F1_GLOBAL_THRESHOLD)\n            current_lr = float(optimizer.param_groups[0][\"lr\"])\n            print(f\"Epoch {epoch+1} | LR: {current_lr:.2e} | Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f} | F1-micro: {f1_micro:.4f}, F1-macro: {f1_macro:.4f}\")\n            history.append(dict(epoch=epoch+1, lr=current_lr, train_loss=train_loss, val_loss=val_loss, val_f1_micro=f1_micro, val_f1_macro=f1_macro))\n            scheduler.step()\n            if val_loss < best_val_loss:\n                best_val_loss = val_loss\n                best_model_path = CFG.OUTPUT_DIR / f\"model_fold{current_fold}_best.pth\"\n                torch.save(model.state_dict(), best_model_path)\n                print(f\"💾 Saved best(ValLoss) to: {best_model_path}\")\n\n        # Save curves\n        def _save_curve(history, name):\n            if not history: return\n            df = pd.DataFrame(history)\n            fig = plt.figure(figsize=(6,4))\n            if name == \"loss\":\n                plt.plot(df[\"epoch\"], df[\"train_loss\"], label=\"train\")\n                plt.plot(df[\"epoch\"], df[\"val_loss\"],   label=\"valid\")\n                plt.ylabel(\"Loss\")\n            else:\n                if \"val_f1_micro\" in df.columns: plt.plot(df[\"epoch\"], df[\"val_f1_micro\"], label=\"F1-micro\")\n                if \"val_f1_macro\" in df.columns: plt.plot(df[\"epoch\"], df[\"val_f1_macro\"], label=\"F1-macro\")\n                plt.ylabel(\"F1\")\n                plt.title(f\"F1 Curve @ thr={CFG.F1_GLOBAL_THRESHOLD}\")\n            plt.xlabel(\"Epoch\"); plt.legend(); plt.tight_layout()\n            plt.savefig(CFG.OUTPUT_DIR / f\"{name}_curve_fold{current_fold}.png\", dpi=150); plt.close(fig)\n            if name == \"loss\": df.to_csv(CFG.OUTPUT_DIR / f\"loss_history_fold{current_fold}.csv\", index=False)\n        _save_curve(history, \"loss\"); _save_curve(history, \"f1\")\n\n        # OOF (Masked + Unmasked)\n        write_oof_predictions(\n            model=model,\n            video_ids=val_ids,\n            meta_df=meta_df,\n            frame_labels=frame_labels,\n            all_tracking_data=all_tracking_data,\n            behaviors_masks=behaviors_masks,\n            fold_idx=current_fold,\n            out_dir=CFG.OUTPUT_DIR,\n            device=CFG.DEVICE\n        )\n\n        # Threshold Optimization (on Masked version)\n        oof_probs_path = CFG.OUTPUT_DIR / f\"oof_probs_fold{current_fold}.npy\"\n        oof_meta_path  = CFG.OUTPUT_DIR / f\"oof_meta_fold{current_fold}.parquet\"\n        thr_json_path  = CFG.OUTPUT_DIR / f\"thr_per_action_oof_fold{current_fold}.json\"\n        thr_payload = compute_and_save_per_action_thresholds_from_oof(\n            oof_probs_path=oof_probs_path,\n            oof_meta_path=oof_meta_path,\n            frame_labels=frame_labels,\n            behaviors_masks=behaviors_masks,\n            class_to_idx=class_to_idx,\n            out_json_path=thr_json_path,\n            threshold_grid=np.linspace(0.05, 0.95, 91)\n        )\n        # Bout Optimization\n        oof_probs = np.load(oof_probs_path).astype(np.float32)\n        oof_meta  = pd.read_parquet(oof_meta_path)\n        per_action_thr = thr_payload.get(\"per_action_thresholds\", {})\n        global_thr     = float(thr_payload.get(\"global_best_threshold\", CFG.F1_GLOBAL_THRESHOLD))\n        per_action_min_len_sec, per_action_max_gap_sec, best_scores = optimize_bout_params_per_action(\n            oof_probs=oof_probs,\n            oof_meta=oof_meta,\n            frame_labels=frame_labels,\n            behaviors_masks=behaviors_masks,\n            class_to_idx=class_to_idx,\n            per_action_thresholds=per_action_thr,\n            meta_df_for_fps=meta_df,\n            oof_frame_step=1,\n            min_len_grid_sec=CFG.MIN_LEN_GRID_SEC,\n            max_gap_grid_sec=CFG.MAX_GAP_GRID_SEC,\n            global_thr=global_thr,\n            min_support_bouts=CFG.MIN_SUPPORT_BOUTS,\n            default_min_len_sec=CFG.DEFAULT_MIN_LEN_SEC,\n            default_max_gap_sec=CFG.DEFAULT_MAX_GAP_SEC,\n        )\n        save_extended_json_and_csv(\n            out_json_path=thr_json_path,\n            base_payload=thr_payload,\n            per_action_min_len_sec=per_action_min_len_sec,\n            per_action_max_gap_sec=per_action_max_gap_sec,\n            defaults={\"min_bout_sec\": CFG.DEFAULT_MIN_LEN_SEC, \"max_gap_sec\": CFG.DEFAULT_MAX_GAP_SEC},\n            meta={\n                \"oof_frame_step\": 1,\n                \"grid_min_len_sec\": list(map(float, CFG.MIN_LEN_GRID_SEC)),\n                \"grid_max_gap_sec\": list(map(float, CFG.MAX_GAP_GRID_SEC)),\n            },\n            out_csv_path=CFG.OUTPUT_DIR / f\"per_action_bout_params_fold{current_fold}.csv\"\n        )\n\n        print(f\"Fold {current_fold} finished in {((time.time() - start_time) / 60):.2f} minutes.\")\n        del train_dataset, val_dataset, train_loader, val_loader, model\n        gc.collect()\n        if torch.cuda.is_available(): torch.cuda.empty_cache()\n\n    # ====== Re-annotation / Pseudo-labeling (Step 2) ======\n    ran_all_folds = sorted(CFG.FOLDS_TO_RUN) == list(range(CFG.N_SPLITS))\n    if ran_all_folds:\n        print(\"\\n[Re-Annotation] Start building train_annotation_Allannotation ...\")\n        allann_root = CFG.OUTPUT_DIR / CFG.ALLANN_ROOT_NAME\n        gt_ann_root = base_path / \"train_annotation\"\n        # 1) Copy existing GT annotations\n        if not CFG.RELABEL_LABELED_VIDEOS:\n            _copy_existing_annotations_to_allann(gt_ann_root, allann_root)\n        else:\n            _copy_existing_annotations_to_allann(gt_ann_root, allann_root)\n            print(\"[INFO] RELABEL_LABELED_VIDEOS=True is not fully implemented; copied GT as-is.\")\n\n        # 2) Generate annotations for unlabeled labs\n        _generate_annotations_for_unlabeled_labs(\n            meta_df=meta_df,\n            class_to_idx=class_to_idx,\n            all_tracking_files=all_tracking_files,\n            gt_ann_root=gt_ann_root,\n            out_allann_root=allann_root,\n            folds=CFG.ENSEMBLE_FOLDS,\n            model_dir=CFG.OUTPUT_DIR\n        )\n        print(f\"\\n✅ All annotations prepared at: {allann_root}\")\n        print(\"   -> Use this path instead of 'train_annotation' for next training.\")\n    else:\n        print(\"\\n[Re-Annotation] Skipped (run all 5 folds to enable).\")\n\n    print(f\"\\n🏁 Run finished.\")\n\n\nif __name__ == '__main__':\n    main()","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}]}