{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"},{"sourceId":13234212,"sourceType":"datasetVersion","datasetId":8387028}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Cell 1: Install required packages\n!pip install -q timm==0.9.12\n\n# ===================================================================\n# MABe: 2.5D CNN - ANTI-OVERFITTING VERSION with F-Score Tracking\n# ===================================================================\n\nimport os, gc, re, json, warnings\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nimport pandas as pd\nimport numpy as np\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn.parallel import DataParallel\nimport timm\n\nwarnings.filterwarnings('ignore')\n\n# ===================================================================\n# 1. Configuration with Regularization\n# ===================================================================\nclass CFG:\n    DATA_DIR = Path(\"/kaggle/input/MABe-mouse-behavior-detection\")\n    OUTPUT_DIR = Path(\"/kaggle/working\")\n    \n    # Model architecture\n    MODEL_NAME = \"efficientnet_b0\"\n    IMG_SIZE = 64\n    TEMPORAL_OFFSETS = [-10, -5, -2, 0, 2, 5, 10]\n    IN_CHANNELS = len(TEMPORAL_OFFSETS)\n    \n    # Regularization\n    DROPOUT = 0.3\n    LABEL_SMOOTHING = 0.1\n    AUGMENT_DATA = True\n    \n    # Training hyperparameters\n    BATCH_SIZE = 48\n    NUM_EPOCHS = 10\n    LEARNING_RATE = 5e-4\n    WEIGHT_DECAY = 1e-4\n    PATIENCE = 3\n    \n    # Data sampling - DOUBLED!\n    FRAME_SAMPLING_RATE = 5\n    MAX_VIDEOS_TRAIN = 100  # 2x increase\n    MAX_VIDEOS_VAL = 20     # 2x increase\n    \n    # Inference\n    INFERENCE_BATCH_SIZE = 96\n    INFERENCE_STRIDE = 10\n    PREDICTION_THRESHOLD = 0.28  # Lowered from 0.45\n    MIN_EVENT_DURATION = 3\n    \n    # System\n    USE_MULTI_GPU = torch.cuda.device_count() > 1\n    NUM_WORKERS = 0\n    POS_WEIGHT_SCALE = 2.0\n    MIXED_PRECISION = True\n    GRADIENT_CLIP = 1.0\n    \n    SEED = 42\n\nprint(f\"🛡️ ANTI-OVERFITTING VERSION with F-Score Tracking 🛡️\")\nprint(f\"PyTorch: {torch.__version__}\")\nprint(f\"CUDA: {torch.cuda.is_available()}\")\nprint(f\"GPUs: {torch.cuda.device_count()}\")\nprint(f\"\\nSettings:\")\nprint(f\"  - Videos: {CFG.MAX_VIDEOS_TRAIN} train, {CFG.MAX_VIDEOS_VAL} val\")\nprint(f\"  - Image size: {CFG.IMG_SIZE}x{CFG.IMG_SIZE}\")\nprint(f\"  - Temporal frames: {len(CFG.TEMPORAL_OFFSETS)} (offsets: {CFG.TEMPORAL_OFFSETS})\")\nprint(f\"  - Frame sampling: every {CFG.FRAME_SAMPLING_RATE}th frame\")\nprint(f\"  - Epochs: {CFG.NUM_EPOCHS} (with early stopping)\")\nprint(f\"  - Batch size: {CFG.BATCH_SIZE}\")\nprint(f\"  - Dropout: {CFG.DROPOUT}\")\nprint(f\"  - Label smoothing: {CFG.LABEL_SMOOTHING}\")\nprint(f\"  - Data augmentation: {CFG.AUGMENT_DATA}\")\nprint(f\"  - Prediction threshold: {CFG.PREDICTION_THRESHOLD}\")\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ntorch.manual_seed(CFG.SEED)\nnp.random.seed(CFG.SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed_all(CFG.SEED)\n\n# ===================================================================\n# 2. Data Functions\n# ===================================================================\n\ndef drop_body_parts(df: pd.DataFrame) -> pd.DataFrame:\n    df = df[~df['bodypart'].str.contains('headpiece', na=False)]\n    df['bodypart'] = df['bodypart'].str.replace('_left', '', regex=False).str.replace('_right', '', regex=False)\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 _precompute_tracking_buffers(df: pd.DataFrame):\n    key = df['bodypart'].str.extract(r'^(nose|ear|neck|lateral|body_center|hip|tail_base|tail_midpoint|tail_tip)')[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,\n                           augment: bool = False, shift_x: float = 0, shift_y: float = 0) -> 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        \n        x = (buf['x'][idx] + shift_x) * sx\n        y = (buf['y'][idx] + shift_y) * sy\n        \n        x = np.clip(x.astype(np.int32), 0, img_size - 1)\n        y = np.clip(y.astype(np.int32), 0, img_size - 1)\n        img[y, x] = buf['gray'][idx]\n    return torch.from_numpy(img).unsqueeze(0).to(dtype=torch.float32).div_(255.0)\n\ndef build_2p5d_clip(tracking_buf: dict, center_frame: int, width: int, height: int, \n                    img_size: int, offsets, augment: bool = False) -> torch.Tensor:\n    if augment and np.random.rand() < 0.5:\n        shift_x = np.random.uniform(-5, 5)\n        shift_y = np.random.uniform(-5, 5)\n    else:\n        shift_x, shift_y = 0, 0\n    \n    imgs = []\n    for off in offsets:\n        f = center_frame + off\n        t = _rasterize_frame_tensor(tracking_buf, f, width, height, img_size, \n                                    augment, shift_x, shift_y)\n        imgs.append(t)\n    \n    clip = torch.cat(imgs, dim=0)\n    \n    if augment and np.random.rand() < 0.3:\n        clip = torch.flip(clip, dims=[2])\n    \n    return clip\n\ndef find_annotated_videos(annotation_dir: Path):\n    video_to_ann_path = {}\n    for lab_dir in annotation_dir.iterdir():\n        if not lab_dir.is_dir():\n            continue\n        for ann_file in lab_dir.glob(\"*.parquet\"):\n            try:\n                video_to_ann_path[int(ann_file.stem)] = ann_file\n            except ValueError:\n                continue\n    print(f\"Found {len(video_to_ann_path)} annotated videos\")\n    return video_to_ann_path\n\ndef create_class_vocabulary(train_meta_df: pd.DataFrame, video_to_ann_path: dict, max_videos: int = None):\n    all_classes = set()\n    video_annotations = {}\n    \n    annotated_video_ids = list(video_to_ann_path.keys())\n    if max_videos:\n        annotated_video_ids = annotated_video_ids[:max_videos]\n    \n    print(f\"Processing {len(annotated_video_ids)} videos...\")\n    for video_id in annotated_video_ids:\n        try:\n            ann_df = pd.read_parquet(video_to_ann_path[video_id])\n            if len(ann_df) == 0:\n                continue\n            video_annotations[video_id] = ann_df\n            \n            for _, ann_row in ann_df.iterrows():\n                agent = str(ann_row['agent_id']).lower()\n                target = str(ann_row['target_id']).lower()\n                action = str(ann_row['action']).lower()\n                all_classes.add(f\"{agent}_{target}_{action}\")\n        except:\n            continue\n    \n    all_classes = sorted(list(all_classes))\n    class_to_idx = {c: i for i, c in enumerate(all_classes)}\n    idx_to_class = {i: c for c, i in class_to_idx.items()}\n    \n    print(f\"Found {len(all_classes)} classes from {len(video_annotations)} videos\")\n    return class_to_idx, idx_to_class, video_annotations\n\n# ===================================================================\n# 3. Datasets with Augmentation\n# ===================================================================\n\nclass MABeTrainingDataset(Dataset):\n    def __init__(self, video_ids, train_meta_df, video_annotations, tracking_dir, \n                 class_to_idx, img_size=96, offsets=None, frame_sampling_rate=10, augment=False):\n        self.video_ids = video_ids\n        self.train_meta_df = train_meta_df.set_index('video_id')\n        self.video_annotations = video_annotations\n        self.tracking_dir = tracking_dir\n        self.class_to_idx = class_to_idx\n        self.img_size = img_size\n        self.offsets = offsets or CFG.TEMPORAL_OFFSETS\n        self.frame_sampling_rate = frame_sampling_rate\n        self.augment = augment\n        self.tracking_cache = {}\n        self.samples = []\n        self._build_samples()\n        \n    def _build_samples(self):\n        for video_id in self.video_ids:\n            if video_id not in self.video_annotations:\n                continue\n            ann_df = self.video_annotations[video_id]\n            if len(ann_df) == 0:\n                continue\n            \n            buffer = max(abs(o) for o in self.offsets)\n            min_frame = max(0, int(ann_df['start_frame'].min()) - buffer)\n            max_frame = int(ann_df['stop_frame'].max()) + buffer\n            \n            for frame in range(min_frame, max_frame + 1, self.frame_sampling_rate):\n                self.samples.append((video_id, frame))\n        \n        print(f\"  Created {len(self.samples):,} samples (augment={self.augment})\")\n    \n    def _load_tracking(self, video_id):\n        if video_id in self.tracking_cache:\n            return self.tracking_cache[video_id]\n        \n        row = self.train_meta_df.loc[video_id]\n        tracking_path = self.tracking_dir / str(row['lab_id']) / f\"{video_id}.parquet\"\n        if not tracking_path.exists():\n            tracking_path = self.tracking_dir / f\"{video_id}.parquet\"\n        \n        tracking_df = pd.read_parquet(tracking_path)\n        tracking_df = drop_body_parts(tracking_df)\n        tracking_buf = _precompute_tracking_buffers(tracking_df)\n        \n        self.tracking_cache[video_id] = {\n            'buf': tracking_buf,\n            'width': int(row.get('video_width_pix', 640)),\n            'height': int(row.get('video_height_pix', 480))\n        }\n        return self.tracking_cache[video_id]\n    \n    def __len__(self):\n        return len(self.samples)\n    \n    def __getitem__(self, idx):\n        video_id, center_frame = self.samples[idx]\n        tracking_data = self._load_tracking(video_id)\n        \n        clip = build_2p5d_clip(tracking_data['buf'], center_frame, \n                               tracking_data['width'], tracking_data['height'],\n                               self.img_size, self.offsets, self.augment)\n        \n        label = torch.zeros(len(self.class_to_idx), dtype=torch.float32)\n        ann_df = self.video_annotations[video_id]\n        active_anns = ann_df[(ann_df['start_frame'] <= center_frame) & (ann_df['stop_frame'] >= center_frame)]\n        \n        for _, ann_row in active_anns.iterrows():\n            class_label = f\"{ann_row['agent_id']}_{ann_row['target_id']}_{ann_row['action']}\".lower()\n            if class_label in self.class_to_idx:\n                label[self.class_to_idx[class_label]] = 1.0\n        \n        if self.augment and CFG.LABEL_SMOOTHING > 0:\n            label = label * (1 - CFG.LABEL_SMOOTHING) + CFG.LABEL_SMOOTHING / 2\n        \n        return clip, label\n\nclass MABeInferenceDataset(Dataset):\n    def __init__(self, video_info, tracking_dir, img_size=96, offsets=None, stride=5):\n        self.video_id = video_info['video_id']\n        self.img_size = img_size\n        self.offsets = offsets or CFG.TEMPORAL_OFFSETS\n        \n        tracking_path = tracking_dir / video_info['lab_id'] / f\"{self.video_id}.parquet\"\n        if not tracking_path.exists():\n            tracking_path = tracking_dir / f\"{self.video_id}.parquet\"\n        \n        tracking_df = pd.read_parquet(tracking_path)\n        tracking_df = drop_body_parts(tracking_df)\n        self.tracking_buf = _precompute_tracking_buffers(tracking_df)\n        \n        self.width = video_info.get('video_width_pix', 640)\n        self.height = video_info.get('video_height_pix', 480)\n        \n        available_frames = sorted(list(self.tracking_buf['frame_to_rows'].keys()))\n        buffer = max(abs(o) for o in self.offsets)\n        self.frames = list(range(min(available_frames) + buffer, \n                                max(available_frames) - buffer + 1, stride))\n        \n    def __len__(self):\n        return len(self.frames)\n    \n    def __getitem__(self, idx):\n        frame = self.frames[idx]\n        clip = build_2p5d_clip(self.tracking_buf, frame, self.width, self.height,\n                               self.img_size, self.offsets, augment=False)\n        return clip, frame\n\n# ===================================================================\n# 4. Model with Dropout\n# ===================================================================\n\nclass MABeEfficientNet2p5D(nn.Module):\n    def __init__(self, model_name, n_classes, pretrained=True, in_chans=7, dropout=0.3):\n        super().__init__()\n        self.backbone = timm.create_model(\n            model_name, \n            pretrained=pretrained, \n            in_chans=in_chans, \n            num_classes=0\n        )\n        self.dropout = nn.Dropout(dropout)\n        self.classifier = nn.Linear(self.backbone.num_features, n_classes)\n    \n    def forward(self, x):\n        features = self.backbone(x)\n        features = self.dropout(features)\n        return self.classifier(features)\n\n# ===================================================================\n# 5. F-Score Calculation\n# ===================================================================\n\ndef calculate_frame_fscore(pred_events, true_events, class_to_idx):\n    \"\"\"\n    Calculate F-score based on frame-level predictions vs ground truth.\n    \"\"\"\n    if len(true_events) == 0:\n        return 0.0, 0.0, 0.0\n    \n    # Build frame-level ground truth\n    true_frames = defaultdict(set)\n    for _, row in true_events.iterrows():\n        class_label = f\"{row['agent_id']}_{row['target_id']}_{row['action']}\".lower()\n        if class_label not in class_to_idx:\n            continue\n        for frame in range(int(row['start_frame']), int(row['stop_frame']) + 1):\n            true_frames[frame].add(class_label)\n    \n    # Build frame-level predictions\n    pred_frames = defaultdict(set)\n    for event in pred_events:\n        class_label = f\"{event['agent_id']}_{event['target_id']}_{event['action']}\".lower()\n        for frame in range(event['start_frame'], event['stop_frame'] + 1):\n            pred_frames[frame].add(class_label)\n    \n    # Calculate frame-level metrics\n    tp = 0\n    fp = 0\n    fn = 0\n    \n    all_frames = set(true_frames.keys()) | set(pred_frames.keys())\n    for frame in all_frames:\n        true_labels = true_frames.get(frame, set())\n        pred_labels = pred_frames.get(frame, set())\n        \n        tp += len(true_labels & pred_labels)\n        fp += len(pred_labels - true_labels)\n        fn += len(true_labels - pred_labels)\n    \n    precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0\n    recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0\n    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0\n    \n    return precision, recall, f1\n\n# ===================================================================\n# 6. Training with F-Score Tracking\n# ===================================================================\n\ndef compute_class_weights(video_annotations, class_to_idx):\n    class_counts = Counter()\n    total = 0\n    for ann_df in video_annotations.values():\n        for _, row in ann_df.iterrows():\n            label = f\"{row['agent_id']}_{row['target_id']}_{row['action']}\".lower()\n            if label in class_to_idx:\n                dur = int(row['stop_frame']) - int(row['start_frame']) + 1\n                class_counts[label] += dur\n                total += dur\n    \n    weights = torch.ones(len(class_to_idx))\n    for label, idx in class_to_idx.items():\n        if label in class_counts:\n            freq = class_counts[label] / total\n            weights[idx] = min(CFG.POS_WEIGHT_SCALE / max(freq, 1e-6), 100.0)\n    return weights\n\ndef train_epoch(model, loader, criterion, optimizer, device, scaler, epoch):\n    model.train()\n    total_loss = 0\n    \n    pbar = tqdm(loader, desc=f\"Epoch {epoch+1} [Train]\")\n    for clips, labels in pbar:\n        clips, labels = clips.to(device), labels.to(device)\n        \n        with torch.cuda.amp.autocast(enabled=scaler is not None):\n            logits = model(clips)\n            loss = criterion(logits, labels)\n        \n        if scaler:\n            scaler.scale(loss).backward()\n            scaler.unscale_(optimizer)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.GRADIENT_CLIP)\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.GRADIENT_CLIP)\n            optimizer.step()\n        \n        optimizer.zero_grad()\n        total_loss += loss.item()\n        pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n    \n    return total_loss / len(loader)\n\n@torch.no_grad()\ndef validate_epoch_with_fscore(model, val_ids, train_meta_df, video_annotations, \n                                tracking_dir, class_to_idx, idx_to_class, \n                                criterion, device, epoch):\n    \"\"\"\n    Validate with both loss and F-score calculation.\n    \"\"\"\n    model.eval()\n    \n    # Calculate loss on validation samples\n    val_dataset = MABeTrainingDataset(\n        val_ids, train_meta_df, video_annotations, tracking_dir,\n        class_to_idx, CFG.IMG_SIZE, CFG.TEMPORAL_OFFSETS, \n        CFG.FRAME_SAMPLING_RATE * 2, augment=False  # Sparse sampling for speed\n    )\n    val_loader = DataLoader(val_dataset, batch_size=CFG.BATCH_SIZE * 2, \n                           shuffle=False, num_workers=0, pin_memory=True)\n    \n    total_loss = 0\n    pbar = tqdm(val_loader, desc=f\"Epoch {epoch+1} [Val Loss]\")\n    for clips, labels in pbar:\n        clips, labels = clips.to(device), labels.to(device)\n        logits = model(clips)\n        loss = criterion(logits, labels)\n        total_loss += loss.item()\n        pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n    \n    avg_loss = total_loss / len(val_loader)\n    \n    # Calculate F-score on a subset of validation videos (for speed)\n    print(f\"  Computing F-score on {min(5, len(val_ids))} validation videos...\")\n    all_precisions, all_recalls, all_f1s = [], [], []\n    \n    for video_id in val_ids[:5]:  # Sample 5 videos for F-score\n        try:\n            row = train_meta_df[train_meta_df['video_id'] == video_id].iloc[0]\n            video_info = {\n                'video_id': video_id,\n                'lab_id': row['lab_id'],\n                'video_width_pix': row.get('video_width_pix', 640),\n                'video_height_pix': row.get('video_height_pix', 480)\n            }\n            \n            pred_events = predict_video_for_validation(\n                model, video_info, tracking_dir, idx_to_class, device\n            )\n            \n            true_events = video_annotations.get(video_id, pd.DataFrame())\n            if len(true_events) > 0:\n                precision, recall, f1 = calculate_frame_fscore(\n                    pred_events, true_events, class_to_idx\n                )\n                all_precisions.append(precision)\n                all_recalls.append(recall)\n                all_f1s.append(f1)\n        except Exception as e:\n            continue\n    \n    avg_precision = np.mean(all_precisions) if all_precisions else 0.0\n    avg_recall = np.mean(all_recalls) if all_recalls else 0.0\n    avg_f1 = np.mean(all_f1s) if all_f1s else 0.0\n    \n    return avg_loss, avg_precision, avg_recall, avg_f1\n\ndef predict_video_for_validation(model, video_info, tracking_dir, idx_to_class, device):\n    \"\"\"\n    Predict events for a single video (used during validation).\n    \"\"\"\n    model.eval()\n    dataset = MABeInferenceDataset(video_info, tracking_dir, CFG.IMG_SIZE, \n                                   CFG.TEMPORAL_OFFSETS, CFG.INFERENCE_STRIDE * 2)\n    if len(dataset) == 0:\n        return []\n    \n    loader = DataLoader(dataset, batch_size=CFG.INFERENCE_BATCH_SIZE, \n                       shuffle=False, num_workers=0, pin_memory=True)\n    \n    all_probs, all_frames = [], []\n    with torch.no_grad():\n        for clips, frames in loader:\n            logits = model(clips.to(device))\n            probs = torch.sigmoid(logits).cpu().numpy()\n            all_probs.append(probs)\n            all_frames.extend(frames.numpy())\n    \n    all_probs = np.concatenate(all_probs, axis=0)\n    events = []\n    \n    for class_idx in range(all_probs.shape[1]):\n        predictions = (all_probs[:, class_idx] >= CFG.PREDICTION_THRESHOLD).astype(int)\n        changes = np.diff(np.concatenate([[0], predictions, [0]]))\n        starts = np.where(changes == 1)[0]\n        ends = np.where(changes == -1)[0]\n        \n        class_label = idx_to_class[class_idx]\n        agent_id, target_id, action = class_label.split('_')\n        \n        for start_idx, end_idx in zip(starts, ends):\n            start_frame = all_frames[start_idx]\n            end_frame = all_frames[end_idx - 1] if end_idx > 0 else all_frames[start_idx]\n            \n            if end_frame - start_frame >= CFG.MIN_EVENT_DURATION:\n                events.append({\n                    'video_id': video_info['video_id'],\n                    'agent_id': agent_id,\n                    'target_id': target_id,\n                    'action': action,\n                    'start_frame': int(start_frame),\n                    'stop_frame': int(end_frame)\n                })\n    \n    return events\n\n# ===================================================================\n# 7. Inference\n# ===================================================================\n\ndef predict_video(model, video_info, tracking_dir, idx_to_class, device):\n    model.eval()\n    dataset = MABeInferenceDataset(video_info, tracking_dir, CFG.IMG_SIZE, \n                                   CFG.TEMPORAL_OFFSETS, CFG.INFERENCE_STRIDE)\n    if len(dataset) == 0:\n        return []\n    \n    loader = DataLoader(dataset, batch_size=CFG.INFERENCE_BATCH_SIZE, \n                       shuffle=False, num_workers=0, pin_memory=True)\n    \n    all_probs, all_frames = [], []\n    with torch.no_grad():\n        for clips, frames in loader:\n            logits = model(clips.to(device))\n            probs = torch.sigmoid(logits).cpu().numpy()\n            all_probs.append(probs)\n            all_frames.extend(frames.numpy())\n    \n    all_probs = np.concatenate(all_probs, axis=0)\n    events = []\n    \n    for class_idx in range(all_probs.shape[1]):\n        predictions = (all_probs[:, class_idx] >= CFG.PREDICTION_THRESHOLD).astype(int)\n        changes = np.diff(np.concatenate([[0], predictions, [0]]))\n        starts = np.where(changes == 1)[0]\n        ends = np.where(changes == -1)[0]\n        \n        class_label = idx_to_class[class_idx]\n        agent_id, target_id, action = class_label.split('_')\n        \n        for start_idx, end_idx in zip(starts, ends):\n            start_frame = all_frames[start_idx]\n            end_frame = all_frames[end_idx - 1] if end_idx > 0 else all_frames[start_idx]\n            \n            if end_frame - start_frame >= CFG.MIN_EVENT_DURATION:\n                events.append({\n                    'video_id': video_info['video_id'],\n                    'agent_id': agent_id,\n                    'target_id': target_id,\n                    'action': action,\n                    'start_frame': int(start_frame),\n                    'stop_frame': int(end_frame)\n                })\n    \n    return events\n\n# ===================================================================\n# 8. Main with F-Score Tracking\n# ===================================================================\n\ndef main():\n    print(f\"\\n{'='*60}\\n🛡️ ANTI-OVERFITTING TRAINING with F-Score 🛡️\\n{'='*60}\\n\")\n    \n    # Load data\n    print(\"Loading metadata...\")\n    train_meta_df = pd.read_csv(CFG.DATA_DIR / \"train.csv\")\n    test_meta_df = pd.read_csv(CFG.DATA_DIR / \"test.csv\")\n    \n    # Build vocabulary\n    print(\"\\nBuilding vocabulary...\")\n    video_to_ann_path = find_annotated_videos(CFG.DATA_DIR / \"train_annotation\")\n    class_to_idx, idx_to_class, video_annotations = create_class_vocabulary(\n        train_meta_df, video_to_ann_path, CFG.MAX_VIDEOS_TRAIN + CFG.MAX_VIDEOS_VAL)\n    \n    with open(CFG.OUTPUT_DIR / 'class_to_idx.json', 'w') as f:\n        json.dump(class_to_idx, f)\n    \n    # Split data\n    all_ids = list(video_annotations.keys())\n    train_ids = all_ids[:CFG.MAX_VIDEOS_TRAIN]\n    val_ids = all_ids[CFG.MAX_VIDEOS_TRAIN:CFG.MAX_VIDEOS_TRAIN + CFG.MAX_VIDEOS_VAL]\n    print(f\"\\nTrain: {len(train_ids)} videos | Val: {len(val_ids)} videos\")\n    \n    # Create training dataset\n    print(\"\\nCreating training dataset...\")\n    train_dataset = MABeTrainingDataset(\n        train_ids, train_meta_df, video_annotations,\n        CFG.DATA_DIR / \"train_tracking\", class_to_idx,\n        CFG.IMG_SIZE, CFG.TEMPORAL_OFFSETS, CFG.FRAME_SAMPLING_RATE,\n        augment=CFG.AUGMENT_DATA\n    )\n    \n    train_loader = DataLoader(train_dataset, batch_size=CFG.BATCH_SIZE, shuffle=True,\n                             num_workers=0, pin_memory=True, drop_last=True)\n    \n    print(f\"Train batches: {len(train_loader)}\")\n    \n    # Initialize model\n    print(f\"\\nInitializing model with dropout={CFG.DROPOUT}...\")\n    model = MABeEfficientNet2p5D(\n        CFG.MODEL_NAME, len(class_to_idx), True, CFG.IN_CHANNELS, CFG.DROPOUT\n    )\n    if CFG.USE_MULTI_GPU:\n        print(f\"Using {torch.cuda.device_count()} GPUs\")\n        model = DataParallel(model)\n    model = model.to(DEVICE)\n    print(f\"✓ Model ready!\")\n    \n    # Training setup\n    print(\"\\nSetting up training...\")\n    train_annotations = {vid: video_annotations[vid] for vid in train_ids if vid in video_annotations}\n    pos_weights = compute_class_weights(train_annotations, class_to_idx).to(DEVICE)\n    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weights)\n    optimizer = torch.optim.AdamW(model.parameters(), lr=CFG.LEARNING_RATE, \n                                  weight_decay=CFG.WEIGHT_DECAY)\n    scaler = torch.cuda.amp.GradScaler() if CFG.MIXED_PRECISION else None\n    \n    # Training loop with F-score tracking\n    print(f\"\\n{'='*60}\\nTraining for up to {CFG.NUM_EPOCHS} epochs\\n{'='*60}\\n\")\n    best_val_loss = float('inf')\n    best_f1 = 0.0\n    patience_counter = 0\n    \n    for epoch in range(CFG.NUM_EPOCHS):\n        train_loss = train_epoch(model, train_loader, criterion, optimizer, DEVICE, scaler, epoch)\n        val_loss, val_precision, val_recall, val_f1 = validate_epoch_with_fscore(\n            model, val_ids, train_meta_df, video_annotations,\n            CFG.DATA_DIR / \"train_tracking\", class_to_idx, idx_to_class,\n            criterion, DEVICE, epoch\n        )\n        \n        print(f\"\\n  Train Loss: {train_loss:.4f}\")\n        print(f\"  Val Loss: {val_loss:.4f}\")\n        print(f\"  Val F1: {val_f1:.4f} (P: {val_precision:.4f}, R: {val_recall:.4f})\")\n        \n        # Save best model based on validation loss\n        if val_loss < best_val_loss:\n            best_val_loss = val_loss\n            patience_counter = 0\n            save_model = model.module if isinstance(model, DataParallel) else model\n            torch.save(save_model.state_dict(), CFG.OUTPUT_DIR / 'model_best.pth')\n            print(f\"  ✓ Saved best model (val_loss: {val_loss:.4f}, val_f1: {val_f1:.4f})\")\n        else:\n            patience_counter += 1\n            print(f\"  ⚠ No improvement ({patience_counter}/{CFG.PATIENCE})\")\n            \n            if patience_counter >= CFG.PATIENCE:\n                print(f\"\\n  🛑 Early stopping triggered at epoch {epoch+1}\")\n                break\n        \n        if val_f1 > best_f1:\n            best_f1 = val_f1\n        \n        print()\n        gc.collect()\n        torch.cuda.empty_cache()\n    \n    # Inference\n    print(f\"{'='*60}\\nRunning Inference\\n{'='*60}\\n\")\n    save_model = model.module if isinstance(model, DataParallel) else model\n    save_model.load_state_dict(torch.load(CFG.OUTPUT_DIR / 'model_best.pth'))\n    model.eval()\n    \n    all_events = []\n    for _, row in test_meta_df.iterrows():\n        video_info = {\n            'video_id': row['video_id'],\n            'lab_id': row['lab_id'],\n            'video_width_pix': row.get('video_width_pix', 640),\n            'video_height_pix': row.get('video_height_pix', 480)\n        }\n        try:\n            events = predict_video(model, video_info, CFG.DATA_DIR / \"test_tracking\", \n                                  idx_to_class, DEVICE)\n            all_events.extend(events)\n            print(f\"✓ Video {video_info['video_id']}: {len(events)} events\")\n        except Exception as e:\n            print(f\"✗ Video {video_info['video_id']}: {str(e)[:80]}\")\n    \n    # Save submission\n    if len(all_events) > 0:\n        submission = pd.DataFrame(all_events)\n    else:\n        submission = pd.DataFrame({\n            'video_id': [test_meta_df.iloc[0]['video_id']],\n            'agent_id': ['mouse1'], 'target_id': ['self'], 'action': ['rear'],\n            'start_frame': [0], 'stop_frame': [100]\n        })\n    \n    submission.index.name = 'row_id'\n    submission.to_csv(CFG.OUTPUT_DIR / 'submission.csv')\n    \n    print(f\"\\n{'='*60}\")\n    print(f\"✅ DONE! {len(submission)} predictions\")\n    print(f\"🛡️ Configuration:\")\n    print(f\"   - Training videos: {len(train_ids)}\")\n    print(f\"   - Validation videos: {len(val_ids)}\")\n    print(f\"   - Best validation F1: {best_f1:.4f}\")\n    print(f\"   - Prediction threshold: {CFG.PREDICTION_THRESHOLD}\")\n    print(f\"   - Dropout: {CFG.DROPOUT}\")\n    print(f\"   - Label smoothing: {CFG.LABEL_SMOOTHING}\")\n    print(f\"{'='*60}\\n\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}