{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ==================== COMPLETE MABe SOLUTION WITH DUAL-SCALE CNN ====================\n\nvalidate_or_submit = 'submit'\nverbose = True\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport itertools\nimport warnings\nimport json\nimport os\nimport gc\nfrom collections import defaultdict\nimport polars as pl\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, WeightedRandomSampler\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\n\nwarnings.filterwarnings('ignore')\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# ==================== SCORING FUNCTIONS ====================\n\nclass HostVisibleError(Exception):\n    pass\n\ndef single_lab_f1(lab_solution: pl.DataFrame, lab_submission: pl.DataFrame, beta: float = 1) -> float:\n    label_frames: defaultdict[str, set[int]] = defaultdict(set)\n    prediction_frames: defaultdict[str, set[int]] = defaultdict(set)\n\n    for row in lab_solution.to_dicts():\n        label_frames[row['label_key']].update(range(row['start_frame'], row['stop_frame']))\n\n    for video in lab_solution['video_id'].unique():\n        active_labels: str = lab_solution.filter(pl.col('video_id') == video)['behaviors_labeled'].first()\n        active_labels: set[str] = set(json.loads(active_labels))\n        predicted_mouse_pairs: defaultdict[str, set[int]] = defaultdict(set)\n\n        for row in lab_submission.filter(pl.col('video_id') == video).to_dicts():\n            if ','.join([str(row['agent_id']), str(row['target_id']), row['action']]) not in active_labels:\n                continue\n           \n            new_frames = set(range(row['start_frame'], row['stop_frame']))\n            new_frames = new_frames.difference(prediction_frames[row['prediction_key']])\n            prediction_pair = ','.join([str(row['agent_id']), str(row['target_id'])])\n            if predicted_mouse_pairs[prediction_pair].intersection(new_frames):\n                raise HostVisibleError('Multiple predictions for the same frame from one agent/target pair')\n            prediction_frames[row['prediction_key']].update(new_frames)\n            predicted_mouse_pairs[prediction_pair].update(new_frames)\n\n    tps = defaultdict(int)\n    fns = defaultdict(int)\n    fps = defaultdict(int)\n    for key, pred_frames in prediction_frames.items():\n        action = key.split('_')[-1]\n        matched_label_frames = label_frames[key]\n        tps[action] += len(pred_frames.intersection(matched_label_frames))\n        fns[action] += len(matched_label_frames.difference(pred_frames))\n        fps[action] += len(pred_frames.difference(matched_label_frames))\n\n    distinct_actions = set()\n    for key, frames in label_frames.items():\n        action = key.split('_')[-1]\n        distinct_actions.add(action)\n        if key not in prediction_frames:\n            fns[action] += len(frames)\n\n    action_f1s = []\n    for action in distinct_actions:\n        if tps[action] + fns[action] + fps[action] == 0:\n            action_f1s.append(0)\n        else:\n            action_f1s.append((1 + beta**2) * tps[action] / ((1 + beta**2) * tps[action] + beta**2 * fns[action] + fps[action]))\n    return sum(action_f1s) / len(action_f1s)\n\n# ==================== DATA LOADING ====================\n\ntrain = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/train.csv')\ntest = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\n\ntrain['n_mice'] = 4 - train[['mouse1_strain', 'mouse2_strain', 'mouse3_strain', 'mouse4_strain']].isna().sum(axis=1)\nbody_parts_tracked_list = list(np.unique(train.body_parts_tracked))\n\ndrop_body_parts = ['headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', \n                   'headpiece_bottomfrontright', 'headpiece_topbackleft', 'headpiece_topbackright', \n                   'headpiece_topfrontleft', 'headpiece_topfrontright', 'spine_1', 'spine_2', \n                   'tail_middle_1', 'tail_middle_2', 'tail_midpoint']\n\ndef generate_mouse_data(dataset, traintest, traintest_directory=None, generate_single=True, generate_pair=True):\n    assert traintest in ['train', 'test']\n    if traintest_directory is None:\n        traintest_directory = f\"/kaggle/input/MABe-mouse-behavior-detection/{traintest}_tracking\"\n    \n    for _, row in dataset.iterrows():\n        lab_id = row.lab_id\n        if lab_id.startswith('MABe22'): \n            continue\n        video_id = row.video_id\n\n        if type(row.behaviors_labeled) != str:\n            if verbose: \n                print('No labeled behaviors:', lab_id, video_id)\n            continue\n\n        path = f\"{traintest_directory}/{lab_id}/{video_id}.parquet\"\n        vid = pd.read_parquet(path)\n        if len(np.unique(vid.bodypart)) > 5:\n            vid = vid.query(\"~ bodypart.isin(@drop_body_parts)\")\n        \n        pvid = vid.pivot(columns=['mouse_id', 'bodypart'], index='video_frame', values=['x', 'y'])\n        pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T\n        pvid /= row.pix_per_cm_approx\n\n        vid_behaviors = json.loads(row.behaviors_labeled)\n        vid_behaviors = sorted(list({b.replace(\"'\", \"\") for b in vid_behaviors}))\n        vid_behaviors = [b.split(',') for b in vid_behaviors]\n        vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action'])\n        \n        if traintest == 'train':\n            try:\n                annot = pd.read_parquet(path.replace('train_tracking', 'train_annotation'))\n            except FileNotFoundError:\n                continue\n\n        if generate_single:\n            vid_behaviors_subset = vid_behaviors.query(\"target == 'self'\")\n            for mouse_id_str in np.unique(vid_behaviors_subset.agent):\n                try:\n                    mouse_id = int(mouse_id_str[-1])\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"agent == @mouse_id_str\").action)\n                    single_mouse = pvid.loc[:, mouse_id]\n                    single_mouse_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': mouse_id_str,\n                        'target_id': 'self',\n                        'video_frame': single_mouse.index\n                    })\n                    if traintest == 'train':\n                        single_mouse_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=single_mouse.index)\n                        annot_subset = annot.query(\"(agent_id == @mouse_id) & (target_id == @mouse_id)\")\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\n                            single_mouse_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'single', single_mouse, single_mouse_meta, single_mouse_label\n                    else:\n                        yield 'single', single_mouse, single_mouse_meta, vid_agent_actions\n                except KeyError:\n                    pass\n\n        if generate_pair:\n            vid_behaviors_subset = vid_behaviors.query(\"target != 'self'\")\n            if len(vid_behaviors_subset) > 0:\n                for agent, target in itertools.permutations(np.unique(pvid.columns.get_level_values('mouse_id')), 2):\n                    agent_str = f\"mouse{agent}\"\n                    target_str = f\"mouse{target}\"\n                    vid_agent_actions = np.unique(vid_behaviors_subset.query(\"(agent == @agent_str) & (target == @target_str)\").action)\n                    if len(vid_agent_actions) == 0:\n                        continue\n                    mouse_pair = pd.concat([pvid[agent], pvid[target]], axis=1, keys=['A', 'B'])\n                    mouse_pair_meta = pd.DataFrame({\n                        'video_id': video_id,\n                        'agent_id': agent_str,\n                        'target_id': target_str,\n                        'video_frame': mouse_pair.index\n                    })\n                    if traintest == 'train':\n                        mouse_pair_label = pd.DataFrame(0.0, columns=vid_agent_actions, index=mouse_pair.index)\n                        annot_subset = annot.query(\"(agent_id == @agent) & (target_id == @target)\")\n                        for i in range(len(annot_subset)):\n                            annot_row = annot_subset.iloc[i]\n                            mouse_pair_label.loc[annot_row['start_frame']:annot_row['stop_frame'], annot_row.action] = 1.0\n                        yield 'pair', mouse_pair, mouse_pair_meta, mouse_pair_label\n                    else:\n                        yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions\n\n# ==================== FEATURE ENGINEERING ====================\n\ndef transform_single(single_mouse, body_parts_tracked):\n    available_body_parts = single_mouse.columns.get_level_values(0)\n    \n    X = pd.DataFrame({\n        f\"{p1}+{p2}\": np.square(single_mouse[p1] - single_mouse[p2]).sum(axis=1, skipna=False)\n        for p1, p2 in itertools.combinations(body_parts_tracked, 2) \n        if p1 in available_body_parts and p2 in available_body_parts\n    })\n    X = X.reindex(columns=[f\"{p1}+{p2}\" for p1, p2 in itertools.combinations(body_parts_tracked, 2)], copy=False)\n\n    if all(p in single_mouse.columns for p in ['ear_left', 'ear_right', 'tail_base']):\n        shifted = single_mouse[['ear_left', 'ear_right', 'tail_base']].shift(10)\n        X['sp_lf'] = np.square(single_mouse['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False)\n        X['sp_rt'] = np.square(single_mouse['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False)\n        X['sp_tb'] = np.square(single_mouse['tail_base'] - shifted['tail_base']).sum(axis=1, skipna=False)\n    \n    if 'body_center' in available_body_parts:\n        cx = single_mouse['body_center']['x']\n        cy = single_mouse['body_center']['y']\n        \n        for w in [5, 15, 30, 60]:\n            X[f'cx_m{w}'] = cx.rolling(w, min_periods=1, center=True).mean()\n            X[f'cy_m{w}'] = cy.rolling(w, min_periods=1, center=True).mean()\n            X[f'cx_s{w}'] = cx.rolling(w, min_periods=1, center=True).std()\n            X[f'cy_s{w}'] = cy.rolling(w, min_periods=1, center=True).std()\n    \n    return X\n\ndef transform_pair(mouse_pair, body_parts_tracked):\n    avail_A = mouse_pair['A'].columns.get_level_values(0)\n    avail_B = mouse_pair['B'].columns.get_level_values(0)\n    \n    X = pd.DataFrame({\n        f\"12+{p1}+{p2}\": np.square(mouse_pair['A'][p1] - mouse_pair['B'][p2]).sum(axis=1, skipna=False)\n        for p1, p2 in itertools.product(body_parts_tracked, repeat=2) \n        if p1 in avail_A and p2 in avail_B\n    })\n    X = X.reindex(columns=[f\"12+{p1}+{p2}\" for p1, p2 in itertools.product(body_parts_tracked, repeat=2)], copy=False)\n\n    if ('A', 'body_center') in mouse_pair.columns and ('B', 'body_center') in mouse_pair.columns:\n        cd = np.square(mouse_pair['A']['body_center'] - mouse_pair['B']['body_center']).sum(axis=1, skipna=False)\n        for w in [5, 15, 30, 60]:\n            X[f'd_m{w}'] = cd.rolling(w, min_periods=1, center=True).mean()\n            X[f'd_s{w}'] = cd.rolling(w, min_periods=1, center=True).std()\n    \n    return X\n\n# ==================== FOCAL LOSS FOR CLASS IMBALANCE ====================\n\nclass FocalLoss(nn.Module):\n    def __init__(self, alpha=0.25, gamma=2.0):\n        super().__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        \n    def forward(self, inputs, targets):\n        bce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')\n        pt = torch.exp(-bce_loss)\n        focal_loss = self.alpha * (1 - pt) ** self.gamma * bce_loss\n        return focal_loss.mean()\n\n# ==================== DUAL-SCALE 1D CNN MODEL ====================\n\nclass DualScaleConv1D(nn.Module):\n    def __init__(self, input_dim, num_classes, local_window=60, \n                 long_window=200, sparse_rate=5, dropout=0.3):\n        super().__init__()\n        self.local_window = local_window\n        self.long_window = long_window\n        self.sparse_rate = sparse_rate\n        self.n_sparse = long_window // sparse_rate\n        \n        # Local branch (fine-grained)\n        self.local_conv1 = nn.Conv1d(input_dim, 128, 5, padding=2)\n        self.local_bn1 = nn.BatchNorm1d(128)\n        self.local_conv2 = nn.Conv1d(128, 256, 5, padding=2)\n        self.local_bn2 = nn.BatchNorm1d(256)\n        self.local_conv3 = nn.Conv1d(256, 256, 3, padding=1)\n        self.local_bn3 = nn.BatchNorm1d(256)\n        \n        # Long branch (sparse)\n        self.long_conv1 = nn.Conv1d(input_dim, 64, 3, padding=1)\n        self.long_bn1 = nn.BatchNorm1d(64)\n        self.long_conv2 = nn.Conv1d(64, 128, 3, padding=1)\n        self.long_bn2 = nn.BatchNorm1d(128)\n        self.long_conv3 = nn.Conv1d(128, 128, 3, padding=1)\n        self.long_bn3 = nn.BatchNorm1d(128)\n        \n        self.pool = nn.MaxPool1d(2)\n        self.dropout = nn.Dropout(dropout)\n        \n        local_sz = local_window // 4\n        long_sz = self.n_sparse // 4\n        fusion_dim = 256 * local_sz + 128 * long_sz\n        \n        self.fc1 = nn.Linear(fusion_dim, 512)\n        self.fc2 = nn.Linear(512, 256)\n        self.fc3 = nn.Linear(256, num_classes)\n    \n    def forward(self, x_local, x_long):\n        # Local branch\n        x_l = x_local.transpose(1, 2)\n        x_l = self.dropout(self.pool(F.relu(self.local_bn1(self.local_conv1(x_l)))))\n        x_l = self.dropout(self.pool(F.relu(self.local_bn2(self.local_conv2(x_l)))))\n        x_l = self.dropout(F.relu(self.local_bn3(self.local_conv3(x_l))))\n        x_l = x_l.flatten(1)\n        \n        # Long branch\n        x_g = x_long.transpose(1, 2)\n        x_g = self.dropout(self.pool(F.relu(self.long_bn1(self.long_conv1(x_g)))))\n        x_g = self.dropout(self.pool(F.relu(self.long_bn2(self.long_conv2(x_g)))))\n        x_g = self.dropout(F.relu(self.long_bn3(self.long_conv3(x_g))))\n        x_g = x_g.flatten(1)\n        \n        # Fusion\n        x = torch.cat([x_l, x_g], 1)\n        x = self.dropout(F.relu(self.fc1(x)))\n        x = self.dropout(F.relu(self.fc2(x)))\n        return self.fc3(x)\n\n# ==================== DUAL-SCALE DATASET ====================\n\nclass DualScaleDataset(Dataset):\n    def __init__(self, features, labels, local_window=60, \n                 long_window=200, sparse_rate=5):\n        self.features = features\n        self.labels = labels\n        self.local_window = local_window\n        self.long_window = long_window\n        self.sparse_rate = sparse_rate\n        self.half_local = local_window // 2\n        self.half_long = long_window // 2\n        self.valid_indices = list(range(self.half_long, len(features) - self.half_long))\n    \n    def __len__(self):\n        return len(self.valid_indices)\n    \n    def __getitem__(self, idx):\n        center = self.valid_indices[idx]\n        local = self.features[center - self.half_local:center + self.half_local]\n        sparse_idx = list(range(center - self.half_long, center + self.half_long, self.sparse_rate))\n        long_range = self.features[sparse_idx]\n        return (torch.FloatTensor(local), \n                torch.FloatTensor(long_range), \n                torch.FloatTensor(self.labels[center]))\n\n# ==================== CLASS-BALANCED SAMPLING ====================\n\ndef create_balanced_sampler(labels, oversample_ratio=2.0):\n    n_pos = labels.sum()\n    n_neg = len(labels) - n_pos\n    if n_pos == 0:\n        return None\n    weights = np.zeros(len(labels))\n    weights[labels == 1] = (n_neg / n_pos) * oversample_ratio\n    weights[labels == 0] = 1.0\n    return WeightedRandomSampler(weights, len(weights), replacement=True)\n\n# ==================== TRAINING FUNCTION ====================\n\ndef train_dual_model(model, train_loader, val_loader, epochs=15, lr=1e-3, patience=3):\n    optimizer = AdamW(model.parameters(), lr=lr, weight_decay=0.01)\n    scheduler = CosineAnnealingLR(optimizer, T_max=epochs)\n    criterion = FocalLoss(alpha=0.25, gamma=2.0)\n    \n    best_loss = float('inf')\n    counter = 0\n    \n    for epoch in range(epochs):\n        model.train()\n        train_loss = 0\n        for local, long_range, labels in train_loader:\n            local, long_range, labels = local.to(device), long_range.to(device), labels.to(device)\n            optimizer.zero_grad()\n            loss = criterion(model(local, long_range), labels)\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            optimizer.step()\n            train_loss += loss.item()\n        \n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for local, long_range, labels in val_loader:\n                local, long_range, labels = local.to(device), long_range.to(device), labels.to(device)\n                val_loss += criterion(model(local, long_range), labels).item()\n        \n        train_loss /= len(train_loader)\n        val_loss /= len(val_loader)\n        \n        if verbose:\n            print(f'  Epoch {epoch+1}/{epochs} - Train: {train_loss:.4f}, Val: {val_loss:.4f}')\n        \n        scheduler.step()\n        \n        if val_loss < best_loss:\n            best_loss = val_loss\n            counter = 0\n            best_state = model.state_dict().copy()\n        else:\n            counter += 1\n            if counter >= patience:\n                if verbose:\n                    print(f'  Early stop at epoch {epoch+1}')\n                break\n    \n    model.load_state_dict(best_state)\n    return model\n\n# ==================== PREDICTION FUNCTION ====================\n\ndef predict_dual(model, features, local_w=60, long_w=200, sparse_r=5):\n    model.eval()\n    half_local = local_w // 2\n    half_long = long_w // 2\n    probs = []\n    \n    with torch.no_grad():\n        for i in range(half_long, len(features) - half_long):\n            local = features[i - half_local:i + half_local]\n            sparse_idx = list(range(i - half_long, i + half_long, sparse_r))\n            long_range = features[sparse_idx]\n            \n            local_t = torch.FloatTensor(local).unsqueeze(0).to(device)\n            long_t = torch.FloatTensor(long_range).unsqueeze(0).to(device)\n            \n            prob = torch.sigmoid(model(local_t, long_t)).cpu().numpy()[0]\n            probs.append(prob)\n    \n    probs = np.array(probs)\n    full = np.zeros((len(features), probs.shape[1]))\n    full[half_long:-half_long] = probs\n    full[:half_long] = probs[0]\n    full[-half_long:] = probs[-1]\n    return full\n\n# ==================== OPTIMAL THRESHOLD CALCULATION ====================\n\ndef find_optimal_threshold(probs, labels):\n    if labels.sum() < 10:\n        return 0.27\n    best_th, best_f1 = 0.27, 0\n    for th in np.arange(0.1, 0.6, 0.05):\n        pred = (probs > th).astype(int)\n        tp = ((pred == 1) & (labels == 1)).sum()\n        fp = ((pred == 1) & (labels == 0)).sum()\n        fn = ((pred == 0) & (labels == 1)).sum()\n        if tp + fp + fn > 0:\n            f1 = 2 * tp / (2 * tp + fp + fn)\n            if f1 > best_f1:\n                best_f1, best_th = f1, th\n    return best_th\n\n# ==================== PREDICTION TO SUBMISSION ====================\n\ndef predict_multiclass_adaptive(pred, meta, action_thresholds, min_duration=3):\n    pred_smoothed = pd.DataFrame(pred, columns=pred.columns if isinstance(pred, pd.DataFrame) else range(pred.shape[1]))\n    pred_smoothed = pred_smoothed.rolling(window=5, min_periods=1, center=True).mean()\n    \n    ama = np.argmax(pred_smoothed.values, axis=1)\n    max_probs = pred_smoothed.values.max(axis=1)\n    \n    threshold_mask = np.zeros(len(pred_smoothed), dtype=bool)\n    for i, action in enumerate(pred_smoothed.columns):\n        action_mask = (ama == i)\n        threshold = action_thresholds.get(action, 0.27)\n        threshold_mask |= (action_mask & (max_probs >= threshold))\n    \n    ama = np.where(threshold_mask, ama, -1)\n    ama = pd.Series(ama, index=meta.video_frame.values if hasattr(meta, 'video_frame') else range(len(ama)))\n    \n    changes_mask = (ama != ama.shift(1)).values\n    ama_changes = ama[changes_mask]\n    meta_changes = meta[changes_mask]\n    mask = ama_changes.values >= 0\n    mask[-1] = False\n    \n    if mask.sum() == 0:\n        return pd.DataFrame(columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])\n    \n    submission_part = pd.DataFrame({\n        'video_id': meta_changes['video_id'].values[mask],\n        'agent_id': meta_changes['agent_id'].values[mask],\n        'target_id': meta_changes['target_id'].values[mask],\n        'action': pred_smoothed.columns[ama_changes.values[mask]],\n        'start_frame': ama_changes.index[mask],\n        'stop_frame': ama_changes.index[1:][mask[:-1]]\n    })\n    \n    for i in range(len(submission_part)):\n        video_id = submission_part.video_id.iloc[i]\n        if i < len(submission_part) - 1:\n            if submission_part.video_id.iloc[i+1] != video_id:\n                new_stop = meta.query(\"video_id == @video_id\").video_frame.max() + 1\n                submission_part.at[submission_part.index[i], 'stop_frame'] = new_stop\n        else:\n            new_stop = meta.query(\"video_id == @video_id\").video_frame.max() + 1\n            submission_part.at[submission_part.index[i], 'stop_frame'] = new_stop\n    \n    duration = submission_part.stop_frame - submission_part.start_frame\n    submission_part = submission_part[duration >= min_duration].reset_index(drop=True)\n    \n    return submission_part\n\n# ==================== MAIN TRAINING AND PREDICTION ====================\n\ndef process_with_dual_scale_cnn(body_parts_tracked_str, switch_tr, X_tr, label, meta):\n    body_parts_tracked = json.loads(body_parts_tracked_str)\n    if len(body_parts_tracked) > 5:\n        body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n    \n    X_tr_filled = X_tr.fillna(0).values\n    mean = X_tr_filled.mean(axis=0, keepdims=True)\n    std = X_tr_filled.std(axis=0, keepdims=True) + 1e-8\n    X_tr_normalized = (X_tr_filled - mean) / std\n    \n    action_thresholds = {}\n    models_dict = {}\n    \n    for action in label.columns:\n        action_mask = ~label[action].isna().values\n        y_action = label[action][action_mask].values\n        \n        if not (y_action == 0).all() and y_action.sum() >= 10:\n            pos_ratio = y_action.sum() / len(y_action) * 100\n            if verbose:\n                print(f\"  Training {action}: {int(y_action.sum())} pos ({pos_ratio:.2f}%)\")\n            \n            X_action = X_tr_normalized[action_mask]\n            \n            dataset = DualScaleDataset(X_action, y_action.reshape(-1, 1))\n            \n            train_size = int(0.8 * len(dataset))\n            val_size = len(dataset) - train_size\n            train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])\n            \n            train_labels = y_action[dataset.valid_indices[:train_size]]\n            sampler = create_balanced_sampler(train_labels, oversample_ratio=2.0)\n            \n            if sampler is not None:\n                train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler, num_workers=0)\n            else:\n                train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0)\n            \n            val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=0)\n            \n            model = DualScaleConv1D(input_dim=X_tr_normalized.shape[1], num_classes=1).to(device)\n            model = train_dual_model(model, train_loader, val_loader, epochs=15, lr=1e-3, patience=3)\n            \n            # Calculate optimal threshold\n            val_probs, val_labels = [], []\n            model.eval()\n            with torch.no_grad():\n                for local, long_range, labels in val_loader:\n                    local, long_range = local.to(device), long_range.to(device)\n                    outputs = model(local, long_range)\n                    probs = torch.sigmoid(outputs).cpu().numpy()\n                    val_probs.extend(probs[:, 0])\n                    val_labels.extend(labels.numpy()[:, 0])\n            \n            threshold = find_optimal_threshold(np.array(val_probs), np.array(val_labels))\n            if verbose:\n                print(f\"    Optimal threshold: {threshold:.3f}\")\n            \n            action_thresholds[action] = threshold\n            models_dict[action] = (model, mean, std)\n            \n            del train_loader, val_loader, dataset\n            gc.collect()\n            torch.cuda.empty_cache()\n    \n    # Generate test predictions\n    test_subset = test[test.body_parts_tracked == body_parts_tracked_str]\n    generator = generate_mouse_data(test_subset, 'test',\n                                    generate_single=(switch_tr == 'single'), \n                                    generate_pair=(switch_tr == 'pair'))\n    \n    submission_list = []\n    \n    for switch_te, data_te, meta_te, actions_te in generator:\n        try:\n            if switch_te == 'single':\n                X_te = transform_single(data_te, body_parts_tracked)\n            else:\n                X_te = transform_pair(data_te, body_parts_tracked)\n            \n            X_te_filled = X_te.fillna(0).values\n            pred = pd.DataFrame(index=meta_te.video_frame)\n            \n            for action in models_dict.keys():\n                if action in actions_te:\n                    model, mean, std = models_dict[action]\n                    X_te_normalized = (X_te_filled - mean) / std\n                    probs = predict_dual(model, X_te_normalized)\n                    pred[action] = probs[:, 0]\n            \n            if pred.shape[1] > 0:\n                sub_part = predict_multiclass_adaptive(pred, meta_te, action_thresholds)\n                submission_list.append(sub_part)\n            \n            del X_te, X_te_filled\n            gc.collect()\n            \n        except Exception as e:\n            if verbose:\n                print(f'  ERROR: {str(e)[:100]}')\n            gc.collect()\n    \n    return submission_list\n\n# ==================== MAIN EXECUTION ====================\n\nsubmission_list = []\n\nprint(f\"\\nProcessing {len(body_parts_tracked_list)} body part configurations...\\n\")\n\n# Process first 3 configurations for speed (increase to process all)\nfor section in range(min(3, len(body_parts_tracked_list))):\n    body_parts_tracked_str = body_parts_tracked_list[section]\n    \n    try:\n        body_parts_tracked = json.loads(body_parts_tracked_str)\n        print(f\"\\n{'='*70}\")\n        print(f\"Section {section+1}: {len(body_parts_tracked)} body parts\")\n        print('='*70)\n        \n        if len(body_parts_tracked) > 5:\n            body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts]\n        \n        train_subset = train[train.body_parts_tracked == body_parts_tracked_str]\n        \n        # Process single mouse behaviors\n        single_list, single_label_list, single_meta_list = [], [], []\n        \n        for switch, data, meta, label in generate_mouse_data(train_subset, 'train'):\n            if switch == 'single':\n                single_list.append(data)\n                single_label_list.append(label)\n                single_meta_list.append(meta)\n            if len(single_list) >= 10:  # Increase for better performance\n                break\n        \n        if len(single_list) > 0:\n            print(f\"\\nProcessing SINGLE mouse behaviors ({len(single_list)} sequences)\")\n            single_mouse = pd.concat(single_list)\n            single_label = pd.concat(single_label_list)\n            single_meta = pd.concat(single_meta_list)\n            \n            X_tr = transform_single(single_mouse, body_parts_tracked)\n            print(f\"Feature shape: {X_tr.shape}\")\n            \n            sub_parts = process_with_dual_scale_cnn(body_parts_tracked_str, 'single', X_tr, single_label, single_meta)\n            submission_list.extend(sub_parts)\n            \n            del single_mouse, single_label, single_meta, X_tr\n            gc.collect()\n        \n        # Process pair behaviors\n        pair_list, pair_label_list, pair_meta_list = [], [], []\n        \n        for switch, data, meta, label in generate_mouse_data(train_subset, 'train'):\n            if switch == 'pair':\n                pair_list.append(data)\n                pair_label_list.append(label)\n                pair_meta_list.append(meta)\n            if len(pair_list) >= 10:  # Increase for better performance\n                break\n        \n        if len(pair_list) > 0:\n            print(f\"\\nProcessing PAIR behaviors ({len(pair_list)} sequences)\")\n            mouse_pair = pd.concat(pair_list)\n            pair_label = pd.concat(pair_label_list)\n            pair_meta = pd.concat(pair_meta_list)\n            \n            X_tr = transform_pair(mouse_pair, body_parts_tracked)\n            print(f\"Feature shape: {X_tr.shape}\")\n            \n            sub_parts = process_with_dual_scale_cnn(body_parts_tracked_str, 'pair', X_tr, pair_label, pair_meta)\n            submission_list.extend(sub_parts)\n            \n            del mouse_pair, pair_label, pair_meta, X_tr\n            gc.collect()\n        \n    except Exception as e:\n        print(f'***Exception*** {str(e)[:200]}')\n    \n    gc.collect()\n    torch.cuda.empty_cache()\n\n# ==================== CREATE FINAL SUBMISSION ====================\n\nif len(submission_list) > 0:\n    submission = pd.concat(submission_list, ignore_index=True)\nelse:\n    submission = pd.DataFrame({\n        'video_id': [438887472],\n        'agent_id': ['mouse1'],\n        'target_id': ['self'],\n        'action': ['rear'],\n        'start_frame': [278],\n        'stop_frame': [500]\n    })\n\nsubmission = submission[submission.start_frame < submission.stop_frame].reset_index(drop=True)\nsubmission.index.name = 'row_id'\nsubmission.to_csv('submission.csv')\n\nprint(f\"\\n{'='*70}\")\nprint(\"SUBMISSION CREATED\")\nprint('='*70)\nprint(f\"Total predictions: {len(submission)}\")\nprint(f\"Unique actions: {submission.action.nunique()}\")\nprint(f\"Actions: {sorted(submission.action.unique())}\")\nprint(f\"Videos: {submission.video_id.nunique()}\")\nprint(f\"\\nSubmission saved to: submission.csv\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}