{"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":13719397,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# %%writefile stgcn_mabe_final_official.py\n# import os\n# import sys\n# import json\n# import logging\n# import warnings\n# from pathlib import Path\n# from typing import Dict, List, Tuple, Optional, Any, DefaultDict\n# from collections import defaultdict\n# import argparse\n\n# import numpy as np\n# import polars as pl\n# import torch\n# import torch.nn as nn\n# import torch.nn.functional as F\n# from torch.utils.data import Dataset, DataLoader, random_split\n# from tqdm.auto import tqdm\n\n# # ========================\n# # CUDA AVAILABILITY CHECK — IN RA NGAY LẬP TỨC\n# # ========================\n# print(\"\\n\" + \"=\"*70)\n# print(\">>> 🧪 SYSTEM CUDA & GPU CHECK (BEFORE ANYTHING ELSE)\")\n# print(f\"PyTorch version: {torch.__version__}\")\n# print(f\"CUDA available: {torch.cuda.is_available()}\")\n# if torch.cuda.is_available():\n#     print(f\"Number of GPUs: {torch.cuda.device_count()}\")\n#     for i in range(torch.cuda.device_count()):\n#         print(f\"  → GPU {i}: {torch.cuda.get_device_name(i)}\")\n# else:\n#     print(\"  → NO GPU AVAILABLE. Training on CPU.\")\n# print(\"<<<\")\n# print(\"=\"*70 + \"\\n\")\n\n# # ========================\n# # Config — auto collect behaviors + FULL DATA MODE\n# # ========================\n\n# class Config:\n#     def __init__(self):\n#         self.data_root = Path(os.getenv(\"MABE_DATA_ROOT\", \"/kaggle/input/MABe-mouse-behavior-detection\"))\n#         self.submission_file = \"submission.csv\"\n#         self.model_save_path = \"stgcn_model.pth\"\n#         self.window_size = 64\n#         self.stride = 32\n#         self.batch_size = 16\n#         self.num_epochs = 3\n#         self.lr = 0.001\n\n#         # Force CUDA if available\n#         self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n#         self.num_mice = 4\n#         self.in_channels = 2\n\n#         # 🚨 AUTO COLLECT BEHAVIORS FROM DATA 🚨\n#         self.class_to_idx = self._collect_all_behaviors()\n#         self.num_classes = len(self.class_to_idx)\n#         self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}\n\n#         logger.info(f\"✅ Found {self.num_classes} unique behaviors: {list(self.class_to_idx.keys())}\")\n#         logger.info(f\"🖥️  Using device: {self.device}\")\n\n#     def _collect_all_behaviors(self) -> Dict[str, int]:\n#         \"\"\"Scan all .parquet files to collect unique behaviors\"\"\"\n#         train_csv_path = self.data_root / \"train.csv\"\n#         if not train_csv_path.exists():\n#             return {\"avoid\": 0}\n\n#         df = pl.read_csv(train_csv_path)\n#         behaviors = set()\n\n#         for row in df.to_dicts():  # ⬅️ DUYỆT TOÀN BỘ — KHÔNG .head(10) NỮA\n#             lab_id = row[\"lab_id\"]\n#             video_id = row[\"video_id\"]\n#             annot_path = self.data_root / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n#             if not annot_path.exists():\n#                 continue\n#             try:\n#                 annot_df = pl.read_parquet(annot_path)\n#                 if \"action\" in annot_df.columns:\n#                     behaviors.update(annot_df[\"action\"].unique().to_list())\n#             except Exception as e:\n#                 logger.warning(f\"Failed to read {annot_path}: {e}\")\n#                 continue\n\n#         sorted_behaviors = sorted(list(behaviors))\n#         return {behavior: idx for idx, behavior in enumerate(sorted_behaviors)}\n\n#     @property\n#     def train_csv(self): return self.data_root / \"train.csv\"\n#     @property\n#     def test_csv(self): return self.data_root / \"test.csv\"\n#     @property\n#     def train_annot_dir(self): return self.data_root / \"train_annotation\"\n#     @property\n#     def train_track_dir(self): return self.data_root / \"train_tracking\"\n#     @property\n#     def test_track_dir(self): return self.data_root / \"test_tracking\"\n\n# logging.basicConfig(level=logging.INFO, format=\"%(asctime)s | %(levelname)s | %(message)s\")\n# logger = logging.getLogger(__name__)\n# warnings.filterwarnings(\"ignore\")\n\n# # ========================\n# # ST-GCN Model — FIXED: no mean(dim=3), use .reshape()\n# # ========================\n\n# class SpatialGraphConv(nn.Module):\n#     def __init__(self, in_channels, out_channels, A):\n#         super().__init__()\n#         self.A = nn.Parameter(A, requires_grad=False)\n#         self.conv = nn.Conv1d(in_channels, out_channels * A.size(0), kernel_size=1)\n\n#     def forward(self, x):\n#         x = self.conv(x)\n#         n, kc, v = x.size()\n#         x = x.view(n, self.A.size(0), kc // self.A.size(0), v)\n#         x = torch.einsum('nvcv,vw->nwc', x, self.A)\n#         return x.contiguous()\n\n# class STGCNBlock(nn.Module):\n#     def __init__(self, in_channels, out_channels, A, temporal_kernel_size=9):\n#         super().__init__()\n#         self.sgc = SpatialGraphConv(in_channels, out_channels, A)\n#         self.tcn = nn.Sequential(\n#             nn.BatchNorm2d(out_channels),\n#             nn.ReLU(),\n#             nn.Conv2d(out_channels, out_channels, (temporal_kernel_size, 1), padding=(temporal_kernel_size//2, 0)),\n#             nn.BatchNorm2d(out_channels),\n#             nn.Dropout(0.5)\n#         )\n#         self.relu = nn.ReLU()\n\n#     def forward(self, x):\n#         # x: (N, C, T, V)\n#         N, C, T, V = x.size()\n#         x = x.permute(0, 2, 1, 3).contiguous().view(N * T, C, V)  # (N*T, C, V)\n#         x = self.sgc(x)  # (N*T, V, C') → (N*T, C', V)\n#         x = x.view(N, T, -1, V).permute(0, 2, 1, 3).contiguous()  # (N, C', T, V)\n#         x = self.tcn(x)  # (N, C', T, V) — no mean!\n#         return self.relu(x)\n\n# class MABeSTGCN(nn.Module):\n#     def __init__(self, num_classes, num_mice=4, in_channels=2):\n#         super().__init__()\n#         self.num_mice = num_mice\n#         self.total_nodes = num_mice\n#         A = torch.ones(num_mice, num_mice)\n#         self.A = nn.Parameter(A, requires_grad=False)\n#         self.block1 = STGCNBlock(in_channels, 64, self.A)\n#         self.block2 = STGCNBlock(64, 128, self.A)\n#         self.block3 = STGCNBlock(128, 256, self.A)\n#         self.fc = nn.Linear(256 * num_mice, num_classes)  # flatten all nodes\n\n#     def forward(self, x):\n#         # x: (N, T, M, 2) → reshape\n#         N, T, M, C = x.size()\n#         x = x.permute(0, 3, 1, 2).contiguous()  # (N, 2, T, M)\n\n#         x = self.block1(x)\n#         x = self.block2(x)\n#         x = self.block3(x)\n\n#         mid_frame = T // 2\n#         x = x[:, :, mid_frame, :]  # (N, 256, M)\n#         x = x.reshape(N, -1)  # ✅ Use .reshape() instead of .view()\n#         return self.fc(x)\n\n# # ========================\n# # Dataset — WITH FULL CACHING + TQDM LOADING\n# # ========================\n\n# class MABeDataset(Dataset):\n#     def __init__(self, video_ids: List[int], lab_ids: List[str], track_dir: Path, annot_dir: Path,\n#                  window_size: int = 64, stride: int = 32, class_to_idx: Dict[str, int] = None, is_test: bool = False, num_mice: int = 4):\n#         self.video_ids = video_ids\n#         self.lab_ids = lab_ids\n#         self.track_dir = track_dir\n#         self.annot_dir = annot_dir\n#         self.window_size = window_size\n#         self.stride = stride\n#         self.class_to_idx = class_to_idx or {}\n#         self.is_test = is_test\n#         self.num_mice = num_mice\n#         self.samples = []\n\n#         # 🚀 CACHING SYSTEM\n#         self.track_cache = {}    # (lab_id, video_id) -> pl.DataFrame\n#         self.xy_cache = {}       # (lab_id, video_id) -> np.ndarray (T, M, 2) — đã chuẩn hóa\n\n#         self._build_samples()\n\n#     def _get_track_path(self, lab_id: str, video_id: int) -> Optional[Path]:\n#         path = self.track_dir / lab_id / f\"{video_id}.parquet\"\n#         if path.exists():\n#             return path\n#         return None\n\n#     def _build_samples(self):\n#         total_video = len(self.video_ids)\n#         logger.info(f\"🚀 Starting to process {total_video} videos...\")\n\n#         count_track = 0\n#         count_parquet = 0\n#         count_valid_labels = 0\n\n#         # ✅ THÊM TQDM VÀO VÒNG LẶP XỬ LÝ VIDEO\n#         for idx, (video_id, lab_id) in enumerate(tqdm(zip(self.video_ids, self.lab_ids), total=total_video, desc=\"📂 Loading & Processing Videos\")):\n#             logger.info(f\"🎬 Processing video {idx+1}/{total_video}: {lab_id}/{video_id}\")\n\n#             track_path = self._get_track_path(lab_id, video_id)\n#             if track_path is None:\n#                 logger.warning(f\"⚠️ No .parquet found in train_tracking for {lab_id}/{video_id}\")\n#                 continue\n#             count_track += 1\n\n#             try:\n#                 track_df = pl.read_parquet(track_path)\n#                 self.track_cache[(lab_id, video_id)] = track_df  # ✅ CACHE TRACK_DF\n#                 logger.info(f\"✅ Loaded & cached tracking {track_df.shape}\")\n#             except Exception as e:\n#                 logger.warning(f\"⚠️ Failed to read tracking {track_path}: {e}\")\n#                 continue\n\n#             if \"video_frame\" not in track_df.columns or \"mouse_id\" not in track_df.columns or \"x\" not in track_df.columns or \"y\" not in track_df.columns:\n#                 logger.warning(f\"⚠️ Missing required columns in {track_path}\")\n#                 continue\n\n#             unique_mice = sorted(track_df[\"mouse_id\"].unique().to_list())\n#             mouse_to_idx = {mouse: i for i, mouse in enumerate(unique_mice[:self.num_mice])}\n#             all_frames = sorted(track_df[\"video_frame\"].unique().to_list())\n#             T = len(all_frames)\n#             frame_to_idx = {frame: i for i, frame in enumerate(all_frames)}\n\n#             xy = np.full((T, self.num_mice, 2), np.nan, dtype=np.float32)\n\n#             for row in track_df.to_dicts():\n#                 frame = row[\"video_frame\"]\n#                 mouse_id = row[\"mouse_id\"]\n#                 x = row[\"x\"]\n#                 y = row[\"y\"]\n#                 if frame in frame_to_idx and mouse_id in mouse_to_idx:\n#                     t = frame_to_idx[frame]\n#                     m = mouse_to_idx[mouse_id]\n#                     xy[t, m, 0] = x\n#                     xy[t, m, 1] = y\n\n#             # Forward fill\n#             for m in range(self.num_mice):\n#                 for t in range(1, T):\n#                     if np.isnan(xy[t, m, 0]):\n#                         xy[t, m] = xy[t-1, m]\n\n#             # Backward fill\n#             for m in range(self.num_mice):\n#                 for t in range(T-2, -1, -1):\n#                     if np.isnan(xy[t, m, 0]):\n#                         xy[t, m] = xy[t+1, m]\n\n#             centroid = np.nanmean(xy, axis=1, keepdims=True)\n#             xy_norm = xy - centroid\n#             xy_norm = np.nan_to_num(xy_norm, nan=0.0)\n\n#             # ✅ CACHE XY_NORM — chỉ tính 1 lần!\n#             self.xy_cache[(lab_id, video_id)] = xy_norm\n\n#             if not self.is_test:  # ✅ CHỈ XỬ LÝ ANNOTATION NẾU KHÔNG PHẢI TEST SET\n#                 annot_path = self.annot_dir / lab_id / f\"{video_id}.parquet\"\n#                 if not annot_path.exists():\n#                     logger.warning(f\"⚠️ No .parquet for {lab_id}/{video_id}\")\n#                     continue\n#                 count_parquet += 1\n\n#                 try:\n#                     annot_df = pl.read_parquet(annot_path)\n#                     logger.info(f\"✅ Loaded annotation {annot_df.shape}\")\n#                 except Exception as e:\n#                     logger.warning(f\"⚠️ Failed to read annotation {annot_path}: {e}\")\n#                     continue\n\n#                 if \"agent_id\" not in annot_df.columns or \"action\" not in annot_df.columns:\n#                     logger.warning(f\"⚠️ Missing required columns in {annot_path}\")\n#                     continue\n\n#                 unique_pairs = annot_df.select([\"agent_id\", \"target_id\"]).unique().to_dicts()\n#                 if len(unique_pairs) == 0:\n#                     logger.warning(f\"⚠️ No pairs for {lab_id}/{video_id}\")\n#                     continue\n\n#                 for pair in unique_pairs:\n#                     agent_id = pair[\"agent_id\"]\n#                     target_id = pair[\"target_id\"]\n#                     pair_df = annot_df.filter((pl.col(\"agent_id\") == agent_id) & (pl.col(\"target_id\") == target_id))\n#                     if len(pair_df) == 0:\n#                         continue\n\n#                     labels = np.full(T, -1, dtype=np.int64)\n#                     for row in pair_df.to_dicts():\n#                         s, e, action = row[\"start_frame\"], row[\"stop_frame\"], row[\"action\"]\n#                         if action in self.class_to_idx:\n#                             labels[s:e] = self.class_to_idx[action]\n\n#                     for start in range(0, T - self.window_size + 1, self.stride):\n#                         window_labels = labels[start:start + self.window_size]\n#                         if np.any(window_labels != -1):\n#                             mid = start + self.window_size // 2\n#                             label = labels[mid] if labels[mid] != -1 else -1\n#                             if label != -1:\n#                                 self.samples.append((video_id, lab_id, agent_id, target_id, start, label))\n#                                 count_valid_labels += 1\n#             else:\n#                 # ✅ TEST SET: tạo dummy samples để DataLoader có thể chạy\n#                 T = xy_norm.shape[0]\n#                 for start in range(0, T - self.window_size + 1, self.stride):\n#                     # agent_id, target_id = -1 (placeholder), label = -1\n#                     self.samples.append((video_id, lab_id, -1, -1, start, -1))\n\n#         logger.info(f\"✅ Found {count_track} .parquet files in train_tracking/\")\n#         if not self.is_test:\n#             logger.info(f\"✅ Found {count_parquet} .parquet files in train_annotation/\")\n#         logger.info(f\"✅ Created {len(self.samples)} samples\")\n\n#     def __len__(self):\n#         return len(self.samples)\n\n#     def __getitem__(self, idx):\n#         video_id, lab_id, agent_id, target_id, start, label = self.samples[idx]\n\n#         # ✅ LẤY TỪ CACHE — KHÔNG ĐỌC FILE, KHÔNG TIỀN XỬ LÝ LẠI\n#         xy_norm = self.xy_cache[(lab_id, video_id)]\n\n#         window_xy = xy_norm[start:start + self.window_size]  # (W, M, 2)\n#         xy_tensor = torch.tensor(window_xy, dtype=torch.float32)\n#         label_tensor = torch.tensor(label, dtype=torch.long)\n\n#         if self.is_test:\n#             return xy_tensor, video_id, lab_id, agent_id, target_id, start\n#         else:\n#             return xy_tensor, label_tensor\n\n# # ========================\n# # Training — OFFICIAL FULL TRAIN\n# # ========================\n\n# def train_model(cfg: Config):\n#     print(\"\\n🚀 [OFFICIAL TRAINING START] — Using FULL dataset\\n\")\n\n#     logger.info(\"Loading train.csv...\")\n#     train_df = pl.read_csv(cfg.train_csv)\n\n#     # ✅ DÙNG TOÀN BỘ VIDEO — KHÔNG GIỚI HẠN\n#     video_ids = train_df[\"video_id\"].to_list()\n#     lab_ids = train_df[\"lab_id\"].to_list()\n\n#     logger.info(\"Creating dataset...\")\n#     dataset = MABeDataset(\n#         video_ids=video_ids,\n#         lab_ids=lab_ids,\n#         track_dir=cfg.train_track_dir,\n#         annot_dir=cfg.train_annot_dir,\n#         window_size=cfg.window_size,\n#         stride=cfg.stride,\n#         class_to_idx=cfg.class_to_idx,\n#         is_test=False,\n#         num_mice=cfg.num_mice\n#     )\n\n#     if len(dataset) == 0:\n#         raise ValueError(\"❌ Dataset is empty!\")\n\n#     logger.info(f\"✅ Total training samples: {len(dataset)}\")\n\n#     train_size = max(1, int(0.8 * len(dataset)))\n#     val_size = len(dataset) - train_size\n#     if val_size == 0:\n#         val_size = 1\n#         train_size = len(dataset) - 1\n\n#     train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n\n#     # ✅ Dùng num_workers=2 — nếu lỗi, tự sửa thành 0\n#     try:\n#         train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True)\n#         val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)\n#     except Exception as e:\n#         logger.warning(f\"⚠️ DataLoader with num_workers=2 failed: {e}. Falling back to num_workers=0.\")\n#         train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=0, pin_memory=True)\n#         val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n\n#     model = MABeSTGCN(\n#         num_classes=cfg.num_classes,\n#         num_mice=cfg.num_mice,\n#         in_channels=cfg.in_channels\n#     ).to(cfg.device)\n\n#     print(\"\\n\" + \"=\"*60)\n#     print(\"🧠 MODEL DEVICE:\", next(model.parameters()).device)\n#     print(\"=\"*60 + \"\\n\")\n\n#     criterion = nn.CrossEntropyLoss()\n#     optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)\n\n#     best_val_acc = 0.0\n\n#     for epoch in range(cfg.num_epochs):\n#         model.train()\n#         train_loss = 0.0\n#         train_correct = 0\n#         train_total = 0\n\n#         pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{cfg.num_epochs} [Train]\")\n#         for batch_idx, (inputs, labels) in enumerate(pbar):\n#             inputs = inputs.to(cfg.device, non_blocking=True)\n#             labels = labels.to(cfg.device, non_blocking=True)\n\n#             if epoch == 0 and batch_idx == 0:\n#                 print(\"\\n\" + \"-\"*50)\n#                 print(\"📦 FIRST BATCH INPUT DEVICE:\", inputs.device)\n#                 print(\"🏷️  FIRST BATCH LABEL DEVICE:\", labels.device)\n#                 print(\"-\"*50 + \"\\n\")\n\n#             optimizer.zero_grad()\n#             outputs = model(inputs)\n#             loss = criterion(outputs, labels)\n#             loss.backward()\n#             optimizer.step()\n\n#             train_loss += loss.item()\n#             _, predicted = outputs.max(1)\n#             train_total += labels.size(0)\n#             train_correct += predicted.eq(labels).sum().item()\n\n#             pbar.set_postfix({'Loss': loss.item(), 'Acc': train_correct/train_total})\n\n#         model.eval()\n#         val_correct = 0\n#         val_total = 0\n#         with torch.no_grad():\n#             for inputs, labels in tqdm(val_loader, desc=\"Validation\"):\n#                 inputs = inputs.to(cfg.device, non_blocking=True)\n#                 labels = labels.to(cfg.device, non_blocking=True)\n#                 outputs = model(inputs)\n#                 _, predicted = outputs.max(1)\n#                 val_total += labels.size(0)\n#                 val_correct += predicted.eq(labels).sum().item()\n\n#         val_acc = val_correct / val_total\n#         logger.info(f\"Epoch {epoch+1}: Val Acc = {val_acc:.4f}\")\n\n#         if val_acc > best_val_acc:\n#             best_val_acc = val_acc\n#             torch.save(model.state_dict(), cfg.model_save_path)\n#             logger.info(f\"✅ Model saved with Val Acc = {val_acc:.4f}\")\n\n#     logger.info(f\"🎉 Training completed. Best Val Acc: {best_val_acc:.4f}\")\n\n# # ========================\n# # Inference & Submission — OFFICIAL FULL TEST SET\n# # ========================\n\n# def predict_and_submit(cfg: Config):\n#     logger.info(\"Loading test.csv...\")\n#     test_df = pl.read_csv(cfg.test_csv)\n\n#     # ✅ DÙNG TOÀN BỘ TEST SET\n#     video_ids = test_df[\"video_id\"].to_list()\n#     lab_ids = test_df[\"lab_id\"].to_list()\n\n#     logger.info(\"Creating test dataset...\")\n#     test_dataset = MABeDataset(\n#         video_ids=video_ids,\n#         lab_ids=lab_ids,\n#         track_dir=cfg.test_track_dir,\n#         annot_dir=cfg.train_annot_dir,\n#         window_size=cfg.window_size,\n#         stride=cfg.stride,\n#         class_to_idx=cfg.class_to_idx,\n#         is_test=True,\n#         num_mice=cfg.num_mice\n#     )\n\n#     if len(test_dataset) == 0:\n#         raise ValueError(\"❌ Test dataset is empty!\")\n\n#     logger.info(f\"✅ Total test samples: {len(test_dataset)}\")\n\n#     try:\n#         test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)\n#     except Exception as e:\n#         logger.warning(f\"⚠️ DataLoader with num_workers=2 failed: {e}. Falling back to num_workers=0.\")\n#         test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n\n#     model = MABeSTGCN(\n#         num_classes=cfg.num_classes,\n#         num_mice=cfg.num_mice,\n#         in_channels=cfg.in_channels\n#     ).to(cfg.device)\n#     model.load_state_dict(torch.load(cfg.model_save_path, map_location=cfg.device))\n#     model.eval()\n\n#     print(\"\\n🧠 INFERENCE MODEL DEVICE:\", next(model.parameters()).device)\n\n#     predictions = defaultdict(list)\n\n#     with torch.no_grad():\n#         for batch in tqdm(test_loader, desc=\"Inference\"):\n#             if len(batch) == 6:\n#                 inputs, video_ids_batch, lab_ids_batch, agent_ids_batch, target_ids_batch, start_frames = batch\n#             else:\n#                 inputs, video_ids_batch, lab_ids_batch, start_frames = batch\n#                 agent_ids_batch = [-1] * len(video_ids_batch)\n#                 target_ids_batch = [-1] * len(video_ids_batch)\n\n#             inputs = inputs.to(cfg.device, non_blocking=True)\n\n#             if len(predictions) == 0:\n#                 print(\"🧪 INFERENCE INPUT DEVICE:\", inputs.device)\n\n#             outputs = model(inputs)\n#             _, predicted = outputs.max(1)\n\n#             for i in range(len(video_ids_batch)):\n#                 vid = int(video_ids_batch[i])\n#                 aid = int(agent_ids_batch[i])\n#                 tid = int(target_ids_batch[i])\n#                 start = int(start_frames[i])\n#                 mid_frame = start + cfg.window_size // 2\n#                 label_idx = predicted[i].item()\n#                 predictions[(vid, aid, tid)].append((mid_frame, label_idx))\n\n#     records = []\n#     for row in test_df.to_dicts():  # ⬅️ DUYỆT TOÀN BỘ TEST SET\n#         video_id = row[\"video_id\"]\n#         lab_id = row[\"lab_id\"]\n#         behaviors = json.loads(row[\"behaviors_labeled\"].replace(\"'\", '\"')) if isinstance(row[\"behaviors_labeled\"], str) else []\n\n#         for behavior in behaviors:\n#             if len(behavior) != 3:\n#                 continue\n#             agent_str, target_str, action = behavior\n#             if action not in cfg.class_to_idx:\n#                 continue\n\n#             try:\n#                 agent_id = int(agent_str)\n#                 target_id = int(target_str)\n#             except:\n#                 continue\n\n#             key = (video_id, agent_id, target_id)\n#             if key not in predictions or len(predictions[key]) == 0:\n#                 # Nếu không có dự đoán, gán nhãn đầu tiên (fallback)\n#                 default_label = list(cfg.class_to_idx.keys())[0]\n#                 records.append((video_id, f\"mouse{agent_id}\", f\"mouse{target_id}\", default_label, 0, 1))\n#                 continue\n\n#             preds = sorted(predictions[key])\n#             if len(preds) == 0:\n#                 default_label = list(cfg.class_to_idx.keys())[0]\n#                 records.append((video_id, f\"mouse{agent_id}\", f\"mouse{target_id}\", default_label, 0, 1))\n#                 continue\n\n#             current_label = preds[0][1]\n#             start_frame = preds[0][0]\n#             for i in range(1, len(preds)):\n#                 frame, label = preds[i]\n#                 if label != current_label or frame != preds[i-1][0] + 1:\n#                     records.append((\n#                         video_id,\n#                         f\"mouse{agent_id}\", f\"mouse{target_id}\",\n#                         cfg.idx_to_class[current_label],\n#                         start_frame, frame\n#                     ))\n#                     current_label = label\n#                     start_frame = frame\n#             records.append((\n#                 video_id,\n#                 f\"mouse{agent_id}\", f\"mouse{target_id}\",\n#                 cfg.idx_to_class[current_label],\n#                 start_frame, preds[-1][0] + 1\n#             ))\n\n#     if len(records) == 0:\n#         logger.warning(\"❌ No predictions generated! Using fallback.\")\n#         default_label = list(cfg.class_to_idx.keys())[0]\n#         records = [(0, \"mouse0\", \"mouse0\", default_label, 0, 1)]\n\n#     submission_df = pl.DataFrame(\n#         records,\n#         schema={\n#             \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n#             \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n#         },\n#         orient=\"row\"\n#     )\n\n#     submission_df = submission_df.with_row_index(\"row_id\")\n#     submission_df.write_csv(cfg.submission_file)\n#     logger.info(f\"✅ Official Submission saved to {cfg.submission_file}\")\n#     logger.info(f\"📊 Submission shape: {submission_df.shape}\")\n#     print(\"\\n📋 Sample predictions:\")\n#     print(submission_df.head(10))\n\n# # ========================\n# # Main\n# # ========================\n\n# def main():\n#     parser = argparse.ArgumentParser()\n#     parser.add_argument(\"--mode\", choices=[\"train\", \"submit\", \"all\"], default=\"all\")\n#     args = parser.parse_args()\n\n#     cfg = Config()\n\n#     if args.mode in (\"train\", \"all\"):\n#         train_model(cfg)\n\n#     if args.mode in (\"submit\", \"all\"):\n#         predict_and_submit(cfg)\n\n# if __name__ == \"__main__\":\n#     main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-21T16:25:03.903447Z","iopub.execute_input":"2025-09-21T16:25:03.903922Z","iopub.status.idle":"2025-09-21T16:25:03.919312Z","shell.execute_reply.started":"2025-09-21T16:25:03.903897Z","shell.execute_reply":"2025-09-21T16:25:03.918669Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#!python stgcn_mabe_final_official.py --mode train","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-21T16:25:07.804096Z","iopub.execute_input":"2025-09-21T16:25:07.804764Z","iopub.status.idle":"2025-09-21T17:01:02.739103Z","shell.execute_reply.started":"2025-09-21T16:25:07.804741Z","shell.execute_reply":"2025-09-21T17:01:02.738382Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%writefile stgcn_mabe_final_official_fixed_v6.py\n# import os\n# import sys\n# import json\n# import logging\n# import warnings\n# from pathlib import Path\n# from typing import Dict, List, Tuple, Optional, Any, DefaultDict\n# from collections import defaultdict\n# import argparse\n\n# import numpy as np\n# import polars as pl\n# import torch\n# import torch.nn as nn\n# import torch.nn.functional as F\n# from torch.utils.data import Dataset, DataLoader, random_split\n# from tqdm.auto import tqdm\n\n# print(\"\\n\" + \"=\"*70)\n# print(\">>> 🧪 SYSTEM CUDA & GPU CHECK (BEFORE ANYTHING ELSE)\")\n# print(f\"PyTorch version: {torch.__version__}\")\n# print(f\"CUDA available: {torch.cuda.is_available()}\")\n# if torch.cuda.is_available():\n#     print(f\"Number of GPUs: {torch.cuda.device_count()}\")\n#     for i in range(torch.cuda.device_count()):\n#         print(f\"  → GPU {i}: {torch.cuda.get_device_name(i)}\")\n# else:\n#     print(\"  → NO GPU AVAILABLE. Training on CPU.\")\n# print(\"<<<\")\n# print(\"=\"*70 + \"\\n\")\n\n# class Config:\n#     def __init__(self):\n#         self.data_root = Path(os.getenv(\"MABE_DATA_ROOT\", \"/kaggle/input/MABe-mouse-behavior-detection\"))\n#         self.submission_file = \"submission.csv\"\n#         self.model_save_path = \"stgcn_model.pth\"\n#         self.window_size = 64\n#         self.stride = 32\n#         self.batch_size = 16\n#         self.num_epochs = 3\n#         self.lr = 0.001\n#         self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#         self.num_mice = 4\n#         self.in_channels = 2\n\n#         self.class_to_idx = self._collect_all_behaviors()\n#         self.num_classes = len(self.class_to_idx)\n#         self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}\n\n#         logger.info(f\"✅ Found {self.num_classes} unique behaviors: {list(self.class_to_idx.keys())}\")\n#         logger.info(f\"🖥️  Using device: {self.device}\")\n\n#     def _collect_all_behaviors(self) -> Dict[str, int]:\n#         train_csv_path = self.data_root / \"train.csv\"\n#         if not train_csv_path.exists():\n#             return {\"avoid\": 0}\n\n#         df = pl.read_csv(train_csv_path)\n#         behaviors = set()\n\n#         for row in df.to_dicts():\n#             lab_id = row[\"lab_id\"]\n#             video_id = row[\"video_id\"]\n#             annot_path = self.data_root / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n#             if not annot_path.exists():\n#                 continue\n#             try:\n#                 annot_df = pl.read_parquet(annot_path)\n#                 if \"action\" in annot_df.columns:\n#                     behaviors.update(annot_df[\"action\"].unique().to_list())\n#             except Exception as e:\n#                 logger.warning(f\"Failed to read {annot_path}: {e}\")\n#                 continue\n\n#         sorted_behaviors = sorted(list(behaviors))\n#         return {behavior: idx for idx, behavior in enumerate(sorted_behaviors)}\n\n#     @property\n#     def train_csv(self): return self.data_root / \"train.csv\"\n#     @property\n#     def test_csv(self): return self.data_root / \"test.csv\"\n#     @property\n#     def train_annot_dir(self): return self.data_root / \"train_annotation\"\n#     @property\n#     def train_track_dir(self): return self.data_root / \"train_tracking\"\n#     @property\n#     def test_track_dir(self): return self.data_root / \"test_tracking\"\n\n# logging.basicConfig(level=logging.INFO, format=\"%(asctime)s | %(levelname)s | %(message)s\")\n# logger = logging.getLogger(__name__)\n# warnings.filterwarnings(\"ignore\")\n\n# # ✅ Hàm safe_json_loads - tối ưu cho data của bạn\n# def safe_json_loads(s: Optional[str]) -> List[str]:\n#     if s is None: return []\n#     if isinstance(s, list): return [str(x) for x in s]\n#     if not isinstance(s, str): return []\n#     s = s.strip()\n#     if not s: return []\n#     try:\n#         loaded = json.loads(s)\n#         if isinstance(loaded, list):\n#             return [str(item) for item in loaded]\n#         else:\n#             return []\n#     except Exception:\n#         try:\n#             # Thử sửa dấu nháy đơn\n#             fixed = s.replace(\"'\", '\"')\n#             loaded = json.loads(fixed)\n#             if isinstance(loaded, list):\n#                 return [str(item) for item in loaded]\n#             else:\n#                 return []\n#         except Exception:\n#             return []\n\n# class SpatialGraphConv(nn.Module):\n#     def __init__(self, in_channels, out_channels, A):\n#         super().__init__()\n#         self.A = nn.Parameter(A, requires_grad=False)\n#         self.conv = nn.Conv1d(in_channels, out_channels * A.size(0), kernel_size=1)\n\n#     def forward(self, x):\n#         x = self.conv(x)\n#         n, kc, v = x.size()\n#         x = x.view(n, self.A.size(0), kc // self.A.size(0), v)\n#         x = torch.einsum('nvcv,vw->nwc', x, self.A)\n#         return x.contiguous()\n\n# class STGCNBlock(nn.Module):\n#     def __init__(self, in_channels, out_channels, A, temporal_kernel_size=9):\n#         super().__init__()\n#         self.sgc = SpatialGraphConv(in_channels, out_channels, A)\n#         self.tcn = nn.Sequential(\n#             nn.BatchNorm2d(out_channels),\n#             nn.ReLU(),\n#             nn.Conv2d(out_channels, out_channels, (temporal_kernel_size, 1), padding=(temporal_kernel_size//2, 0)),\n#             nn.BatchNorm2d(out_channels),\n#             nn.Dropout(0.5)\n#         )\n#         self.relu = nn.ReLU()\n\n#     def forward(self, x):\n#         N, C, T, V = x.size()\n#         x = x.permute(0, 2, 1, 3).contiguous().view(N * T, C, V)\n#         x = self.sgc(x)\n#         x = x.view(N, T, -1, V).permute(0, 2, 1, 3).contiguous()\n#         x = self.tcn(x)\n#         return self.relu(x)\n\n# class MABeSTGCN(nn.Module):\n#     def __init__(self, num_classes, num_mice=4, in_channels=2):\n#         super().__init__()\n#         self.num_mice = num_mice\n#         self.total_nodes = num_mice\n#         A = torch.ones(num_mice, num_mice)\n#         self.A = nn.Parameter(A, requires_grad=False)\n#         self.block1 = STGCNBlock(in_channels, 64, self.A)\n#         self.block2 = STGCNBlock(64, 128, self.A)\n#         self.block3 = STGCNBlock(128, 256, self.A)\n#         self.fc = nn.Linear(256 * num_mice, num_classes)\n\n#     def forward(self, x):\n#         N, T, M, C = x.size()\n#         x = x.permute(0, 3, 1, 2).contiguous()\n#         x = self.block1(x)\n#         x = self.block2(x)\n#         x = self.block3(x)\n#         mid_frame = T // 2\n#         x = x[:, :, mid_frame, :]\n#         x = x.reshape(N, -1)\n#         return self.fc(x)\n\n# class MABeDataset(Dataset):\n#     def __init__(self, video_ids: List[int], lab_ids: List[str], track_dir: Path, annot_dir: Path,\n#                  window_size: int = 64, stride: int = 32, class_to_idx: Dict[str, int] = None, is_test: bool = False, num_mice: int = 4):\n#         self.video_ids = video_ids\n#         self.lab_ids = lab_ids\n#         self.track_dir = track_dir\n#         self.annot_dir = annot_dir\n#         self.window_size = window_size\n#         self.stride = stride\n#         self.class_to_idx = class_to_idx or {}\n#         self.is_test = is_test\n#         self.num_mice = num_mice\n#         self.samples = []\n#         self.track_cache = {}\n#         self.xy_cache = {}\n\n#         self._build_samples()\n\n#     def _get_track_path(self, lab_id: str, video_id: int) -> Optional[Path]:\n#         path = self.track_dir / lab_id / f\"{video_id}.parquet\"\n#         if path.exists():\n#             return path\n#         return None\n\n#     def _build_samples(self):\n#         total_video = len(self.video_ids)\n#         logger.info(f\"🚀 Starting to process {total_video} videos...\")\n\n#         count_track = 0\n#         count_parquet = 0\n#         count_valid_labels = 0\n\n#         for idx, (video_id, lab_id) in enumerate(tqdm(zip(self.video_ids, self.lab_ids), total=total_video, desc=\"📂 Loading & Processing Videos\")):\n#             track_path = self._get_track_path(lab_id, video_id)\n#             if track_path is None:\n#                 continue\n#             count_track += 1\n\n#             try:\n#                 track_df = pl.read_parquet(track_path)\n#                 self.track_cache[(lab_id, video_id)] = track_df\n#             except Exception as e:\n#                 logger.warning(f\"Failed to cache tracking for {lab_id}/{video_id}: {e}\")\n#                 continue\n\n#             required_cols = [\"video_frame\", \"mouse_id\", \"x\", \"y\"]\n#             if not all(col in track_df.columns for col in required_cols):\n#                 logger.warning(f\"Missing required columns in {track_path}\")\n#                 continue\n\n#             # Lấy danh sách mouse_id thực tế — CHUYỂN VỀ INT nếu có thể\n#             unique_mice_raw = track_df[\"mouse_id\"].unique().to_list()\n#             unique_mice = []\n#             for m in unique_mice_raw:\n#                 try:\n#                     mid = int(m)\n#                     unique_mice.append(mid)\n#                 except (ValueError, TypeError):\n#                     logger.warning(f\"Cannot convert mouse_id {m} to int, skipping.\")\n#                     continue\n\n#             unique_mice = sorted(set(unique_mice))[:self.num_mice]\n#             if len(unique_mice) == 0:\n#                 continue\n\n#             mouse_to_idx = {mouse: i for i, mouse in enumerate(unique_mice)}\n#             all_frames = sorted(track_df[\"video_frame\"].unique().to_list())\n#             T = len(all_frames)\n#             frame_to_idx = {frame: i for i, frame in enumerate(all_frames)}\n\n#             xy = np.full((T, self.num_mice, 2), np.nan, dtype=np.float32)\n\n#             for row in track_df.to_dicts():\n#                 frame = row[\"video_frame\"]\n#                 mouse_id = row[\"mouse_id\"]\n#                 try:\n#                     mouse_id = int(mouse_id)\n#                 except (ValueError, TypeError):\n#                     continue\n#                 x = row[\"x\"]\n#                 y = row[\"y\"]\n#                 if frame in frame_to_idx and mouse_id in mouse_to_idx:\n#                     t = frame_to_idx[frame]\n#                     m = mouse_to_idx[mouse_id]\n#                     xy[t, m, 0] = x\n#                     xy[t, m, 1] = y\n\n#             # Interpolation forward\n#             for m in range(self.num_mice):\n#                 for t in range(1, T):\n#                     if np.isnan(xy[t, m, 0]):\n#                         xy[t, m] = xy[t-1, m]\n\n#             # Interpolation backward\n#             for m in range(self.num_mice):\n#                 for t in range(T-2, -1, -1):\n#                     if np.isnan(xy[t, m, 0]):\n#                         xy[t, m] = xy[t+1, m]\n\n#             centroid = np.nanmean(xy, axis=1, keepdims=True)\n#             xy_norm = xy - centroid\n#             xy_norm = np.nan_to_num(xy_norm, nan=0.0)\n#             self.xy_cache[(lab_id, video_id)] = xy_norm\n\n#             if not self.is_test:\n#                 annot_path = self.annot_dir / lab_id / f\"{video_id}.parquet\"\n#                 if not annot_path.exists():\n#                     continue\n#                 count_parquet += 1\n\n#                 try:\n#                     annot_df = pl.read_parquet(annot_path)\n#                 except Exception as e:\n#                     logger.warning(f\"Failed to read annotation {annot_path}: {e}\")\n#                     continue\n\n#                 if \"agent_id\" not in annot_df.columns or \"action\" not in annot_df.columns:\n#                     continue\n\n#                 unique_pairs = annot_df.select([\"agent_id\", \"target_id\"]).unique().to_dicts()\n#                 if len(unique_pairs) == 0:\n#                     continue\n\n#                 for pair in unique_pairs:\n#                     agent_id = pair[\"agent_id\"]\n#                     target_id = pair[\"target_id\"]\n#                     try:\n#                         agent_id = int(agent_id)\n#                         target_id = int(target_id)\n#                     except (ValueError, TypeError):\n#                         continue\n\n#                     pair_df = annot_df.filter((pl.col(\"agent_id\") == agent_id) & (pl.col(\"target_id\") == target_id))\n#                     if len(pair_df) == 0:\n#                         continue\n\n#                     labels = np.full(T, -1, dtype=np.int64)\n#                     for row in pair_df.to_dicts():\n#                         s, e, action = row[\"start_frame\"], row[\"stop_frame\"], row[\"action\"]\n#                         if action in self.class_to_idx:\n#                             labels[s:e] = self.class_to_idx[action]\n\n#                     for start in range(0, T - self.window_size + 1, self.stride):\n#                         window_labels = labels[start:start + self.window_size]\n#                         if np.any(window_labels != -1):\n#                             mid = start + self.window_size // 2\n#                             label = labels[mid] if labels[mid] != -1 else -1\n#                             if label != -1:\n#                                 self.samples.append((video_id, lab_id, agent_id, target_id, start, label))\n#                                 count_valid_labels += 1\n#             else:\n#                 # TEST SET: tạo sample cho tất cả cặp (agent, target) có thể — CHỈ DÙNG INT\n#                 for agent_id in unique_mice:\n#                     for target_id in unique_mice:\n#                         for start in range(0, T - self.window_size + 1, self.stride):\n#                             self.samples.append((video_id, lab_id, agent_id, target_id, start, -1))\n\n#         logger.info(f\"✅ Found {count_track} .parquet files in train_tracking/\")\n#         if not self.is_test:\n#             logger.info(f\"✅ Found {count_parquet} .parquet files in train_annotation/\")\n#         logger.info(f\"✅ Created {len(self.samples)} samples\")\n\n#     def __len__(self):\n#         return len(self.samples)\n\n#     def __getitem__(self, idx):\n#         video_id, lab_id, agent_id, target_id, start, label = self.samples[idx]\n#         xy_norm = self.xy_cache[(lab_id, video_id)]\n#         window_xy = xy_norm[start:start + self.window_size]\n#         xy_tensor = torch.tensor(window_xy, dtype=torch.float32)\n#         label_tensor = torch.tensor(label, dtype=torch.long)\n#         return xy_tensor, label_tensor, video_id, lab_id, agent_id, target_id, start\n\n# def train_model(cfg: Config):\n#     print(\"\\n🚀 [OFFICIAL TRAINING START] — Using FULL dataset\\n\")\n\n#     train_df = pl.read_csv(cfg.train_csv)\n#     video_ids = train_df[\"video_id\"].to_list()\n#     lab_ids = train_df[\"lab_id\"].to_list()\n\n#     dataset = MABeDataset(\n#         video_ids=video_ids,\n#         lab_ids=lab_ids,\n#         track_dir=cfg.train_track_dir,\n#         annot_dir=cfg.train_annot_dir,\n#         window_size=cfg.window_size,\n#         stride=cfg.stride,\n#         class_to_idx=cfg.class_to_idx,\n#         is_test=False,\n#         num_mice=cfg.num_mice\n#     )\n\n#     if len(dataset) == 0:\n#         raise ValueError(\"❌ Dataset is empty!\")\n\n#     logger.info(f\"✅ Total training samples: {len(dataset)}\")\n\n#     train_size = max(1, int(0.8 * len(dataset)))\n#     val_size = len(dataset) - train_size\n#     if val_size == 0:\n#         val_size = 1\n#         train_size = len(dataset) - 1\n\n#     train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n\n#     try:\n#         train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True)\n#         val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)\n#     except Exception as e:\n#         logger.warning(f\"Failed to use num_workers=2: {e}\")\n#         train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=0, pin_memory=True)\n#         val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n\n#     model = MABeSTGCN(\n#         num_classes=cfg.num_classes,\n#         num_mice=cfg.num_mice,\n#         in_channels=cfg.in_channels\n#     ).to(cfg.device)\n\n#     print(\"\\n\" + \"=\"*60)\n#     print(\"🧠 MODEL DEVICE:\", next(model.parameters()).device)\n#     print(\"=\"*60 + \"\\n\")\n\n#     criterion = nn.CrossEntropyLoss()\n#     optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)\n\n#     best_val_acc = 0.0\n\n#     for epoch in range(cfg.num_epochs):\n#         model.train()\n#         train_loss = 0.0\n#         train_correct = 0\n#         train_total = 0\n\n#         pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{cfg.num_epochs} [Train]\")\n#         for batch_idx, (inputs, labels, _, _, _, _, _) in enumerate(pbar):\n#             inputs = inputs.to(cfg.device, non_blocking=True)\n#             labels = labels.to(cfg.device, non_blocking=True)\n\n#             optimizer.zero_grad()\n#             outputs = model(inputs)\n#             loss = criterion(outputs, labels)\n#             loss.backward()\n#             optimizer.step()\n\n#             train_loss += loss.item()\n#             _, predicted = outputs.max(1)\n#             train_total += labels.size(0)\n#             train_correct += predicted.eq(labels).sum().item()\n\n#             pbar.set_postfix({'Loss': loss.item(), 'Acc': train_correct/train_total})\n\n#         model.eval()\n#         val_correct = 0\n#         val_total = 0\n#         with torch.no_grad():\n#             for inputs, labels, _, _, _, _, _ in tqdm(val_loader, desc=\"Validation\"):\n#                 inputs = inputs.to(cfg.device, non_blocking=True)\n#                 labels = labels.to(cfg.device, non_blocking=True)\n#                 outputs = model(inputs)\n#                 _, predicted = outputs.max(1)\n#                 val_total += labels.size(0)\n#                 val_correct += predicted.eq(labels).sum().item()\n\n#         val_acc = val_correct / val_total\n#         logger.info(f\"Epoch {epoch+1}: Val Acc = {val_acc:.4f}\")\n\n#         if val_acc > best_val_acc:\n#             best_val_acc = val_acc\n#             torch.save(model.state_dict(), cfg.model_save_path)\n#             logger.info(f\"✅ Model saved with Val Acc = {val_acc:.4f}\")\n\n#     logger.info(f\"🎉 Training completed. Best Val Acc: {best_val_acc:.4f}\")\n\n# def predict_and_submit(cfg: Config):\n#     logger.info(\"Loading test.csv...\")\n#     test_df = pl.read_csv(cfg.test_csv)\n#     video_ids = test_df[\"video_id\"].to_list()\n#     lab_ids = test_df[\"lab_id\"].to_list()\n\n#     # ✅ Tạo mapping: video_id → set of allowed (agent_str, target_str, action)\n#     video_to_allowed_triples = {}  # (agent_str, target_str, action) → bool\n\n#     for row in test_df.to_dicts():\n#         video_id = row[\"video_id\"]\n#         behaviors_str = row[\"behaviors_labeled\"]\n#         behaviors_list = safe_json_loads(behaviors_str)\n#         allowed_triples = set()\n#         for item in behaviors_list:\n#             parts = [p.strip() for p in str(item).replace(\"'\", \"\").split(\",\")]\n#             if len(parts) == 3:\n#                 agent_str, target_str, action = parts\n#                 allowed_triples.add((agent_str, target_str, action))\n#         video_to_allowed_triples[video_id] = allowed_triples\n\n#     logger.info(\"Creating test dataset...\")\n#     test_dataset = MABeDataset(\n#         video_ids=video_ids,\n#         lab_ids=lab_ids,\n#         track_dir=cfg.test_track_dir,\n#         annot_dir=cfg.train_annot_dir,\n#         window_size=cfg.window_size,\n#         stride=cfg.stride,\n#         class_to_idx=cfg.class_to_idx,\n#         is_test=True,\n#         num_mice=cfg.num_mice\n#     )\n\n#     if len(test_dataset) == 0:\n#         raise ValueError(\"❌ Test dataset is empty!\")\n\n#     logger.info(f\"✅ Total test samples: {len(test_dataset)}\")\n\n#     try:\n#         test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)\n#     except Exception as e:\n#         logger.warning(f\"Failed to use num_workers=2: {e}\")\n#         test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n\n#     model = MABeSTGCN(\n#         num_classes=cfg.num_classes,\n#         num_mice=cfg.num_mice,\n#         in_channels=cfg.in_channels\n#     ).to(cfg.device)\n#     model.load_state_dict(torch.load(cfg.model_save_path, map_location=cfg.device))\n#     model.eval()\n\n#     predictions = defaultdict(list)\n\n#     with torch.no_grad():\n#         for batch in tqdm(test_loader, desc=\"Inference\"):\n#             inputs, _, video_ids_batch, lab_ids_batch, agent_ids_batch, target_ids_batch, start_frames = batch\n#             inputs = inputs.to(cfg.device, non_blocking=True)\n#             outputs = model(inputs)\n#             _, predicted = outputs.max(1)\n\n#             for i in range(len(video_ids_batch)):\n#                 vid = int(video_ids_batch[i])\n#                 aid = int(agent_ids_batch[i])\n#                 tid = int(target_ids_batch[i])\n#                 start = int(start_frames[i])\n#                 mid_frame = start + cfg.window_size // 2\n#                 label_idx = predicted[i].item()\n#                 predictions[(vid, aid, tid)].append((mid_frame, label_idx))\n\n#     records = []\n\n#     # ✅ CHỈ DỰ ĐOÁN NẾU (agent_str, target_str, action) CÓ TRONG behaviors_labeled\n#     for (video_id, agent_id, target_id), preds in predictions.items():\n#         if len(preds) == 0:\n#             continue\n\n#         allowed_triples = video_to_allowed_triples.get(video_id, set())\n\n#         # Chuyển agent_id, target_id (int) → string có tiền tố \"mouse\"\n#         agent_str = f\"mouse{agent_id}\"\n#         target_str = f\"mouse{target_id}\"\n\n#         preds = sorted(preds)\n#         current_label = preds[0][1]\n#         current_action = cfg.idx_to_class[current_label]\n\n#         # ✅ Kiểm tra triple đầy đủ: (agent_str, target_str, action)\n#         current_triple = (agent_str, target_str, current_action)\n#         if current_triple not in allowed_triples:\n#             current_label = -1  # skip\n\n#         start_frame = preds[0][0]\n\n#         for i in range(1, len(preds)):\n#             frame, label = preds[i]\n#             action = cfg.idx_to_class[label]\n#             triple = (agent_str, target_str, action)\n\n#             # ✅ Nếu thay đổi label, không liên tục, hoặc triple không được phép → đóng segment cũ\n#             if label != current_label or frame != preds[i-1][0] + 1 or triple not in allowed_triples:\n#                 if current_label != -1:\n#                     # ✅ Kiểm tra lần cuối trước khi ghi\n#                     final_triple = (agent_str, target_str, cfg.idx_to_class[current_label])\n#                     if final_triple in allowed_triples:\n#                         records.append((\n#                             video_id,\n#                             agent_str, target_str,\n#                             cfg.idx_to_class[current_label],\n#                             start_frame, frame\n#                         ))\n#                 # Bắt đầu segment mới — chỉ nếu triple được phép\n#                 current_label = label if triple in allowed_triples else -1\n#                 start_frame = frame\n\n#         # ✅ Đóng segment cuối cùng — kiểm tra triple\n#         if current_label != -1:\n#             final_triple = (agent_str, target_str, cfg.idx_to_class[current_label])\n#             if final_triple in allowed_triples:\n#                 records.append((\n#                     video_id,\n#                     agent_str, target_str,\n#                     cfg.idx_to_class[current_label],\n#                     start_frame, preds[-1][0] + 1\n#                 ))\n\n#     if len(records) == 0:\n#         logger.warning(\"❌ No valid predictions generated!\")\n#         # Thử tìm một hành vi hợp lệ bất kỳ để dùng làm mặc định\n#         found = False\n#         for vid, triples in video_to_allowed_triples.items():\n#             if triples:\n#                 agent_str, target_str, action = next(iter(triples))\n#                 records = [(vid, agent_str, target_str, action, 0, 1)]\n#                 found = True\n#                 break\n#         if not found:\n#             default_label = list(cfg.class_to_idx.keys())[0]\n#             records = [(0, \"mouse1\", \"mouse2\", default_label, 0, 1)]\n\n#     submission_df = pl.DataFrame(\n#         records,\n#         schema={\n#             \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n#             \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n#         },\n#         orient=\"row\"\n#     )\n\n#     submission_df = submission_df.with_row_index(\"row_id\")\n#     submission_df.write_csv(cfg.submission_file)\n#     logger.info(f\"✅ Official Submission saved to {cfg.submission_file}\")\n#     logger.info(f\"📊 Submission shape: {submission_df.shape}\")\n#     print(\"\\n📋 Sample predictions:\")\n#     print(submission_df.head(10))\n\n       \n\n\n# def main():\n#     parser = argparse.ArgumentParser()\n#     parser.add_argument(\"--mode\", choices=[\"train\", \"submit\", \"all\"], default=\"all\")\n#     args = parser.parse_args()\n\n#     cfg = Config()\n\n#     if args.mode in (\"train\", \"all\"):\n#         train_model(cfg)\n\n#     if args.mode in (\"submit\", \"all\"):\n#         predict_and_submit(cfg)\n\n# if __name__ == \"__main__\":\n#     main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T02:55:58.887753Z","iopub.execute_input":"2025-09-22T02:55:58.888108Z","iopub.status.idle":"2025-09-22T02:55:58.902418Z","shell.execute_reply.started":"2025-09-22T02:55:58.888084Z","shell.execute_reply":"2025-09-22T02:55:58.901732Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %%writefile stgcn_mabe_v2_final.py\n# import os\n# import sys\n# import json\n# import warnings\n# from pathlib import Path\n# from typing import Dict, List, Tuple, Optional, Any, DefaultDict\n# from collections import defaultdict\n# import argparse\n\n# import numpy as np\n# import polars as pl\n# import torch\n# import torch.nn as nn\n# import torch.nn.functional as F\n# from torch.utils.data import Dataset, DataLoader, random_split\n# from tqdm.auto import tqdm\n# from sklearn.metrics import f1_score\n\n# print(\"\\n\" + \"=\"*70)\n# print(\">>> 🧪 SYSTEM CUDA & GPU CHECK (BEFORE ANYTHING ELSE)\")\n# print(f\"PyTorch version: {torch.__version__}\")\n# print(f\"CUDA available: {torch.cuda.is_available()}\")\n# if torch.cuda.is_available():\n#     print(f\"Number of GPUs: {torch.cuda.device_count()}\")\n#     for i in range(torch.cuda.device_count()):\n#         print(f\"  → GPU {i}: {torch.cuda.get_device_name(i)}\")\n# else:\n#     print(\"  → NO GPU AVAILABLE. Training on CPU.\")\n# print(\"<<<\")\n# print(\"=\"*70 + \"\\n\")\n\n# class Config:\n#     def __init__(self):\n#         self.data_root = Path(os.getenv(\"MABE_DATA_ROOT\", \"/kaggle/input/MABe-mouse-behavior-detection\"))\n#         self.submission_file = \"submission.csv\"\n#         self.model_save_path = \"stgcn_v2_model.pth\"\n#         self.window_size = 64\n#         self.stride = 32\n#         self.batch_size = 16\n#         self.num_epochs = 5\n#         self.lr = 0.001\n#         self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#         self.num_mice = 4\n#         self.in_channels = 5  # x, y, vx, vy, angle\n#         self.embed_dim = 32\n#         self.lab_embed_dim = 16\n\n#         # Thu thập tất cả hành vi\n#         self.class_to_idx = self._collect_all_behaviors()\n#         self.num_classes = len(self.class_to_idx)\n#         self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}\n\n#         # Thu thập tất cả lab_id để tạo lab embedding\n#         self.lab_to_idx = self._collect_all_labs()\n#         self.num_labs = len(self.lab_to_idx)\n\n#         print(f\"✅ Found {self.num_classes} unique behaviors: {list(self.class_to_idx.keys())}\")\n#         print(f\"✅ Found {self.num_labs} unique labs: {list(self.lab_to_idx.keys())}\")\n#         print(f\"🖥️  Using device: {self.device}\")\n\n#     def _collect_all_behaviors(self) -> Dict[str, int]:\n#         train_csv_path = self.data_root / \"train.csv\"\n#         if not train_csv_path.exists():\n#             return {\"avoid\": 0}\n\n#         df = pl.read_csv(train_csv_path)\n#         behaviors = set()\n\n#         for row in df.to_dicts():\n#             lab_id = row[\"lab_id\"]\n#             video_id = row[\"video_id\"]\n#             annot_path = self.data_root / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n#             if not annot_path.exists():\n#                 continue\n#             try:\n#                 annot_df = pl.read_parquet(annot_path)\n#                 if \"action\" in annot_df.columns:\n#                     behaviors.update(annot_df[\"action\"].unique().to_list())\n#             except Exception as e:\n#                 print(f\"⚠️  Failed to read {annot_path}: {e}\")\n#                 continue\n\n#         sorted_behaviors = sorted(list(behaviors))\n#         return {behavior: idx for idx, behavior in enumerate(sorted_behaviors)}\n\n#     def _collect_all_labs(self) -> Dict[str, int]:\n#         train_csv_path = self.data_root / \"train.csv\"\n#         if not train_csv_path.exists():\n#             return {\"unknown\": 0}\n#         df = pl.read_csv(train_csv_path)\n#         labs = df[\"lab_id\"].unique().to_list()\n#         return {lab: idx for idx, lab in enumerate(labs)}\n\n#     @property\n#     def train_csv(self): return self.data_root / \"train.csv\"\n#     @property\n#     def test_csv(self): return self.data_root / \"test.csv\"\n#     @property\n#     def train_annot_dir(self): return self.data_root / \"train_annotation\"\n#     @property\n#     def train_track_dir(self): return self.data_root / \"train_tracking\"\n#     @property\n#     def test_track_dir(self): return self.data_root / \"test_tracking\"\n\n# warnings.filterwarnings(\"ignore\")\n\n# # ✅ Hàm safe_json_loads - tối ưu cho data của bạn\n# def safe_json_loads(s: Optional[str]) -> List[str]:\n#     if s is None: return []\n#     if isinstance(s, list): return [str(x) for x in s]\n#     if not isinstance(s, str): return []\n#     s = s.strip()\n#     if not s: return []\n#     try:\n#         loaded = json.loads(s)\n#         if isinstance(loaded, list):\n#             return [str(item) for item in loaded]\n#         else:\n#             return []\n#     except Exception:\n#         try:\n#             fixed = s.replace(\"'\", '\"')\n#             loaded = json.loads(fixed)\n#             if isinstance(loaded, list):\n#                 return [str(item) for item in loaded]\n#             else:\n#                 return []\n#         except Exception:\n#             return []\n\n# class SpatialGraphConv(nn.Module):\n#     def __init__(self, in_channels, out_channels, A):\n#         super().__init__()\n#         self.A = nn.Parameter(A, requires_grad=True)  # ← ✅ Learnable adjacency\n#         self.conv = nn.Conv1d(in_channels, out_channels * A.size(0), kernel_size=1)\n\n#     def forward(self, x):\n#         x = self.conv(x)\n#         n, kc, v = x.size()\n#         x = x.view(n, self.A.size(0), kc // self.A.size(0), v)\n#         x = torch.einsum('nkcv,vw->nwc', x, self.A)\n#         return x.contiguous()\n\n# class STGCNBlock(nn.Module):\n#     def __init__(self, in_channels, out_channels, A, temporal_kernel_size=9):\n#         super().__init__()\n#         self.sgc = SpatialGraphConv(in_channels, out_channels, A)\n#         self.tcn = nn.Sequential(\n#             nn.InstanceNorm2d(out_channels),  # ← ✅ InstanceNorm thay BatchNorm\n#             nn.ReLU(),\n#             nn.Conv2d(out_channels, out_channels, (temporal_kernel_size, 1), padding=(temporal_kernel_size//2, 0)),\n#             nn.InstanceNorm2d(out_channels),\n#             nn.Dropout(0.5)\n#         )\n#         self.relu = nn.ReLU()\n\n#     def forward(self, x):\n#         N, C, T, V = x.size()\n#         x = x.permute(0, 2, 1, 3).contiguous().view(N * T, C, V)\n#         x = self.sgc(x)\n#         x = x.view(N, T, -1, V).permute(0, 2, 1, 3).contiguous()\n#         x = self.tcn(x)\n#         return self.relu(x)\n\n# class MABeSTGCNv2(nn.Module):\n#     def __init__(self, num_classes, num_mice=4, in_channels=5, embed_dim=32, lab_embed_dim=16, num_labs=10):\n#         super().__init__()\n#         self.num_mice = num_mice\n#         self.in_channels = in_channels\n#         self.embed_dim = embed_dim\n#         self.lab_embed_dim = lab_embed_dim\n\n#         # Adaptive adjacency matrix (learnable)\n#         A = torch.ones(num_mice, num_mice)\n#         self.A = nn.Parameter(A, requires_grad=True)\n\n#         # ST-GCN Blocks\n#         self.block1 = STGCNBlock(in_channels, 64, self.A)\n#         self.block2 = STGCNBlock(64, 128, self.A)\n#         self.block3 = STGCNBlock(128, 256, self.A)\n\n#         # Pair embedding: 4x4 = 16 pairs\n#         self.pair_embedding = nn.Embedding(num_mice * num_mice, embed_dim)\n\n#         # Lab embedding\n#         self.lab_embedding = nn.Embedding(num_labs, lab_embed_dim)\n\n#         # Final classifier\n#         self.fc = nn.Sequential(\n#             nn.Linear(256 * num_mice + embed_dim + lab_embed_dim, 512),\n#             nn.ReLU(),\n#             nn.Dropout(0.5),\n#             nn.Linear(512, num_classes)\n#         )\n\n#     def forward(self, x, agent_id, target_id, lab_id):\n#         # x: (N, T, M, C) — C=5: x,y,vx,vy,angle\n#         N, T, M, C = x.size()\n#         x = x.permute(0, 3, 1, 2).contiguous()  # (N, C, T, M)\n\n#         x = self.block1(x)\n#         x = self.block2(x)\n#         x = self.block3(x)  # (N, 256, T, M)\n\n#         # Global average pooling over time\n#         x = x.mean(dim=2)  # (N, 256, M)\n#         x = x.view(N, -1)  # (N, 256*M)\n\n#         # Pair embedding\n#         pair_idx = agent_id * self.num_mice + target_id  # (N,)\n#         pair_emb = self.pair_embedding(pair_idx)  # (N, embed_dim)\n\n#         # Lab embedding\n#         lab_emb = self.lab_embedding(lab_id)  # (N, lab_embed_dim)\n\n#         # Concat all\n#         x = torch.cat([x, pair_emb, lab_emb], dim=1)  # (N, 256*M + embed_dim + lab_embed_dim)\n#         return self.fc(x)\n\n# class MABeDataset(Dataset):\n#     def __init__(self, video_ids: List[int], lab_ids: List[str], track_dir: Path, annot_dir: Path,\n#                  window_size: int = 64, stride: int = 32, class_to_idx: Dict[str, int] = None,\n#                  lab_to_idx: Dict[str, int] = None, is_test: bool = False, num_mice: int = 4):\n#         self.video_ids = video_ids\n#         self.lab_ids = lab_ids\n#         self.track_dir = track_dir\n#         self.annot_dir = annot_dir\n#         self.window_size = window_size\n#         self.stride = stride\n#         self.class_to_idx = class_to_idx or {}\n#         self.lab_to_idx = lab_to_idx or {}\n#         self.is_test = is_test\n#         self.num_mice = num_mice\n#         self.samples = []\n#         self.track_cache = {}\n#         self.feature_cache = {}  # cache xy + velocity + angle\n\n#         self._build_samples()\n\n#     def _get_track_path(self, lab_id: str, video_id: int) -> Optional[Path]:\n#         path = self.track_dir / lab_id / f\"{video_id}.parquet\"\n#         if path.exists():\n#             return path\n#         return None\n\n#     def _build_samples(self):\n#         total_video = len(self.video_ids)\n#         print(f\"🚀 Starting to process {total_video} videos...\")\n\n#         count_track = 0\n#         count_parquet = 0\n#         count_valid_labels = 0\n\n#         for idx, (video_id, lab_id) in enumerate(tqdm(zip(self.video_ids, self.lab_ids), total=total_video, desc=\"📂 Loading & Processing Videos\")):\n#             track_path = self._get_track_path(lab_id, video_id)\n#             if track_path is None:\n#                 continue\n#             count_track += 1\n\n#             try:\n#                 track_df = pl.read_parquet(track_path)\n#                 self.track_cache[(lab_id, video_id)] = track_df\n#             except Exception as e:\n#                 print(f\"⚠️  Failed to cache tracking for {lab_id}/{video_id}: {e}\")\n#                 continue\n\n#             required_cols = [\"video_frame\", \"mouse_id\", \"x\", \"y\"]\n#             if not all(col in track_df.columns for col in required_cols):\n#                 print(f\"⚠️  Missing required columns in {track_path}\")\n#                 continue\n\n#             # Lấy danh sách mouse_id thực tế — CHUYỂN VỀ INT\n#             unique_mice_raw = track_df[\"mouse_id\"].unique().to_list()\n#             unique_mice = []\n#             for m in unique_mice_raw:\n#                 try:\n#                     mid = int(m)\n#                     unique_mice.append(mid)\n#                 except (ValueError, TypeError):\n#                     continue\n\n#             unique_mice = sorted(set(unique_mice))[:self.num_mice]\n#             if len(unique_mice) == 0:\n#                 continue\n\n#             mouse_to_idx = {mouse: i for i, mouse in enumerate(unique_mice)}\n#             all_frames = sorted(track_df[\"video_frame\"].unique().to_list())\n#             T = len(all_frames)\n#             frame_to_idx = {frame: i for i, frame in enumerate(all_frames)}\n\n#             # Khởi tạo mảng xy\n#             xy = np.full((T, self.num_mice, 2), np.nan, dtype=np.float32)\n\n#             for row in track_df.to_dicts():\n#                 frame = row[\"video_frame\"]\n#                 mouse_id = row[\"mouse_id\"]\n#                 try:\n#                     mouse_id = int(mouse_id)\n#                 except (ValueError, TypeError):\n#                     continue\n#                 x = row[\"x\"]\n#                 y = row[\"y\"]\n#                 if frame in frame_to_idx and mouse_id in mouse_to_idx:\n#                     t = frame_to_idx[frame]\n#                     m = mouse_to_idx[mouse_id]\n#                     xy[t, m, 0] = x\n#                     xy[t, m, 1] = y\n\n#             # Interpolation forward/backward\n#             for m in range(self.num_mice):\n#                 for t in range(1, T):\n#                     if np.isnan(xy[t, m, 0]):\n#                         xy[t, m] = xy[t-1, m]\n#                 for t in range(T-2, -1, -1):\n#                     if np.isnan(xy[t, m, 0]):\n#                         xy[t, m] = xy[t+1, m]\n\n#             # Tính centroid và chuẩn hóa\n#             centroid = np.nanmean(xy, axis=1, keepdims=True)\n#             xy_norm = xy - centroid\n#             xy_norm = np.nan_to_num(xy_norm, nan=0.0)\n\n#             # Tính velocity (sai phân)\n#             vx = np.zeros_like(xy_norm[..., 0])\n#             vy = np.zeros_like(xy_norm[..., 1])\n#             vx[1:] = xy_norm[1:, :, 0] - xy_norm[:-1, :, 0]\n#             vy[1:] = xy_norm[1:, :, 1] - xy_norm[:-1, :, 1]\n\n#             # Tính angle (giả sử: angle = arctan2(vy, vx))\n#             angle = np.arctan2(vy, vx)  # shape (T, M)\n\n#             # Ghép thành feature 5D: [x, y, vx, vy, angle]\n#             features = np.stack([xy_norm[..., 0], xy_norm[..., 1], vx, vy, angle], axis=-1)  # (T, M, 5)\n\n#             self.feature_cache[(lab_id, video_id)] = features\n\n#             if not self.is_test:\n#                 annot_path = self.annot_dir / lab_id / f\"{video_id}.parquet\"\n#                 if not annot_path.exists():\n#                     continue\n#                 count_parquet += 1\n\n#                 try:\n#                     annot_df = pl.read_parquet(annot_path)\n#                 except Exception as e:\n#                     print(f\"⚠️  Failed to read annotation {annot_path}: {e}\")\n#                     continue\n\n#                 if \"agent_id\" not in annot_df.columns or \"action\" not in annot_df.columns:\n#                     continue\n\n#                 unique_pairs = annot_df.select([\"agent_id\", \"target_id\"]).unique().to_dicts()\n#                 if len(unique_pairs) == 0:\n#                     continue\n\n#                 for pair in unique_pairs:\n#                     agent_id = pair[\"agent_id\"]\n#                     target_id = pair[\"target_id\"]\n#                     try:\n#                         agent_id = int(agent_id)\n#                         target_id = int(target_id)\n#                     except (ValueError, TypeError):\n#                         continue\n\n#                     if agent_id not in mouse_to_idx or target_id not in mouse_to_idx:\n#                         continue\n\n#                     pair_df = annot_df.filter((pl.col(\"agent_id\") == agent_id) & (pl.col(\"target_id\") == target_id))\n#                     if len(pair_df) == 0:\n#                         continue\n\n#                     labels = np.full(T, -1, dtype=np.int64)\n#                     for row in pair_df.to_dicts():\n#                         s, e, action = row[\"start_frame\"], row[\"stop_frame\"], row[\"action\"]\n#                         if action in self.class_to_idx:\n#                             labels[s:e+1] = self.class_to_idx[action]  # ⚠️ Sửa: e+1 để bao gồm stop_frame\n\n#                     for start in range(0, T - self.window_size + 1, self.stride):\n#                         window_labels = labels[start:start + self.window_size]\n#                         if np.any(window_labels != -1):\n#                             mid = start + self.window_size // 2\n#                             label = labels[mid] if labels[mid] != -1 else -1\n#                             if label != -1:\n#                                 self.samples.append((video_id, lab_id, agent_id, target_id, start, label))\n#                                 count_valid_labels += 1\n#             else:\n#                 # TEST SET: tạo sample cho tất cả cặp (agent, target) có thể\n#                 for agent_id in unique_mice:\n#                     for target_id in unique_mice:\n#                         for start in range(0, T - self.window_size + 1, self.stride):\n#                             self.samples.append((video_id, lab_id, agent_id, target_id, start, -1))\n\n#         print(f\"✅ Found {count_track} .parquet files in train_tracking/\")\n#         if not self.is_test:\n#             print(f\"✅ Found {count_parquet} .parquet files in train_annotation/\")\n#         print(f\"✅ Created {len(self.samples)} samples\")\n\n#     def __len__(self):\n#         return len(self.samples)\n\n#     def __getitem__(self, idx):\n#         video_id, lab_id, agent_id, target_id, start, label = self.samples[idx]\n#         features = self.feature_cache[(lab_id, video_id)]  # (T, M, 5)\n#         window_feat = features[start:start + self.window_size]  # (window, M, 5)\n#         feat_tensor = torch.tensor(window_feat, dtype=torch.float32)\n#         label_tensor = torch.tensor(label, dtype=torch.long)\n#         agent_tensor = torch.tensor(agent_id, dtype=torch.long)\n#         target_tensor = torch.tensor(target_id, dtype=torch.long)\n#         lab_idx = self.lab_to_idx.get(lab_id, 0)\n#         lab_tensor = torch.tensor(lab_idx, dtype=torch.long)\n#         return feat_tensor, label_tensor, agent_tensor, target_tensor, lab_tensor, video_id, lab_id, start\n\n# def train_model(cfg: Config):\n#     print(\"\\n🚀 [OFFICIAL TRAINING START] — Using FULL dataset\\n\")\n\n#     train_df = pl.read_csv(cfg.train_csv)\n#     video_ids = train_df[\"video_id\"].to_list()\n#     lab_ids = train_df[\"lab_id\"].to_list()\n\n#     dataset = MABeDataset(\n#         video_ids=video_ids,\n#         lab_ids=lab_ids,\n#         track_dir=cfg.train_track_dir,\n#         annot_dir=cfg.train_annot_dir,\n#         window_size=cfg.window_size,\n#         stride=cfg.stride,\n#         class_to_idx=cfg.class_to_idx,\n#         lab_to_idx=cfg.lab_to_idx,\n#         is_test=False,\n#         num_mice=cfg.num_mice\n#     )\n\n#     if len(dataset) == 0:\n#         raise ValueError(\"❌ Dataset is empty!\")\n\n#     print(f\"✅ Total training samples: {len(dataset)}\")\n\n#     train_size = max(1, int(0.8 * len(dataset)))\n#     val_size = len(dataset) - train_size\n#     if val_size == 0:\n#         val_size = 1\n#         train_size = len(dataset) - 1\n\n#     train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n\n#     try:\n#         train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True)\n#         val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)\n#     except Exception as e:\n#         print(f\"⚠️  Failed to use num_workers=2: {e}\")\n#         train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=0, pin_memory=True)\n#         val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n\n#     model = MABeSTGCNv2(\n#         num_classes=cfg.num_classes,\n#         num_mice=cfg.num_mice,\n#         in_channels=cfg.in_channels,\n#         embed_dim=cfg.embed_dim,\n#         lab_embed_dim=cfg.lab_embed_dim,\n#         num_labs=cfg.num_labs\n#     ).to(cfg.device)\n\n#     print(f\"\\n🧠 MODEL DEVICE: {next(model.parameters()).device}\\n\")\n\n#     criterion = nn.CrossEntropyLoss()\n#     optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)\n\n#     best_val_f1 = 0.0\n\n#     for epoch in range(cfg.num_epochs):\n#         model.train()\n#         train_loss = 0.0\n#         train_correct = 0\n#         train_total = 0\n\n#         pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{cfg.num_epochs} [Train]\")\n#         for batch_idx, (inputs, labels, agent_ids, target_ids, lab_ids, _, _, _) in enumerate(pbar):\n#             inputs = inputs.to(cfg.device, non_blocking=True)\n#             labels = labels.to(cfg.device, non_blocking=True)\n#             agent_ids = agent_ids.to(cfg.device, non_blocking=True)\n#             target_ids = target_ids.to(cfg.device, non_blocking=True)\n#             lab_ids = lab_ids.to(cfg.device, non_blocking=True)\n\n#             optimizer.zero_grad()\n#             outputs = model(inputs, agent_ids, target_ids, lab_ids)\n#             loss = criterion(outputs, labels)\n#             loss.backward()\n#             optimizer.step()\n\n#             train_loss += loss.item()\n#             _, predicted = outputs.max(1)\n#             train_total += labels.size(0)\n#             train_correct += predicted.eq(labels).sum().item()\n\n#             pbar.set_postfix({'Loss': loss.item(), 'Acc': train_correct/train_total})\n\n#         # ✅ Validation with F1 Score\n#         model.eval()\n#         val_preds = []\n#         val_targets = []\n\n#         with torch.no_grad():\n#             for inputs, labels, agent_ids, target_ids, lab_ids, _, _, _ in tqdm(val_loader, desc=\"Validation\"):\n#                 inputs = inputs.to(cfg.device, non_blocking=True)\n#                 labels = labels.to(cfg.device, non_blocking=True)\n#                 agent_ids = agent_ids.to(cfg.device, non_blocking=True)\n#                 target_ids = target_ids.to(cfg.device, non_blocking=True)\n#                 lab_ids = lab_ids.to(cfg.device, non_blocking=True)\n\n#                 outputs = model(inputs, agent_ids, target_ids, lab_ids)\n#                 _, predicted = outputs.max(1)\n\n#                 # Chỉ tính F1 cho các mẫu có nhãn hợp lệ (label != -1)\n#                 mask = labels != -1\n#                 if mask.sum() > 0:\n#                     val_preds.extend(predicted[mask].cpu().numpy())\n#                     val_targets.extend(labels[mask].cpu().numpy())\n\n#         if len(val_preds) > 0 and len(val_targets) > 0:\n#             val_f1 = f1_score(val_targets, val_preds, average='macro')\n#             print(f\"🎯 Epoch {epoch+1}: Validation F1 Score = {val_f1:.4f}\")\n#         else:\n#             val_f1 = 0.0\n#             print(f\"⚠️  Epoch {epoch+1}: No valid validation samples!\")\n\n#         if val_f1 > best_val_f1:\n#             best_val_f1 = val_f1\n#             torch.save(model.state_dict(), cfg.model_save_path)\n#             print(f\"✅ Model saved with Val F1 = {val_f1:.4f}\")\n\n#     print(f\"🎉 Training completed. Best Validation F1: {best_val_f1:.4f}\")\n\n# def predict_and_submit(cfg: Config):\n#     print(\"Loading test.csv...\")\n#     test_df = pl.read_csv(cfg.test_csv)\n#     video_ids = test_df[\"video_id\"].to_list()\n#     lab_ids = test_df[\"lab_id\"].to_list()\n\n#     # ✅ Tạo mapping: video_id → set of allowed (agent_str, target_str, action)\n#     video_to_allowed_triples = {}\n\n#     for row in test_df.to_dicts():\n#         video_id = row[\"video_id\"]\n#         behaviors_str = row[\"behaviors_labeled\"]\n#         behaviors_list = safe_json_loads(behaviors_str)\n#         allowed_triples = set()\n#         for item in behaviors_list:\n#             if isinstance(item, str):\n#                 parts = [p.strip() for p in item.split(\",\")]\n#                 if len(parts) == 3:\n#                     agent_str, target_str, action = parts\n#                     allowed_triples.add((agent_str, target_str, action))\n#         video_to_allowed_triples[video_id] = allowed_triples\n\n#     print(\"Creating test dataset...\")\n#     test_dataset = MABeDataset(\n#         video_ids=video_ids,\n#         lab_ids=lab_ids,\n#         track_dir=cfg.test_track_dir,\n#         annot_dir=cfg.train_annot_dir,\n#         window_size=cfg.window_size,\n#         stride=cfg.stride,\n#         class_to_idx=cfg.class_to_idx,\n#         lab_to_idx=cfg.lab_to_idx,\n#         is_test=True,\n#         num_mice=cfg.num_mice\n#     )\n\n#     if len(test_dataset) == 0:\n#         raise ValueError(\"❌ Test dataset is empty!\")\n\n#     print(f\"✅ Total test samples: {len(test_dataset)}\")\n\n#     try:\n#         test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)\n#     except Exception as e:\n#         print(f\"⚠️  Failed to use num_workers=2: {e}\")\n#         test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n\n#     model = MABeSTGCNv2(\n#         num_classes=cfg.num_classes,\n#         num_mice=cfg.num_mice,\n#         in_channels=cfg.in_channels,\n#         embed_dim=cfg.embed_dim,\n#         lab_embed_dim=cfg.lab_embed_dim,\n#         num_labs=cfg.num_labs\n#     ).to(cfg.device)\n#     model.load_state_dict(torch.load(cfg.model_save_path, map_location=cfg.device))\n#     model.eval()\n\n#     predictions = defaultdict(list)\n\n#     with torch.no_grad():\n#         for batch in tqdm(test_loader, desc=\"Inference\"):\n#             inputs, _, agent_ids, target_ids, lab_ids, video_ids_batch, lab_ids_str, start_frames = batch\n#             inputs = inputs.to(cfg.device, non_blocking=True)\n#             agent_ids = agent_ids.to(cfg.device, non_blocking=True)\n#             target_ids = target_ids.to(cfg.device, non_blocking=True)\n#             lab_ids = lab_ids.to(cfg.device, non_blocking=True)\n\n#             outputs = model(inputs, agent_ids, target_ids, lab_ids)\n#             _, predicted = outputs.max(1)\n\n#             for i in range(len(video_ids_batch)):\n#                 vid = int(video_ids_batch[i].item()) if torch.is_tensor(video_ids_batch[i]) else int(video_ids_batch[i])\n#                 aid = int(agent_ids[i].item()) if torch.is_tensor(agent_ids[i]) else int(agent_ids[i])\n#                 tid = int(target_ids[i].item()) if torch.is_tensor(target_ids[i]) else int(target_ids[i])\n#                 start = int(start_frames[i].item()) if torch.is_tensor(start_frames[i]) else int(start_frames[i])\n#                 mid_frame = start + cfg.window_size // 2\n#                 label_idx = predicted[i].item()\n#                 predictions[(vid, aid, tid)].append((mid_frame, label_idx))\n\n#     records = []\n\n#     for (video_id, agent_id, target_id), preds in predictions.items():\n#         if len(preds) == 0:\n#             continue\n\n#         allowed_triples = video_to_allowed_triples.get(video_id, set())\n#         agent_str = f\"mouse{agent_id}\"\n#         target_str = f\"mouse{target_id}\"\n\n#         # ✅ Lọc chỉ giữ lại các dự đoán có triple hợp lệ\n#         filtered_preds = []\n#         for frame, label_idx in preds:\n#             if label_idx == -1 or label_idx >= cfg.num_classes:\n#                 continue\n#             action = cfg.idx_to_class[label_idx]\n#             triple = (agent_str, target_str, action)\n#             if triple in allowed_triples:\n#                 filtered_preds.append((frame, label_idx, action))\n\n#         if len(filtered_preds) == 0:\n#             continue\n\n#         # ✅ Tạo segment từ filtered_preds\n#         filtered_preds.sort(key=lambda x: x[0])  # sort by frame\n#         current_label = filtered_preds[0][1]\n#         current_action = filtered_preds[0][2]\n#         start_frame = filtered_preds[0][0]\n\n#         for i in range(1, len(filtered_preds)):\n#             frame, label, action = filtered_preds[i]\n#             if label != current_label or frame != filtered_preds[i-1][0] + 1:\n#                 # Đóng segment cũ\n#                 records.append((\n#                     video_id,\n#                     agent_str, target_str,\n#                     cfg.idx_to_class[current_label],\n#                     start_frame, frame\n#                 ))\n#                 # Mở segment mới\n#                 current_label = label\n#                 start_frame = frame\n\n#         # Đóng segment cuối\n#         records.append((\n#             video_id,\n#             agent_str, target_str,\n#             cfg.idx_to_class[current_label],\n#             start_frame, filtered_preds[-1][0] + 1\n#         ))\n\n#     if len(records) == 0:\n#         print(\"❌ No valid predictions generated!\")\n#         found = False\n#         for vid, triples in video_to_allowed_triples.items():\n#             if triples:\n#                 agent_str, target_str, action = next(iter(triples))\n#                 records = [(vid, agent_str, target_str, action, 0, 1)]\n#                 found = True\n#                 break\n#         if not found:\n#             default_label = list(cfg.class_to_idx.keys())[0]\n#             records = [(0, \"mouse1\", \"mouse2\", default_label, 0, 1)]\n\n#     submission_df = pl.DataFrame(\n#         records,\n#         schema={\n#             \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n#             \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n#         },\n#         orient=\"row\"\n#     )\n\n#     submission_df = submission_df.with_row_index(\"row_id\")\n#     submission_df.write_csv(cfg.submission_file)\n#     print(f\"✅ Official Submission saved to {cfg.submission_file}\")\n#     print(f\"📊 Submission shape: {submission_df.shape}\")\n#     print(\"\\n📋 Sample predictions:\")\n#     print(submission_df.head(10))\n\n# def main():\n#     parser = argparse.ArgumentParser()\n#     parser.add_argument(\"--mode\", choices=[\"train\", \"submit\", \"all\"], default=\"all\")\n#     args = parser.parse_args()\n\n#     cfg = Config()\n\n#     if args.mode in (\"train\", \"all\"):\n#         train_model(cfg)\n\n#     if args.mode in (\"submit\", \"all\"):\n#         predict_and_submit(cfg)\n\n# if __name__ == \"__main__\":\n#     main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T06:00:29.157566Z","iopub.execute_input":"2025-09-22T06:00:29.158231Z","iopub.status.idle":"2025-09-22T06:00:29.184222Z","shell.execute_reply.started":"2025-09-22T06:00:29.158198Z","shell.execute_reply":"2025-09-22T06:00:29.183385Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile stgcn_mabe_final_official_fixed_v8.py\nimport os\nimport sys\nimport json\nimport logging\nimport warnings\nfrom pathlib import Path\nfrom typing import Dict, List, Tuple, Optional, Any, DefaultDict\nfrom collections import defaultdict\nimport argparse\n\nimport numpy as np\nimport polars as pl\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom tqdm.auto import tqdm\n\nprint(\"\\n\" + \"=\"*70)\nprint(\">>> 🧪 SYSTEM CUDA & GPU CHECK (BEFORE ANYTHING ELSE)\")\nprint(f\"PyTorch version: {torch.__version__}\")\nprint(f\"CUDA available: {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    print(f\"Number of GPUs: {torch.cuda.device_count()}\")\n    for i in range(torch.cuda.device_count()):\n        print(f\"  → GPU {i}: {torch.cuda.get_device_name(i)}\")\nelse:\n    print(\"  → NO GPU AVAILABLE. Training on CPU.\")\nprint(\"<<<\")\nprint(\"=\"*70 + \"\\n\")\n\nclass Config:\n    def __init__(self):\n        self.data_root = Path(os.getenv(\"MABE_DATA_ROOT\", \"/kaggle/input/MABe-mouse-behavior-detection\"))\n        self.submission_file = \"submission.csv\"\n        self.model_save_path = \"stgcn_model.pth\"\n        self.window_size = 64\n        self.stride = 32\n        self.batch_size = 16\n        self.num_epochs = 3\n        self.lr = 0.001\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.num_mice = 4\n        self.in_channels = 2\n\n        self.class_to_idx = self._collect_all_behaviors()\n        self.num_classes = len(self.class_to_idx)\n        self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}\n\n        logger.info(f\"✅ Found {self.num_classes} unique behaviors: {list(self.class_to_idx.keys())}\")\n        logger.info(f\"🖥️  Using device: {self.device}\")\n\n    def _collect_all_behaviors(self) -> Dict[str, int]:\n        train_csv_path = self.data_root / \"train.csv\"\n        if not train_csv_path.exists():\n            return {\"avoid\": 0}\n\n        df = pl.read_csv(train_csv_path)\n        behaviors = set()\n\n        for row in df.to_dicts():\n            lab_id = row[\"lab_id\"]\n            video_id = row[\"video_id\"]\n            annot_path = self.data_root / \"train_annotation\" / lab_id / f\"{video_id}.parquet\"\n            if not annot_path.exists():\n                continue\n            try:\n                annot_df = pl.read_parquet(annot_path)\n                if \"action\" in annot_df.columns:\n                    behaviors.update(annot_df[\"action\"].unique().to_list())\n            except Exception as e:\n                logger.warning(f\"Failed to read {annot_path}: {e}\")\n                continue\n\n        sorted_behaviors = sorted(list(behaviors))\n        return {behavior: idx for idx, behavior in enumerate(sorted_behaviors)}\n\n    @property\n    def train_csv(self): return self.data_root / \"train.csv\"\n    @property\n    def test_csv(self): return self.data_root / \"test.csv\"\n    @property\n    def train_annot_dir(self): return self.data_root / \"train_annotation\"\n    @property\n    def train_track_dir(self): return self.data_root / \"train_tracking\"\n    @property\n    def test_track_dir(self): return self.data_root / \"test_tracking\"\n\nlogging.basicConfig(level=logging.INFO, format=\"%(asctime)s | %(levelname)s | %(message)s\")\nlogger = logging.getLogger(__name__)\nwarnings.filterwarnings(\"ignore\")\n\n# ✅ Hàm safe_json_loads - tối ưu cho data của bạn\ndef safe_json_loads(s: Optional[str]) -> List[str]:\n    if s is None: return []\n    if isinstance(s, list): return [str(x) for x in s]\n    if not isinstance(s, str): return []\n    s = s.strip()\n    if not s: return []\n    try:\n        loaded = json.loads(s)\n        if isinstance(loaded, list):\n            return [str(item) for item in loaded]\n        else:\n            return []\n    except Exception:\n        try:\n            # Thử sửa dấu nháy đơn\n            fixed = s.replace(\"'\", '\"')\n            loaded = json.loads(fixed)\n            if isinstance(loaded, list):\n                return [str(item) for item in loaded]\n            else:\n                return []\n        except Exception:\n            return []\n\nclass SpatialGraphConv(nn.Module):\n    def __init__(self, in_channels, out_channels, A):\n        super().__init__()\n        self.A = nn.Parameter(A, requires_grad=False)\n        self.conv = nn.Conv1d(in_channels, out_channels * A.size(0), kernel_size=1)\n\n    def forward(self, x):\n        x = self.conv(x)\n        n, kc, v = x.size()\n        x = x.view(n, self.A.size(0), kc // self.A.size(0), v)\n        x = torch.einsum('nvcv,vw->nwc', x, self.A)\n        return x.contiguous()\n\nclass STGCNBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, A, temporal_kernel_size=9):\n        super().__init__()\n        self.sgc = SpatialGraphConv(in_channels, out_channels, A)\n        self.tcn = nn.Sequential(\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(),\n            nn.Conv2d(out_channels, out_channels, (temporal_kernel_size, 1), padding=(temporal_kernel_size//2, 0)),\n            nn.BatchNorm2d(out_channels),\n            nn.Dropout(0.5)\n        )\n        self.relu = nn.ReLU()\n\n    def forward(self, x):\n        N, C, T, V = x.size()\n        x = x.permute(0, 2, 1, 3).contiguous().view(N * T, C, V)\n        x = self.sgc(x)\n        x = x.view(N, T, -1, V).permute(0, 2, 1, 3).contiguous()\n        x = self.tcn(x)\n        return self.relu(x)\n\nclass MABeSTGCN(nn.Module):\n    def __init__(self, num_classes, num_mice=4, in_channels=2):\n        super().__init__()\n        self.num_mice = num_mice\n        self.total_nodes = num_mice\n        A = torch.ones(num_mice, num_mice)\n        self.A = nn.Parameter(A, requires_grad=False)\n        self.block1 = STGCNBlock(in_channels, 64, self.A)\n        self.block2 = STGCNBlock(64, 128, self.A)\n        self.block3 = STGCNBlock(128, 256, self.A)\n        self.fc = nn.Linear(256 * num_mice + num_mice * num_mice, num_classes)\n\n    def forward(self, x, pair_enc):\n        N, T, M, C = x.size()\n        x = x.permute(0, 3, 1, 2).contiguous()\n        x = self.block1(x)\n        x = self.block2(x)\n        x = self.block3(x)\n        mid_frame = T // 2\n        x = x[:, :, mid_frame, :]\n        x = x.reshape(N, -1)\n        x = torch.cat([x, pair_enc], dim=1)\n        return self.fc(x)\n\nclass MABeDataset(Dataset):\n    def __init__(self, video_ids: List[int], lab_ids: List[str], track_dir: Path, annot_dir: Path,\n                 window_size: int = 64, stride: int = 32, class_to_idx: Dict[str, int] = None, is_test: bool = False, num_mice: int = 4):\n        self.video_ids = video_ids\n        self.lab_ids = lab_ids\n        self.track_dir = track_dir\n        self.annot_dir = annot_dir\n        self.window_size = window_size\n        self.stride = stride\n        self.class_to_idx = class_to_idx or {}\n        self.is_test = is_test\n        self.num_mice = num_mice\n        self.samples = []\n        self.track_cache = {}\n        self.xy_cache = {}\n\n        self.video_to_allowed_triples = {}\n        if not self.is_test:\n            train_meta_df = pl.read_csv(Config().train_csv)\n            for row in train_meta_df.to_dicts():\n                vid = row[\"video_id\"]\n                behaviors_str = row.get(\"behaviors_labeled\", \"\")\n                behaviors_list = safe_json_loads(behaviors_str)\n                allowed_triples = set()\n                for item in behaviors_list:\n                    parts = [p.strip() for p in str(item).replace(\"'\", \"\").split(\",\")]\n                    if len(parts) == 3:\n                        agent_str, target_str, action = parts\n                        allowed_triples.add((agent_str, target_str, action))\n                self.video_to_allowed_triples[vid] = allowed_triples\n\n        self._build_samples()\n\n    def _get_track_path(self, lab_id: str, video_id: int) -> Optional[Path]:\n        path = self.track_dir / lab_id / f\"{video_id}.parquet\"\n        if path.exists():\n            return path\n        return None\n\n    def _build_samples(self):\n        total_video = len(self.video_ids)\n        logger.info(f\"🚀 Starting to process {total_video} videos...\")\n\n        count_track = 0\n        count_parquet = 0\n        count_valid_labels = 0\n\n        for idx, (video_id, lab_id) in enumerate(tqdm(zip(self.video_ids, self.lab_ids), total=total_video, desc=\"📂 Loading & Processing Videos\")):\n            track_path = self._get_track_path(lab_id, video_id)\n            if track_path is None:\n                continue\n            count_track += 1\n\n            try:\n                track_df = pl.read_parquet(track_path)\n                self.track_cache[(lab_id, video_id)] = track_df\n            except Exception as e:\n                logger.warning(f\"Failed to cache tracking for {lab_id}/{video_id}: {e}\")\n                continue\n\n            required_cols = [\"video_frame\", \"mouse_id\", \"x\", \"y\"]\n            if not all(col in track_df.columns for col in required_cols):\n                logger.warning(f\"Missing required columns in {track_path}\")\n                continue\n\n            unique_mice_raw = track_df[\"mouse_id\"].unique().to_list()\n            unique_mice = []\n            for m in unique_mice_raw:\n                try:\n                    mid = int(m)\n                    unique_mice.append(mid)\n                except (ValueError, TypeError):\n                    logger.warning(f\"Cannot convert mouse_id {m} to int, skipping.\")\n                    continue\n\n            unique_mice = sorted(set(unique_mice))[:self.num_mice]\n            if len(unique_mice) == 0:\n                continue\n\n            mouse_to_idx = {mouse: i for i, mouse in enumerate(unique_mice)}\n            all_frames = sorted(track_df[\"video_frame\"].unique().to_list())\n            T = len(all_frames)\n            frame_to_idx = {frame: i for i, frame in enumerate(all_frames)}\n\n            xy = np.full((T, self.num_mice, 2), np.nan, dtype=np.float32)\n\n            for row in track_df.to_dicts():\n                frame = row[\"video_frame\"]\n                mouse_id = row[\"mouse_id\"]\n                try:\n                    mouse_id = int(mouse_id)\n                except (ValueError, TypeError):\n                    continue\n                x = row[\"x\"]\n                y = row[\"y\"]\n                if frame in frame_to_idx and mouse_id in mouse_to_idx:\n                    t = frame_to_idx[frame]\n                    m = mouse_to_idx[mouse_id]\n                    xy[t, m, 0] = x\n                    xy[t, m, 1] = y\n\n            for m in range(self.num_mice):\n                for t in range(1, T):\n                    if np.isnan(xy[t, m, 0]):\n                        xy[t, m] = xy[t-1, m]\n\n            for m in range(self.num_mice):\n                for t in range(T-2, -1, -1):\n                    if np.isnan(xy[t, m, 0]):\n                        xy[t, m] = xy[t+1, m]\n\n            centroid = np.nanmean(xy, axis=1, keepdims=True)\n            xy_norm = xy - centroid\n            xy_norm = np.nan_to_num(xy_norm, nan=0.0)\n            self.xy_cache[(lab_id, video_id)] = xy_norm\n\n            if not self.is_test:\n                annot_path = self.annot_dir / lab_id / f\"{video_id}.parquet\"\n                if not annot_path.exists():\n                    continue\n                count_parquet += 1\n\n                try:\n                    annot_df = pl.read_parquet(annot_path)\n                except Exception as e:\n                    logger.warning(f\"Failed to read annotation {annot_path}: {e}\")\n                    continue\n\n                if \"agent_id\" not in annot_df.columns or \"action\" not in annot_df.columns:\n                    continue\n\n                unique_pairs = annot_df.select([\"agent_id\", \"target_id\"]).unique().to_dicts()\n                if len(unique_pairs) == 0:\n                    continue\n\n                for pair in unique_pairs:\n                    agent_id = pair[\"agent_id\"]\n                    target_id = pair[\"target_id\"]\n                    try:\n                        agent_id = int(agent_id)\n                        target_id = int(target_id)\n                    except (ValueError, TypeError):\n                        continue\n\n                    pair_df = annot_df.filter((pl.col(\"agent_id\") == agent_id) & (pl.col(\"target_id\") == target_id))\n                    if len(pair_df) == 0:\n                        continue\n\n                    labels = np.full(T, -1, dtype=np.int64)\n                    for row in pair_df.to_dicts():\n                        s, e, action = row[\"start_frame\"], row[\"stop_frame\"], row[\"action\"]\n                        if action not in self.class_to_idx:\n                            continue\n\n                        agent_str = f\"mouse{agent_id}\"\n                        target_str = f\"mouse{target_id}\"\n                        triple = (agent_str, target_str, action)\n                        if triple not in self.video_to_allowed_triples.get(video_id, set()):\n                            continue\n\n                        labels[s:e] = self.class_to_idx[action]\n\n                    for start in range(0, T - self.window_size + 1, self.stride):\n                        window_labels = labels[start:start + self.window_size]\n                        valid_labels = window_labels[window_labels != -1]\n                        if len(valid_labels) > 0:\n                            label = valid_labels[0]\n                            self.samples.append((video_id, lab_id, agent_id, target_id, start, label))\n                            count_valid_labels += 1\n            else:\n                for agent_id in unique_mice:\n                    for target_id in unique_mice:\n                        for start in range(0, T - self.window_size + 1, self.stride):\n                            self.samples.append((video_id, lab_id, agent_id, target_id, start, -1))\n\n        logger.info(f\"✅ Found {count_track} .parquet files in train_tracking/\")\n        if not self.is_test:\n            logger.info(f\"✅ Found {count_parquet} .parquet files in train_annotation/\")\n        logger.info(f\"✅ Created {len(self.samples)} samples\")\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        video_id, lab_id, agent_id, target_id, start, label = self.samples[idx]\n        xy_norm = self.xy_cache[(lab_id, video_id)]\n        window_xy = xy_norm[start:start + self.window_size]\n        xy_tensor = torch.tensor(window_xy, dtype=torch.float32)\n        label_tensor = torch.tensor(label, dtype=torch.long)\n\n        try:\n            agent_idx = int(agent_id) - 1\n            target_idx = int(target_id) - 1\n        except:\n            agent_idx = 0\n            target_idx = 0\n\n        if not (0 <= agent_idx < self.num_mice and 0 <= target_idx < self.num_mice):\n            agent_idx = 0\n            target_idx = 0\n\n        pair_encoding = torch.zeros(self.num_mice * self.num_mice, dtype=torch.float32)\n        pair_encoding[agent_idx * self.num_mice + target_idx] = 1.0\n\n        return xy_tensor, pair_encoding, label_tensor, video_id, lab_id, agent_id, target_id, start\n\ndef train_model(cfg: Config):\n    print(\"\\n🚀 [OFFICIAL TRAINING START] — Using FULL dataset\\n\")\n\n    train_df = pl.read_csv(cfg.train_csv)\n    video_ids = train_df[\"video_id\"].to_list()\n    lab_ids = train_df[\"lab_id\"].to_list()\n\n    dataset = MABeDataset(\n        video_ids=video_ids,\n        lab_ids=lab_ids,\n        track_dir=cfg.train_track_dir,\n        annot_dir=cfg.train_annot_dir,\n        window_size=cfg.window_size,\n        stride=cfg.stride,\n        class_to_idx=cfg.class_to_idx,\n        is_test=False,\n        num_mice=cfg.num_mice\n    )\n\n    if len(dataset) == 0:\n        raise ValueError(\"❌ Dataset is empty!\")\n\n    logger.info(f\"✅ Total training samples: {len(dataset)}\")\n\n    all_labels = [sample[-1] for sample in dataset.samples if sample[-1] != -1]\n    if len(all_labels) == 0:\n        class_weights = None\n    else:\n        class_counts = np.bincount(all_labels, minlength=cfg.num_classes)\n        class_weights = 1.0 / (class_counts + 1e-6)\n        class_weights = torch.tensor(class_weights, dtype=torch.float32).to(cfg.device)\n\n    train_size = max(1, int(0.8 * len(dataset)))\n    val_size = len(dataset) - train_size\n    if val_size == 0:\n        val_size = 1\n        train_size = len(dataset) - 1\n\n    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n\n    try:\n        train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True)\n        val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)\n    except Exception as e:\n        logger.warning(f\"Failed to use num_workers=2: {e}\")\n        train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=0, pin_memory=True)\n        val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n\n    model = MABeSTGCN(\n        num_classes=cfg.num_classes,\n        num_mice=cfg.num_mice,\n        in_channels=cfg.in_channels\n    ).to(cfg.device)\n\n    print(\"\\n\" + \"=\"*60)\n    print(\"🧠 MODEL DEVICE:\", next(model.parameters()).device)\n    print(\"=\"*60 + \"\\n\")\n\n    criterion = nn.CrossEntropyLoss(weight=class_weights)\n    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)\n\n    best_val_acc = 0.0\n\n    for epoch in range(cfg.num_epochs):\n        model.train()\n        train_loss = 0.0\n        train_correct = 0\n        train_total = 0\n\n        pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{cfg.num_epochs} [Train]\")\n        for batch_idx, (inputs, pair_enc, labels, _, _, _, _, _) in enumerate(pbar):  # ✅ SỬA: 8 giá trị\n            inputs = inputs.to(cfg.device, non_blocking=True)\n            pair_enc = pair_enc.to(cfg.device, non_blocking=True)\n            labels = labels.to(cfg.device, non_blocking=True)\n\n            optimizer.zero_grad()\n            outputs = model(inputs, pair_enc)\n            loss = criterion(outputs, labels)\n            loss.backward()\n            optimizer.step()\n\n            train_loss += loss.item()\n            _, predicted = outputs.max(1)\n            train_total += labels.size(0)\n            train_correct += predicted.eq(labels).sum().item()\n\n            pbar.set_postfix({'Loss': loss.item(), 'Acc': train_correct/train_total})\n\n        model.eval()\n        val_correct = 0\n        val_total = 0\n        with torch.no_grad():\n            for inputs, pair_enc, labels, _, _, _, _, _ in tqdm(val_loader, desc=\"Validation\"):  # ✅ SỬA: 8 giá trị\n                inputs = inputs.to(cfg.device, non_blocking=True)\n                pair_enc = pair_enc.to(cfg.device, non_blocking=True)\n                labels = labels.to(cfg.device, non_blocking=True)\n                outputs = model(inputs, pair_enc)\n                _, predicted = outputs.max(1)\n                val_total += labels.size(0)\n                val_correct += predicted.eq(labels).sum().item()\n\n        val_acc = val_correct / val_total\n        logger.info(f\"Epoch {epoch+1}: Val Acc = {val_acc:.4f}\")\n\n        if val_acc > best_val_acc:\n            best_val_acc = val_acc\n            torch.save(model.state_dict(), cfg.model_save_path)\n            logger.info(f\"✅ Model saved with Val Acc = {val_acc:.4f}\")\n\n    logger.info(f\"🎉 Training completed. Best Val Acc: {best_val_acc:.4f}\")\n\ndef predict_and_submit(cfg: Config):\n    logger.info(\"Loading test.csv...\")\n    test_df = pl.read_csv(cfg.test_csv)\n    video_ids = test_df[\"video_id\"].to_list()\n    lab_ids = test_df[\"lab_id\"].to_list()\n\n    video_to_allowed_triples = {}\n    for row in test_df.to_dicts():\n        video_id = row[\"video_id\"]\n        behaviors_str = row[\"behaviors_labeled\"]\n        behaviors_list = safe_json_loads(behaviors_str)\n        allowed_triples = set()\n        for item in behaviors_list:\n            parts = [p.strip() for p in str(item).replace(\"'\", \"\").split(\",\")]\n            if len(parts) == 3:\n                agent_str, target_str, action = parts\n                allowed_triples.add((agent_str, target_str, action))\n        video_to_allowed_triples[video_id] = allowed_triples\n\n    logger.info(\"Creating test dataset...\")\n    test_dataset = MABeDataset(\n        video_ids=video_ids,\n        lab_ids=lab_ids,\n        track_dir=cfg.test_track_dir,\n        annot_dir=cfg.train_annot_dir,\n        window_size=cfg.window_size,\n        stride=cfg.stride,\n        class_to_idx=cfg.class_to_idx,\n        is_test=True,\n        num_mice=cfg.num_mice\n    )\n\n    if len(test_dataset) == 0:\n        raise ValueError(\"❌ Test dataset is empty!\")\n\n    logger.info(f\"✅ Total test samples: {len(test_dataset)}\")\n\n    try:\n        test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)\n    except Exception as e:\n        logger.warning(f\"Failed to use num_workers=2: {e}\")\n        test_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0, pin_memory=True)\n\n    model = MABeSTGCN(\n        num_classes=cfg.num_classes,\n        num_mice=cfg.num_mice,\n        in_channels=cfg.in_channels\n    ).to(cfg.device)\n    model.load_state_dict(torch.load(cfg.model_save_path, map_location=cfg.device))\n    model.eval()\n\n    predictions = defaultdict(list)\n\n    with torch.no_grad():\n        for batch in tqdm(test_loader, desc=\"Inference\"):\n            inputs, pair_enc, _, video_ids_batch, lab_ids_batch, agent_ids_batch, target_ids_batch, start_frames = batch\n            inputs = inputs.to(cfg.device, non_blocking=True)\n            pair_enc = pair_enc.to(cfg.device, non_blocking=True)\n            outputs = model(inputs, pair_enc)\n            _, predicted = outputs.max(1)\n\n            for i in range(len(video_ids_batch)):\n                vid = int(video_ids_batch[i])\n                aid = int(agent_ids_batch[i])\n                tid = int(target_ids_batch[i])\n                start = int(start_frames[i])\n                mid_frame = start + cfg.window_size // 2\n                label_idx = predicted[i].item()\n                predictions[(vid, aid, tid)].append((mid_frame, label_idx))\n\n    records = []\n\n    for (video_id, agent_id, target_id), preds in predictions.items():\n        if len(preds) == 0:\n            continue\n\n        allowed_triples = video_to_allowed_triples.get(video_id, set())\n        agent_str = f\"mouse{agent_id}\"\n        target_str = f\"mouse{target_id}\"\n\n        preds = sorted(preds)\n        current_label = preds[0][1]\n        current_action = cfg.idx_to_class[current_label]\n        current_triple = (agent_str, target_str, current_action)\n\n        if current_triple not in allowed_triples:\n            current_label = -1\n\n        start_frame = preds[0][0]\n\n        for i in range(1, len(preds)):\n            frame, label = preds[i]\n            action = cfg.idx_to_class[label]\n            triple = (agent_str, target_str, action)\n\n            if label != current_label or frame != preds[i-1][0] + 1 or triple not in allowed_triples:\n                if current_label != -1:\n                    final_triple = (agent_str, target_str, cfg.idx_to_class[current_label])\n                    if final_triple in allowed_triples:\n                        records.append((\n                            video_id,\n                            agent_str, target_str,\n                            cfg.idx_to_class[current_label],\n                            start_frame, frame\n                        ))\n                current_label = label if triple in allowed_triples else -1\n                start_frame = frame\n\n        if current_label != -1:\n            final_triple = (agent_str, target_str, cfg.idx_to_class[current_label])\n            if final_triple in allowed_triples:\n                records.append((\n                    video_id,\n                    agent_str, target_str,\n                    cfg.idx_to_class[current_label],\n                    start_frame, preds[-1][0] + 1\n                ))\n\n    if len(records) == 0:\n        logger.warning(\"❌ No valid predictions generated!\")\n        found = False\n        for vid, triples in video_to_allowed_triples.items():\n            if triples:\n                agent_str, target_str, action = next(iter(triples))\n                records = [(vid, agent_str, target_str, action, 0, 1)]\n                found = True\n                break\n        if not found:\n            default_label = list(cfg.class_to_idx.keys())[0]\n            records = [(0, \"mouse1\", \"mouse2\", default_label, 0, 1)]\n\n    submission_df = pl.DataFrame(\n        records,\n        schema={\n            \"video_id\": pl.Int64, \"agent_id\": pl.Utf8, \"target_id\": pl.Utf8,\n            \"action\": pl.Utf8, \"start_frame\": pl.Int64, \"stop_frame\": pl.Int64,\n        },\n        orient=\"row\"\n    )\n\n    submission_df = submission_df.with_row_index(\"row_id\")\n    submission_df.write_csv(cfg.submission_file)\n    logger.info(f\"✅ Official Submission saved to {cfg.submission_file}\")\n    logger.info(f\"📊 Submission shape: {submission_df.shape}\")\n    print(\"\\n📋 Sample predictions:\")\n    print(submission_df.head(10))\n\ndef main():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--mode\", choices=[\"train\", \"submit\", \"all\"], default=\"all\")\n    args = parser.parse_args()\n\n    cfg = Config()\n\n    if args.mode in (\"train\", \"all\"):\n        train_model(cfg)\n\n    if args.mode in (\"submit\", \"all\"):\n        predict_and_submit(cfg)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T10:16:18.366305Z","iopub.execute_input":"2025-09-22T10:16:18.367309Z","iopub.status.idle":"2025-09-22T10:16:18.383772Z","shell.execute_reply.started":"2025-09-22T10:16:18.367254Z","shell.execute_reply":"2025-09-22T10:16:18.383259Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python stgcn_mabe_final_official_fixed_v8.py --mode all","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-22T10:16:23.387820Z","iopub.execute_input":"2025-09-22T10:16:23.388333Z","iopub.status.idle":"2025-09-22T10:54:50.360928Z","shell.execute_reply.started":"2025-09-22T10:16:23.388307Z","shell.execute_reply":"2025-09-22T10:54:50.360241Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}