{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================\n# MABe Mouse Behavior Detection\n# Folder-by-folder, video-by-video training\n# Pose + Graph (across mice) + Temporal Transformer\n# ============================================================\n\nimport os\nimport gc\nimport random\nfrom collections import Counter\n\nimport numpy as np\nimport pandas as pd\nfrom sklearn.metrics import accuracy_score\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\n# ------------------------------------------------------------\n# 0. Config\n# ------------------------------------------------------------\n\nROOT = \"/kaggle/input/MABe-mouse-behavior-detection\"\nTRACK_DIR = os.path.join(ROOT, \"train_tracking\")\nANN_DIR   = os.path.join(ROOT, \"train_annotation\")\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"Using device:\", DEVICE)\n\n# Temporal window length\nT = 32\n\n# Max mice and bodyparts (truncate extras)\nM_MAX = 4\nB_MAX = 16\n\n# How many TOTAL videos to process (across all folders)\nMAX_VIDEOS = 40   # increase to use more data\n\n# Max windows (samples) per video\nMAX_SAMPLES_PER_VIDEO = 200\n\n# Seeds\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\n\nprint(\"Tracking labs:\", os.listdir(TRACK_DIR))\nprint(\"Annotation labs:\", os.listdir(ANN_DIR))\n\n# ------------------------------------------------------------\n# 1. Loading tracking & annotations\n# ------------------------------------------------------------\n\ndef load_tracking_file(path):\n    df = pd.read_parquet(path)\n    vid = int(os.path.basename(path).replace(\".parquet\", \"\"))\n    df[\"video_id\"] = vid\n    return df, vid\n\ndef load_annotation_for_video(video_id):\n    fname = f\"{video_id}.parquet\"\n    for lab in os.listdir(ANN_DIR):\n        p = os.path.join(ANN_DIR, lab, fname)\n        if os.path.exists(p):\n            df = pd.read_parquet(p)\n            df[\"video_id\"] = video_id\n            return df\n    return None\n\ndef build_pose_tensor(tracking_df):\n    \"\"\"\n    Build pose tensor: [F, M_max, B_max, 2], mask: [F, M_max, B_max]\n    F = number of frames (contiguous index)\n    M_max = up to M_MAX mice\n    B_max = up to B_MAX bodyparts\n    \"\"\"\n    frames = sorted(tracking_df[\"video_frame\"].unique().tolist())\n    frame_to_idx = {f:i for i,f in enumerate(frames)}\n    F = len(frames)\n\n    mice = sorted(tracking_df[\"mouse_id\"].unique().tolist())\n    if len(mice) == 0 or len(mice) > M_MAX:\n        return None, None, None, None  # skip videos with 0 or too many mice\n\n    mice = mice[:M_MAX]\n    mouse_to_idx = {m:i for i,m in enumerate(mice)}\n\n    bodyparts = sorted(tracking_df[\"bodypart\"].unique().tolist())\n    bodyparts = bodyparts[:B_MAX]\n    bp_to_idx = {b:i for i,b in enumerate(bodyparts)}\n    B = len(bodyparts)\n\n    pose = np.zeros((F, M_MAX, B_MAX, 2), dtype=np.float32)\n    mask = np.zeros((F, M_MAX, B_MAX), dtype=np.float32)\n\n    for _, row in tracking_df.iterrows():\n        f = frame_to_idx[row[\"video_frame\"]]\n        m_id = row[\"mouse_id\"]\n        b_name = row[\"bodypart\"]\n        if m_id not in mouse_to_idx or b_name not in bp_to_idx:\n            continue\n        mi = mouse_to_idx[m_id]\n        bi = bp_to_idx[b_name]\n        pose[f, mi, bi, 0] = row[\"x\"]\n        pose[f, mi, bi, 1] = row[\"y\"]\n        mask[f, mi, bi] = 1.0\n\n    return pose, mask, frames, mouse_to_idx\n\n# ------------------------------------------------------------\n# 2. Model (stable version)\n# ------------------------------------------------------------\n\nclass FullAGIModel(nn.Module):\n    def __init__(self, d_pose=2, d_hidden=64, num_heads=4, num_classes=10, T=T, M=M_MAX, B=B_MAX):\n        super().__init__()\n        self.T = T\n        self.M = M\n        self.B = B\n        self.d_hidden = d_hidden\n\n        # Simple bodypart embedding\n        self.bodypart_proj = nn.Linear(d_pose, d_hidden)\n\n        # Graph attention across mice per frame\n        self.node_attn = nn.MultiheadAttention(\n            embed_dim=d_hidden,\n            num_heads=num_heads,\n            batch_first=True\n        )\n\n        # Temporal Transformer over time steps\n        self.temporal_pos = nn.Parameter(torch.randn(1, T, d_hidden*3))\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=d_hidden*3,\n            nhead=num_heads,\n            dim_feedforward=d_hidden*4,\n            batch_first=True\n        )\n        self.temporal_encoder = nn.TransformerEncoder(encoder_layer, num_layers=2)\n\n        # Classification head\n        self.cls_head = nn.Sequential(\n            nn.Linear(d_hidden*3, d_hidden),\n            nn.ReLU(),\n            nn.Linear(d_hidden, num_classes)\n        )\n\n    def forward(self, pose, mask, agent_idx, target_idx):\n        \"\"\"\n        pose: [B,T,M,B,2]\n        mask: [B,T,M,B]\n        agent_idx: [B]\n        target_idx: [B]\n        \"\"\"\n        B, T_, M_, B_, Dp = pose.shape\n        assert T_ == self.T\n        assert M_ == self.M\n\n        # 1) Bodypart embedding\n        # pose -> [B*T*M, B, 2]\n        x = pose.view(B*T_*M_, B_, Dp)\n        m = mask.view(B*T_*M_, B_)   # [B*T*M,B]\n        bp_emb = self.bodypart_proj(x)  # [B*T*M,B,d_hidden]\n\n        # masked mean over bodyparts\n        m_sum = m.sum(dim=1, keepdim=True) + 1e-6  # [B*T*M,1]\n        h_mouse = (bp_emb * m.unsqueeze(-1)).sum(dim=1) / m_sum  # [B*T*M,d_hidden]\n\n        # reshape to [B,T,M,d_hidden]\n        h_mouse = h_mouse.view(B, T_, M_, self.d_hidden)\n\n        # 2) Graph attention across mice per frame\n        # flatten B*T\n        h_bt = h_mouse.view(B*T_, M_, self.d_hidden)  # [B*T,M,H]\n        node_attn_out, _ = self.node_attn(h_bt, h_bt, h_bt)\n        h_nodes = node_attn_out.view(B, T_, M_, self.d_hidden)  # [B,T,M,H]\n\n        # 3) Extract agent, target, global per time step\n        idx_agent = agent_idx.view(B,1,1).expand(-1, T_, 1)\n        idx_target = target_idx.view(B,1,1).expand(-1, T_, 1)\n\n        agent_vec = torch.gather(\n            h_nodes,\n            2,\n            idx_agent.unsqueeze(-1).expand(-1,-1,-1,self.d_hidden)\n        ).squeeze(2)\n\n        target_vec = torch.gather(\n            h_nodes,\n            2,\n            idx_target.unsqueeze(-1).expand(-1,-1,-1,self.d_hidden)\n        ).squeeze(2)\n\n        global_vec = h_nodes.mean(dim=2)\n\n        z = torch.cat([agent_vec, target_vec, global_vec], dim=-1)  # [B,T,3H]\n\n        # 4) Temporal Transformer\n        z = z + self.temporal_pos[:, :T_, :]\n        z_enc = self.temporal_encoder(z)  # [B,T,3H]\n\n        # 5) Pool & classify\n        z_pool = z_enc.mean(dim=1)  # [B,3H]\n        logits = self.cls_head(z_pool)  # [B,num_classes]\n        return logits\n\n# ------------------------------------------------------------\n# 3. Small dataset class (per video)\n# ------------------------------------------------------------\n\nclass PoseBehaviorDataset(Dataset):\n    def __init__(self, samples, label_encoder):\n        self.samples = samples\n        self.le = label_encoder\n        self.labels = self.le.transform([s[\"action\"] for s in samples])\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        s = self.samples[idx]\n        pose = torch.from_numpy(s[\"pose\"]).float()    # [T,M,B,2]\n        mask = torch.from_numpy(s[\"mask\"]).float()    # [T,M,B]\n        agent_idx = torch.tensor(s[\"agent_idx\"], dtype=torch.long)\n        target_idx = torch.tensor(s[\"target_idx\"], dtype=torch.long)\n        y = torch.tensor(self.labels[idx], dtype=torch.long)\n        return pose, mask, agent_idx, target_idx, y\n\n# ------------------------------------------------------------\n# 4. Training utilities (per video, incremental)\n# ------------------------------------------------------------\n\ndef build_samples_for_video(tracking_df, ann_df):\n    \"\"\"Build windowed samples for one video. Returns list of dicts.\"\"\"\n    pose, mask, frame_values, mouse_to_idx = build_pose_tensor(tracking_df)\n    if pose is None:\n        return []\n\n    # ---- Normalize pose per-video to [-1,1] ----\n    valid_mask = (mask == 1)\n    if valid_mask.sum() == 0:\n        return []\n\n    x_vals = pose[..., 0][valid_mask]\n    y_vals = pose[..., 1][valid_mask]\n    x_min, x_max = x_vals.min(), x_vals.max()\n    y_min, y_max = y_vals.min(), y_vals.max()\n    x_range = max(x_max - x_min, 1e-6)\n    y_range = max(y_max - y_min, 1e-6)\n    pose[..., 0] = 2 * (pose[..., 0] - x_min) / x_range - 1\n    pose[..., 1] = 2 * (pose[..., 1] - y_min) / y_range - 1\n\n    F = pose.shape[0]\n    min_frame_idx = 0\n    max_frame_idx = F - 1\n\n    frame_to_idx = {fv:i for i,fv in enumerate(frame_values)}\n\n    # frames with any behavior\n    behavior_frame_idx = set()\n    for _, r in ann_df.iterrows():\n        s, e = int(r[\"start_frame\"]), int(r[\"stop_frame\"])\n        for fr in range(s, e+1):\n            if fr in frame_to_idx:\n                behavior_frame_idx.add(frame_to_idx[fr])\n\n    samples_video = []\n    actions_set = set()\n\n    # positive windows\n    for (agent_id, target_id, action), grp in ann_df.groupby([\"agent_id\",\"target_id\",\"action\"]):\n        if agent_id not in mouse_to_idx or target_id not in mouse_to_idx:\n            continue\n        agent_idx = mouse_to_idx[agent_id]\n        target_idx = mouse_to_idx[target_id]\n\n        pos_idx = []\n        for _, r in grp.iterrows():\n            s, e = int(r[\"start_frame\"]), int(r[\"stop_frame\"])\n            for fr in range(s, e+1):\n                if fr in frame_to_idx:\n                    pos_idx.append(frame_to_idx[fr])\n        pos_idx = sorted(set(pos_idx))\n        if len(pos_idx) == 0:\n            continue\n\n        centers = np.random.choice(pos_idx, size=min(10, len(pos_idx)), replace=False)\n\n        for c in centers:\n            start = c - T//2\n            end   = start + T - 1\n            if start < min_frame_idx or end > max_frame_idx:\n                continue\n            window_pose = pose[start:end+1]       # [T,M,B,2]\n            window_mask = mask[start:end+1]       # [T,M,B]\n            samples_video.append({\n                \"agent_idx\": agent_idx,\n                \"target_idx\": target_idx,\n                \"pose\": window_pose,\n                \"mask\": window_mask,\n                \"action\": action\n            })\n            actions_set.add(action)\n\n    if len(samples_video) == 0:\n        return []\n\n    # negative windows\n    all_idx = set(range(F))\n    neg_candidates = list(all_idx - behavior_frame_idx)\n    if len(neg_candidates) > 0:\n        neg_centers = np.random.choice(\n            neg_candidates,\n            size=min(len(samples_video), len(neg_candidates)),\n            replace=False\n        )\n\n        mouse_ids_sorted = sorted(mouse_to_idx.items(), key=lambda x:x[1])\n        if len(mouse_ids_sorted) >= 2:\n            agent_idx = mouse_ids_sorted[0][1]\n            target_idx = mouse_ids_sorted[1][1]\n        else:\n            agent_idx = mouse_ids_sorted[0][1]\n            target_idx = agent_idx\n\n        for c in neg_centers:\n            start = c - T//2\n            end   = start + T - 1\n            if start < min_frame_idx or end > max_frame_idx:\n                continue\n            window_pose = pose[start:end+1]\n            window_mask = mask[start:end+1]\n            samples_video.append({\n                \"agent_idx\": agent_idx,\n                \"target_idx\": target_idx,\n                \"pose\": window_pose,\n                \"mask\": window_mask,\n                \"action\": \"none\"\n            })\n            actions_set.add(\"none\")\n\n    if len(samples_video) > MAX_SAMPLES_PER_VIDEO:\n        samples_video = random.sample(samples_video, MAX_SAMPLES_PER_VIDEO)\n\n    return samples_video, actions_set\n\n# ------------------------------------------------------------\n# 5. Global label encoder (collect actions first pass)\n# ------------------------------------------------------------\n\n# First pass over some videos just to discover all actions\nprint(\"\\n=== Discovering action set across some videos ===\")\nglobal_actions = set()\nvideo_paths = []\n\nfor lab in os.listdir(TRACK_DIR):\n    lab_path = os.path.join(TRACK_DIR, lab)\n    track_files = [f for f in os.listdir(lab_path) if f.endswith(\".parquet\")]\n    for fname in track_files:\n        video_paths.append(os.path.join(lab_path, fname))\n\n# sort for reproducibility\nvideo_paths = sorted(video_paths)\n\ntmp_video_count = 0\nfor tpath in video_paths:\n    tracking_df, vid = load_tracking_file(tpath)\n    ann_df = load_annotation_for_video(vid)\n    if ann_df is None or ann_df.empty:\n        continue\n    global_actions.update(set(ann_df[\"action\"].unique().tolist()))\n    tmp_video_count += 1\n    if tmp_video_count >= 10:  # just a few to initialize\n        break\n\nglobal_actions.add(\"none\")\nglobal_actions = sorted(list(global_actions))\nprint(\"Discovered actions:\", global_actions)\n\n# global label encoder\nle = torch.nn.Module()  # dummy holder to keep API consistent\nfrom sklearn.preprocessing import LabelEncoder as SkLabelEnc\nlabel_encoder = SkLabelEnc()\nlabel_encoder.fit(global_actions)\nprint(\"Label classes:\", list(label_encoder.classes_))\n\n# ------------------------------------------------------------\n# 6. Initialize model & optimizer ONCE\n# ------------------------------------------------------------\n\nnum_classes = len(label_encoder.classes_)\nmodel = FullAGIModel(\n    d_pose=2,\n    d_hidden=64,\n    num_heads=4,\n    num_classes=num_classes,\n    T=T,\n    M=M_MAX,\n    B=B_MAX\n).to(DEVICE)\n\ncriterion = nn.CrossEntropyLoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\n# ------------------------------------------------------------\n# 7. Folder-by-folder, video-by-video training\n# ------------------------------------------------------------\n\nTOTAL_VIDEOS_USED = 0\nall_train_losses, all_val_losses = [], []\nall_train_accs, all_val_accs = [], []\n\nfor tpath in video_paths:\n    if TOTAL_VIDEOS_USED >= MAX_VIDEOS:\n        break\n\n    tracking_df, vid = load_tracking_file(tpath)\n    ann_df = load_annotation_for_video(vid)\n    if ann_df is None or ann_df.empty:\n        continue\n\n    samples_video_tuple = build_samples_for_video(tracking_df, ann_df)\n    if not samples_video_tuple:\n        continue\n    samples_video, actions_this = samples_video_tuple\n\n    # filter out samples whose action not in global label_encoder (shouldn't happen, but safe)\n    samples_video = [s for s in samples_video if s[\"action\"] in label_encoder.classes_]\n    if len(samples_video) == 0:\n        continue\n\n    # train/val split for this video\n    random.shuffle(samples_video)\n    split = int(0.8 * len(samples_video))\n    train_samples = samples_video[:split]\n    val_samples   = samples_video[split:]\n\n    train_ds = PoseBehaviorDataset(train_samples, label_encoder)\n    val_ds   = PoseBehaviorDataset(val_samples,   label_encoder)\n\n    train_loader = DataLoader(train_ds, batch_size=8, shuffle=True,  num_workers=0)\n    val_loader   = DataLoader(val_ds,   batch_size=8, shuffle=False, num_workers=0)\n\n    # one small local epoch for this video\n    model.train()\n    running_loss, running_correct, running_total = 0.0, 0, 0\n    for pose, mask, agent_idx, target_idx, y in train_loader:\n        pose = pose.to(DEVICE)\n        mask = mask.to(DEVICE)\n        agent_idx = agent_idx.to(DEVICE)\n        target_idx = target_idx.to(DEVICE)\n        y = y.to(DEVICE)\n\n        optimizer.zero_grad()\n        logits = model(pose, mask, agent_idx, target_idx)\n        loss = criterion(logits, y)\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)\n        optimizer.step()\n\n        running_loss += loss.item() * pose.size(0)\n        preds = logits.argmax(dim=-1)\n        running_correct += (preds == y).sum().item()\n        running_total += pose.size(0)\n\n    train_loss = running_loss / max(running_total, 1)\n    train_acc  = running_correct / max(running_total, 1)\n\n    # validation on this video\n    model.eval()\n    val_loss, val_correct, val_total = 0.0, 0, 0\n    with torch.no_grad():\n        for pose, mask, agent_idx, target_idx, y in val_loader:\n            pose = pose.to(DEVICE)\n            mask = mask.to(DEVICE)\n            agent_idx = agent_idx.to(DEVICE)\n            target_idx = target_idx.to(DEVICE)\n            y = y.to(DEVICE)\n\n            logits = model(pose, mask, agent_idx, target_idx)\n            loss = criterion(logits, y)\n\n            val_loss += loss.item() * pose.size(0)\n            preds = logits.argmax(dim=-1)\n            val_correct += (preds == y).sum().item()\n            val_total += pose.size(0)\n\n    val_loss = val_loss / max(val_total, 1)\n    val_acc  = val_correct / max(val_total, 1)\n\n    all_train_losses.append(train_loss)\n    all_val_losses.append(val_loss)\n    all_train_accs.append(train_acc)\n    all_val_accs.append(val_acc)\n\n    TOTAL_VIDEOS_USED += 1\n\n    print(f\"\\n[Video {vid}] \"\n          f\"train_loss={train_loss:.4f}, train_acc={train_acc:.3f} | \"\n          f\"val_loss={val_loss:.4f}, val_acc={val_acc:.3f} | \"\n          f\"videos_used={TOTAL_VIDEOS_USED}/{MAX_VIDEOS}\")\n\n    del tracking_df, ann_df, samples_video, train_ds, val_ds\n    gc.collect()\n\nprint(\"\\n=== Training complete ===\")\nif all_train_losses:\n    print(\"Average train loss:\", sum(all_train_losses)/len(all_train_losses))\n    print(\"Average val loss:\",   sum(all_val_losses)/len(all_val_losses))\n    print(\"Average train acc:\",  sum(all_train_accs)/len(all_train_accs))\n    print(\"Average val acc:\",    sum(all_val_accs)/len(all_val_accs))\nelse:\n    print(\"No videos were successfully processed.\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-15T07:43:37.282113Z","iopub.execute_input":"2025-12-15T07:43:37.283312Z","iopub.status.idle":"2025-12-15T08:49:28.288471Z","shell.execute_reply.started":"2025-12-15T07:43:37.283263Z","shell.execute_reply":"2025-12-15T08:49:28.286405Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\nROOT = \"/kaggle/input/MABe-mouse-behavior-detection\"\nTRACK = os.path.join(ROOT, \"train_tracking\")\n\nfor lab in os.listdir(TRACK):\n    lab_path = os.path.join(TRACK, lab)\n    print(\"\\nLAB:\", lab)\n    print(\"Files:\", os.listdir(lab_path)[:10])\n    break   # only show first lab\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-11T09:43:38.409481Z","iopub.execute_input":"2025-12-11T09:43:38.409896Z","iopub.status.idle":"2025-12-11T09:43:38.421855Z","shell.execute_reply.started":"2025-12-11T09:43:38.409851Z","shell.execute_reply":"2025-12-11T09:43:38.420946Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2","metadata":{}},{"cell_type":"code","source":"# ============================================================\n# MABe Mouse Behavior Detection - Full End-to-End Pipeline\n# Training + Inference + F-Beta + Submission\n# ============================================================\n\nimport os\nimport gc\nimport random\nimport json\nfrom collections import defaultdict\n\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nfrom sklearn.preprocessing import LabelEncoder as SkLabelEnc\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\n# ============================================================\n# 0. Config\n# ============================================================\n\nROOT = \"/kaggle/input/MABe-mouse-behavior-detection\"\nTRACK_DIR = os.path.join(ROOT, \"train_tracking\")\nANN_DIR   = os.path.join(ROOT, \"train_annotation\")\nTEST_TRACK = os.path.join(ROOT, \"test_tracking\")  # for inference\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"Using device:\", DEVICE)\n\nT = 32\nM_MAX = 4\nB_MAX = 16\nMAX_VIDEOS = 40\nMAX_SAMPLES_PER_VIDEO = 200\nSEED = 42\n\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\n\n# ============================================================\n# 1. Utilities\n# ============================================================\n\ndef load_tracking_file(path):\n    df = pd.read_parquet(path)\n    vid = int(os.path.basename(path).replace(\".parquet\",\"\"))\n    df[\"video_id\"] = vid\n    return df, vid\n\ndef load_annotation_for_video(video_id):\n    fname = f\"{video_id}.parquet\"\n    for lab in os.listdir(ANN_DIR):\n        p = os.path.join(ANN_DIR, lab, fname)\n        if os.path.exists(p):\n            df = pd.read_parquet(p)\n            df[\"video_id\"] = video_id\n            return df\n    return None\n\ndef build_pose_tensor(tracking_df):\n    frames = sorted(tracking_df[\"video_frame\"].unique().tolist())\n    frame_to_idx = {f:i for i,f in enumerate(frames)}\n    F = len(frames)\n\n    mice = sorted(tracking_df[\"mouse_id\"].unique().tolist())\n    if len(mice) == 0 or len(mice) > M_MAX:\n        return None, None, None, None\n    mice = mice[:M_MAX]\n    mouse_to_idx = {m:i for i,m in enumerate(mice)}\n\n    bodyparts = sorted(tracking_df[\"bodypart\"].unique().tolist())\n    bodyparts = bodyparts[:B_MAX]\n    bp_to_idx = {b:i for i,b in enumerate(bodyparts)}\n    B = len(bodyparts)\n\n    pose = np.zeros((F, M_MAX, B_MAX, 2), dtype=np.float32)\n    mask = np.zeros((F, M_MAX, B_MAX), dtype=np.float32)\n\n    for _, row in tracking_df.iterrows():\n        f = frame_to_idx[row[\"video_frame\"]]\n        m_id = row[\"mouse_id\"]\n        b_name = row[\"bodypart\"]\n        if m_id not in mouse_to_idx or b_name not in bp_to_idx:\n            continue\n        mi = mouse_to_idx[m_id]\n        bi = bp_to_idx[b_name]\n        pose[f, mi, bi, 0] = row[\"x\"]\n        pose[f, mi, bi, 1] = row[\"y\"]\n        mask[f, mi, bi] = 1.0\n\n    return pose, mask, frames, mouse_to_idx\n\n# ============================================================\n# 2. Dataset Class\n# ============================================================\n\nclass PoseBehaviorDataset(Dataset):\n    def __init__(self, samples, label_encoder):\n        self.samples = samples\n        self.le = label_encoder\n        self.labels = self.le.transform([s[\"action\"] for s in samples])\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        s = self.samples[idx]\n        pose = torch.from_numpy(s[\"pose\"]).float()\n        mask = torch.from_numpy(s[\"mask\"]).float()\n        agent_idx = torch.tensor(s[\"agent_idx\"], dtype=torch.long)\n        target_idx = torch.tensor(s[\"target_idx\"], dtype=torch.long)\n        y = torch.tensor(self.labels[idx], dtype=torch.long)\n        return pose, mask, agent_idx, target_idx, y\n\n# ============================================================\n# 3. Model\n# ============================================================\n\nclass FullAGIModel(nn.Module):\n    def __init__(self, d_pose=2, d_hidden=64, num_heads=4, num_classes=10, T=T, M=M_MAX, B=B_MAX):\n        super().__init__()\n        self.T = T\n        self.M = M\n        self.B = B\n        self.d_hidden = d_hidden\n        self.bodypart_proj = nn.Linear(d_pose, d_hidden)\n        self.node_attn = nn.MultiheadAttention(embed_dim=d_hidden, num_heads=num_heads, batch_first=True)\n        self.temporal_pos = nn.Parameter(torch.randn(1, T, d_hidden*3))\n        encoder_layer = nn.TransformerEncoderLayer(d_model=d_hidden*3, nhead=num_heads, dim_feedforward=d_hidden*4, batch_first=True)\n        self.temporal_encoder = nn.TransformerEncoder(encoder_layer, num_layers=2)\n        self.cls_head = nn.Sequential(\n            nn.Linear(d_hidden*3, d_hidden),\n            nn.ReLU(),\n            nn.Linear(d_hidden, num_classes)\n        )\n\n    def forward(self, pose, mask, agent_idx=None, target_idx=None):\n        # pose: [B,T,M,B,2], mask: [B,T,M,B]\n        B, T_, M_, B_, Dp = pose.shape\n        assert T_ == self.T and M_ == self.M\n\n        x = pose.view(B*T_*M_, B_, Dp)\n        m = mask.view(B*T_*M_, B_)\n        bp_emb = self.bodypart_proj(x)\n        m_sum = m.sum(dim=1, keepdim=True) + 1e-6\n        h_mouse = (bp_emb * m.unsqueeze(-1)).sum(dim=1)/m_sum\n        h_mouse = h_mouse.view(B, T_, M_, self.d_hidden)\n        h_bt = h_mouse.view(B*T_, M_, self.d_hidden)\n        node_attn_out, _ = self.node_attn(h_bt, h_bt, h_bt)\n        h_nodes = node_attn_out.view(B, T_, M_, self.d_hidden)\n\n        # agent/target selection\n        if agent_idx is None:\n            agent_idx = torch.zeros(B, dtype=torch.long, device=pose.device)\n        if target_idx is None:\n            target_idx = torch.ones(B, dtype=torch.long, device=pose.device)\n\n        idx_agent = agent_idx.view(B,1,1).expand(-1,T_,1)\n        idx_target = target_idx.view(B,1,1).expand(-1,T_,1)\n        agent_vec = torch.gather(h_nodes, 2, idx_agent.unsqueeze(-1).expand(-1,-1,-1,self.d_hidden)).squeeze(2)\n        target_vec = torch.gather(h_nodes, 2, idx_target.unsqueeze(-1).expand(-1,-1,-1,self.d_hidden)).squeeze(2)\n        global_vec = h_nodes.mean(dim=2)\n        z = torch.cat([agent_vec, target_vec, global_vec], dim=-1)\n        z = z + self.temporal_pos[:, :T_, :]\n        z_enc = self.temporal_encoder(z)\n        z_pool = z_enc.mean(dim=1)\n        logits = self.cls_head(z_pool)\n        return logits\n\n# ============================================================\n# 4. MABe F-Beta Metric\n# ============================================================\n\ndef mouse_fbeta(solution: pd.DataFrame, submission: pd.DataFrame, beta: float = 1) -> float:\n    if len(solution)==0 or len(submission)==0:\n        return 0.0\n    solution = pl.DataFrame(solution)\n    submission = pl.DataFrame(submission)\n    solution = solution.with_columns(\n        pl.concat_str([\n            pl.col('video_id').cast(pl.Utf8),\n            pl.col('agent_id').cast(pl.Utf8),\n            pl.col('target_id').cast(pl.Utf8),\n            pl.col('action')], separator='_').alias('label_key')\n    )\n    submission = submission.with_columns(\n        pl.concat_str([\n            pl.col('video_id').cast(pl.Utf8),\n            pl.col('agent_id').cast(pl.Utf8),\n            pl.col('target_id').cast(pl.Utf8),\n            pl.col('action')], separator='_').alias('prediction_key')\n    )\n\n    label_frames = defaultdict(set)\n    prediction_frames = defaultdict(set)\n    for row in solution.to_dicts():\n        label_frames[row['label_key']].update(range(row['start_frame'], row['stop_frame']))\n    for row in submission.to_dicts():\n        prediction_frames[row['prediction_key']].update(range(row['start_frame'], row['stop_frame']))\n\n    tps, fps, fns = defaultdict(int), defaultdict(int), defaultdict(int)\n    for key, pred in prediction_frames.items():\n        action = key.split('_')[-1]\n        matched = label_frames.get(key, set())\n        tps[action] += len(pred & matched)\n        fps[action] += len(pred - matched)\n        fns[action] += len(matched - pred)\n    all_actions = set([k.split('_')[-1] for k in label_frames.keys()])\n    f1s = []\n    for a in all_actions:\n        denom = (1+beta**2)*tps[a] + fps[a] + beta**2*fns[a]\n        f1s.append((1+beta**2)*tps[a]/denom if denom>0 else 0)\n    return sum(f1s)/len(f1s) if f1s else 0.0\n\n# ============================================================\n# 5. Build Samples\n# ============================================================\n\ndef build_samples_for_video(tracking_df, ann_df):\n    pose, mask, frames, mouse_to_idx = build_pose_tensor(tracking_df)\n    if pose is None:\n        return []\n    valid_mask = mask==1\n    if valid_mask.sum()==0:\n        return []\n    x_vals = pose[...,0][valid_mask]\n    y_vals = pose[...,1][valid_mask]\n    x_min,x_max = x_vals.min(), x_vals.max()\n    y_min,y_max = y_vals.min(), y_vals.max()\n    pose[...,0] = 2*(pose[...,0]-x_min)/(x_max-x_min+1e-6)-1\n    pose[...,1] = 2*(pose[...,1]-y_min)/(y_max-y_min+1e-6)-1\n\n    F = pose.shape[0]\n    frame_to_idx = {fv:i for i,fv in enumerate(frames)}\n    behavior_frame_idx = set()\n    for _, r in ann_df.iterrows():\n        for fr in range(int(r[\"start_frame\"]), int(r[\"stop_frame\"])+1):\n            if fr in frame_to_idx:\n                behavior_frame_idx.add(frame_to_idx[fr])\n\n    samples_video = []\n    actions_set = set()\n    for (agent_id, target_id, action), grp in ann_df.groupby([\"agent_id\",\"target_id\",\"action\"]):\n        if agent_id not in mouse_to_idx or target_id not in mouse_to_idx:\n            continue\n        agent_idx = mouse_to_idx[agent_id]\n        target_idx = mouse_to_idx[target_id]\n        pos_idx = []\n        for _, r in grp.iterrows():\n            for fr in range(int(r[\"start_frame\"]), int(r[\"stop_frame\"])+1):\n                if fr in frame_to_idx:\n                    pos_idx.append(frame_to_idx[fr])\n        pos_idx = sorted(set(pos_idx))\n        if len(pos_idx)==0:\n            continue\n        centers = np.random.choice(pos_idx, size=min(10,len(pos_idx)), replace=False)\n        for c in centers:\n            start = c-T//2\n            end = start+T-1\n            if start<0 or end>=F:\n                continue\n            samples_video.append({\n                \"agent_idx\": agent_idx,\n                \"target_idx\": target_idx,\n                \"pose\": pose[start:end+1],\n                \"mask\": mask[start:end+1],\n                \"action\": action\n            })\n            actions_set.add(action)\n\n    # negative windows\n    all_idx = set(range(F))\n    neg_candidates = list(all_idx - behavior_frame_idx)\n    if len(neg_candidates)>0:\n        neg_centers = np.random.choice(neg_candidates, size=min(len(samples_video), len(neg_candidates)), replace=False)\n        mouse_ids_sorted = sorted(mouse_to_idx.items(), key=lambda x:x[1])\n        agent_idx = mouse_ids_sorted[0][1]\n        target_idx = mouse_ids_sorted[1][1] if len(mouse_ids_sorted)>1 else agent_idx\n        for c in neg_centers:\n            start = c-T//2\n            end = start+T-1\n            if start<0 or end>=F:\n                continue\n            samples_video.append({\n                \"agent_idx\": agent_idx,\n                \"target_idx\": target_idx,\n                \"pose\": pose[start:end+1],\n                \"mask\": mask[start:end+1],\n                \"action\": \"none\"\n            })\n            actions_set.add(\"none\")\n    if len(samples_video)>MAX_SAMPLES_PER_VIDEO:\n        samples_video = random.sample(samples_video, MAX_SAMPLES_PER_VIDEO)\n    return samples_video, actions_set\n\n# ============================================================\n# 6. Training Loop\n# ============================================================\n\n# Discover actions\nglobal_actions = set()\nvideo_paths = []\nfor lab in os.listdir(TRACK_DIR):\n    lab_path = os.path.join(TRACK_DIR, lab)\n    for f in os.listdir(lab_path):\n        if f.endswith(\".parquet\"):\n            video_paths.append(os.path.join(lab_path,f))\nvideo_paths = sorted(video_paths)\n\ntmp_video_count=0\nfor tpath in video_paths:\n    tracking_df, vid = load_tracking_file(tpath)\n    ann_df = load_annotation_for_video(vid)\n    if ann_df is None or ann_df.empty:\n        continue\n    global_actions.update(set(ann_df[\"action\"].unique().tolist()))\n    tmp_video_count+=1\n    if tmp_video_count>=10:\n        break\nglobal_actions.add(\"none\")\nglobal_actions = sorted(list(global_actions))\nlabel_encoder = SkLabelEnc()\nlabel_encoder.fit(global_actions)\n\nnum_classes = len(label_encoder.classes_)\nmodel = FullAGIModel(num_classes=num_classes).to(DEVICE)\ncriterion = nn.CrossEntropyLoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\nall_train_losses, all_val_losses = [], []\n\nTOTAL_VIDEOS_USED=0\nfor tpath in video_paths:\n    if TOTAL_VIDEOS_USED>=MAX_VIDEOS:\n        break\n    tracking_df, vid = load_tracking_file(tpath)\n    ann_df = load_annotation_for_video(vid)\n    if ann_df is None or ann_df.empty:\n        continue\n    samples_video_tuple = build_samples_for_video(tracking_df, ann_df)\n    if not samples_video_tuple:\n        continue\n    samples_video, _ = samples_video_tuple\n    samples_video = [s for s in samples_video if s[\"action\"] in label_encoder.classes_]\n    if len(samples_video)==0:\n        continue\n    random.shuffle(samples_video)\n    split = int(0.8*len(samples_video))\n    train_samples = samples_video[:split]\n    val_samples = samples_video[split:]\n    train_ds = PoseBehaviorDataset(train_samples, label_encoder)\n    val_ds = PoseBehaviorDataset(val_samples, label_encoder)\n    train_loader = DataLoader(train_ds, batch_size=8, shuffle=True)\n    val_loader = DataLoader(val_ds, batch_size=8, shuffle=False)\n\n    # Training\n    model.train()\n    running_loss=0\n    for pose, mask, agent_idx, target_idx, y in train_loader:\n        pose, mask, agent_idx, target_idx, y = pose.to(DEVICE), mask.to(DEVICE), agent_idx.to(DEVICE), target_idx.to(DEVICE), y.to(DEVICE)\n        optimizer.zero_grad()\n        logits = model(pose, mask, agent_idx, target_idx)\n        loss = criterion(logits, y)\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)\n        optimizer.step()\n        running_loss += loss.item()*pose.size(0)\n    all_train_losses.append(running_loss/len(train_ds))\n\n    # Validation\n    model.eval()\n    val_loss=0\n    val_preds, val_labels=[],[]\n    with torch.no_grad():\n        for pose, mask, agent_idx, target_idx, y in val_loader:\n            pose, mask, agent_idx, target_idx, y = pose.to(DEVICE), mask.to(DEVICE), agent_idx.to(DEVICE), target_idx.to(DEVICE), y.to(DEVICE)\n            logits = model(pose, mask, agent_idx, target_idx)\n            val_loss += criterion(logits, y).item()*pose.size(0)\n            val_preds.extend(logits.argmax(dim=-1).cpu().numpy())\n            val_labels.extend(y.cpu().numpy())\n    all_val_losses.append(val_loss/len(val_ds))\n    TOTAL_VIDEOS_USED+=1\n    print(f\"Video {vid}: train_loss={all_train_losses[-1]:.4f}, val_loss={all_val_losses[-1]:.4f}, videos_used={TOTAL_VIDEOS_USED}/{MAX_VIDEOS}\")\n    del tracking_df, ann_df, samples_video, train_ds, val_ds\n    gc.collect()\n\nprint(\"Training complete.\")\n\n# ============================================================\n# 7. Inference + Submission\n# ============================================================\n\nsubmission_rows=[]\nfor lab in os.listdir(TEST_TRACK):\n    lab_dir = os.path.join(TEST_TRACK, lab)\n    for f in os.listdir(lab_dir):\n        if not f.endswith(\".parquet\"):\n            continue\n        video_id = int(f.replace(\".parquet\",\"\"))\n        df = pd.read_parquet(os.path.join(lab_dir,f))\n        pose, mask, frames, mouse_to_idx = build_pose_tensor(df)\n        if pose is None:\n            continue\n        F = pose.shape[0]\n        STRIDE = 4\n        probs_per_frame = np.zeros((F, num_classes))\n        model.eval()\n        with torch.no_grad():\n            for s in range(0, F-T+1, STRIDE):\n                clip = torch.tensor(pose[s:s+T]).unsqueeze(0).to(DEVICE)\n                mask_clip = torch.tensor(mask[s:s+T]).unsqueeze(0).to(DEVICE)\n                logits = model(clip, mask_clip)\n                probs = torch.softmax(logits, -1)[0].cpu().numpy()\n                for i in range(T):\n                    if s+i<F:\n                        probs_per_frame[s+i] += probs\n        # assign actions based on max probability\n        for m1 in range(M_MAX):\n            for m2 in range(M_MAX):\n                if m1==m2:\n                    continue\n                for action_idx, action in enumerate(label_encoder.classes_):\n                    if action==\"none\":\n                        continue\n                    frames_above= np.where(probs_per_frame[:, action_idx]>0.5)[0]\n                    if len(frames_above)==0:\n                        continue\n                    submission_rows.append({\n                        \"video_id\": video_id,\n                        \"agent_id\": m1,\n                        \"target_id\": m2,\n                        \"start_frame\": int(frames_above[0]),\n                        \"stop_frame\": int(frames_above[-1]),\n                        \"action\": action\n                    })\n\nsubmission_df = pd.DataFrame(submission_rows)\nsubmission_df.to_csv(\"submission.csv\", index=False)\nprint(\"Saved submission.csv\")\n\n# ============================================================\n# 8. Compute MABe F-Beta\n# ============================================================\n\nsolution_rows=[]\nfor lab in os.listdir(ANN_DIR):\n    lab_dir = os.path.join(ANN_DIR, lab)\n    for f in os.listdir(lab_dir):\n        if not f.endswith(\".parquet\"):\n            continue\n        video_id = int(f.replace(\".parquet\",\"\"))\n        df = pd.read_parquet(os.path.join(lab_dir,f))\n        for _, r in df.iterrows():\n            solution_rows.append({\n                \"video_id\": video_id,\n                \"agent_id\": int(r[\"agent_id\"]),\n                \"target_id\": int(r[\"target_id\"]),\n                \"start_frame\": int(r[\"start_frame\"]),\n                \"stop_frame\": int(r[\"stop_frame\"]),\n                \"action\": r[\"action\"]\n            })\nsolution_df = pd.DataFrame(solution_rows)\n\nfbeta_score = mouse_fbeta(solution_df, submission_df, beta=1)\nprint(f\"MABe F-Beta Score: {fbeta_score:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T12:47:31.773609Z","iopub.execute_input":"2025-12-15T12:47:31.774278Z","iopub.status.idle":"2025-12-15T13:37:46.510593Z","shell.execute_reply.started":"2025-12-15T12:47:31.774243Z","shell.execute_reply":"2025-12-15T13:37:46.50991Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}