{"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":"# ==================== COMPLETE IMPROVED 1D CNN SOLUTION WITH RESIDUAL CONNECTIONS ====================\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 OneCycleLR\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    \"\"\"Feature extraction for single mouse\"\"\"\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    \"\"\"Feature extraction for mouse pairs\"\"\"\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 EXTREME CLASS IMBALANCE ====================\n\nclass FocalLoss(nn.Module):\n    \"\"\"Focal Loss for addressing class imbalance\"\"\"\n    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):\n        super().__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.reduction = reduction\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        \n        if self.reduction == 'mean':\n            return focal_loss.mean()\n        elif self.reduction == 'sum':\n            return focal_loss.sum()\n        else:\n            return focal_loss\n\n# ==================== RESIDUAL BLOCK ====================\n\nclass ResidualBlock1D(nn.Module):\n    \"\"\"1D Residual block with batch normalization and dropout\"\"\"\n    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dropout=0.3):\n        super().__init__()\n        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, \n                               stride=stride, padding=kernel_size//2)\n        self.bn1 = nn.BatchNorm1d(out_channels)\n        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, \n                               stride=1, padding=kernel_size//2)\n        self.bn2 = nn.BatchNorm1d(out_channels)\n        self.dropout = nn.Dropout(dropout)\n        \n        self.shortcut = nn.Sequential()\n        if stride != 1 or in_channels != out_channels:\n            self.shortcut = nn.Sequential(\n                nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride),\n                nn.BatchNorm1d(out_channels)\n            )\n    \n    def forward(self, x):\n        residual = x\n        \n        out = F.relu(self.bn1(self.conv1(x)))\n        out = self.dropout(out)\n        out = self.bn2(self.conv2(out))\n        \n        out += self.shortcut(residual)\n        out = F.relu(out)\n        out = self.dropout(out)\n        \n        return out\n\n# ==================== IMPROVED CNN WITH RESIDUAL CONNECTIONS ====================\n\nclass ImprovedConv1DClassifier(nn.Module):\n    \"\"\"Enhanced 1D CNN with residual connections and multi-scale features\"\"\"\n    def __init__(self, input_dim, num_classes, window_size=60, dropout=0.3):\n        super().__init__()\n        self.window_size = window_size\n        \n        self.conv_init = nn.Conv1d(input_dim, 64, kernel_size=7, padding=3)\n        self.bn_init = nn.BatchNorm1d(64)\n        \n        self.res_block1 = ResidualBlock1D(64, 128, kernel_size=5, dropout=dropout)\n        self.res_block2 = ResidualBlock1D(128, 128, kernel_size=5, dropout=dropout)\n        self.pool1 = nn.MaxPool1d(2)\n        \n        self.res_block3 = ResidualBlock1D(128, 256, kernel_size=3, dropout=dropout)\n        self.res_block4 = ResidualBlock1D(256, 256, kernel_size=3, dropout=dropout)\n        self.pool2 = nn.MaxPool1d(2)\n        \n        self.res_block5 = ResidualBlock1D(256, 256, kernel_size=3, dropout=dropout)\n        self.res_block6 = ResidualBlock1D(256, 128, kernel_size=3, dropout=dropout)\n        \n        self.global_avg_pool = nn.AdaptiveAvgPool1d(1)\n        self.global_max_pool = nn.AdaptiveMaxPool1d(1)\n        \n        pooled_size = window_size // (2 ** 2)\n        \n        self.fc1 = nn.Linear(128 * pooled_size + 128 * 2, 512)\n        self.bn_fc1 = nn.BatchNorm1d(512)\n        self.dropout_fc = nn.Dropout(dropout)\n        \n        self.fc2 = nn.Linear(512, 256)\n        self.bn_fc2 = nn.BatchNorm1d(256)\n        \n        self.fc3 = nn.Linear(256, 128)\n        self.bn_fc3 = nn.BatchNorm1d(128)\n        \n        self.fc_out = nn.Linear(128, num_classes)\n        \n    def forward(self, x):\n        x = x.transpose(1, 2)\n        \n        x = F.relu(self.bn_init(self.conv_init(x)))\n        \n        x = self.res_block1(x)\n        x = self.res_block2(x)\n        x = self.pool1(x)\n        \n        x = self.res_block3(x)\n        x = self.res_block4(x)\n        x = self.pool2(x)\n        \n        x = self.res_block5(x)\n        x = self.res_block6(x)\n        \n        x_flat = x.flatten(1)\n        x_avg = self.global_avg_pool(x).squeeze(-1)\n        x_max = self.global_max_pool(x).squeeze(-1)\n        \n        x = torch.cat([x_flat, x_avg, x_max], dim=1)\n        \n        x = self.dropout_fc(F.relu(self.bn_fc1(self.fc1(x))))\n        x = self.dropout_fc(F.relu(self.bn_fc2(self.fc2(x))))\n        x = self.dropout_fc(F.relu(self.bn_fc3(self.fc3(x))))\n        x = self.fc_out(x)\n        \n        return x\n\n# ==================== IMPROVED DATASET WITH AUGMENTATION ====================\n\nclass ImprovedMouseBehaviorDataset(Dataset):\n    \"\"\"Dataset with support for weighted sampling and augmentation\"\"\"\n    def __init__(self, features, labels, window_size=60, augment=False):\n        self.features = features\n        self.labels = labels\n        self.window_size = window_size\n        self.half_window = window_size // 2\n        self.augment = augment\n        \n        self.valid_indices = list(range(self.half_window, len(features) - self.half_window))\n        self.sample_weights = self._calculate_weights()\n        \n    def _calculate_weights(self):\n        \"\"\"Calculate sample weights based on label frequency\"\"\"\n        weights = np.ones(len(self.valid_indices))\n        \n        for i, idx in enumerate(self.valid_indices):\n            if self.labels[idx, 0] == 1:\n                weights[i] = 10.0\n        \n        return weights\n    \n    def __len__(self):\n        return len(self.valid_indices)\n    \n    def __getitem__(self, idx):\n        center_idx = self.valid_indices[idx]\n        \n        window = self.features[center_idx - self.half_window:center_idx + self.half_window]\n        label = self.labels[center_idx]\n        \n        if self.augment and label[0] == 1:\n            window = self._augment(window)\n        \n        return torch.FloatTensor(window), torch.FloatTensor(label)\n    \n    def _augment(self, window):\n        \"\"\"Apply data augmentation\"\"\"\n        if np.random.rand() < 0.3:\n            noise = np.random.normal(0, 0.01, window.shape)\n            window = window + noise\n        \n        if np.random.rand() < 0.3:\n            factor = np.random.uniform(0.9, 1.1)\n            new_len = int(len(window) * factor)\n            indices = np.linspace(0, len(window)-1, new_len)\n            window = np.array([window[int(i)] for i in indices])\n            if len(window) < self.window_size:\n                pad = self.window_size - len(window)\n                window = np.pad(window, ((0, pad), (0, 0)), mode='edge')\n            else:\n                window = window[:self.window_size]\n        \n        return window\n\n# ==================== TRAINING FUNCTIONS ====================\n\ndef calculate_f1(preds, labels):\n    \"\"\"Helper function to calculate F1 score\"\"\"\n    tp = ((preds == 1) & (labels == 1)).sum().item()\n    fp = ((preds == 1) & (labels == 0)).sum().item()\n    fn = ((preds == 0) & (labels == 1)).sum().item()\n    \n    precision = tp / (tp + fp + 1e-8)\n    recall = tp / (tp + fn + 1e-8)\n    f1 = 2 * precision * recall / (precision + recall + 1e-8)\n    \n    return f1\n\ndef train_model_improved(model, train_loader, val_loader, num_epochs=25, lr=1e-3, \n                        patience=5, use_focal_loss=False, pos_weight=None):\n    \"\"\"Enhanced training with multiple strategies for class imbalance\"\"\"\n    \n    optimizer = AdamW(model.parameters(), lr=lr, weight_decay=0.01)\n    \n    scheduler = OneCycleLR(optimizer, max_lr=lr*10, \n                          steps_per_epoch=len(train_loader), \n                          epochs=num_epochs, pct_start=0.3)\n    \n    if use_focal_loss:\n        criterion = FocalLoss(alpha=0.25, gamma=2.0)\n        if verbose:\n            print(\"    Using Focal Loss\")\n    elif pos_weight is not None:\n        criterion = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([pos_weight]).to(device))\n        if verbose:\n            print(f\"    Using Weighted BCE Loss (weight={pos_weight:.2f})\")\n    else:\n        criterion = nn.BCEWithLogitsLoss()\n    \n    best_val_f1 = 0\n    patience_counter = 0\n    best_state = None\n    \n    for epoch in range(num_epochs):\n        model.train()\n        train_loss = 0\n        train_preds_all = []\n        train_labels_all = []\n        \n        for batch_x, batch_y in train_loader:\n            batch_x = batch_x.to(device)\n            batch_y = batch_y.to(device)\n            \n            optimizer.zero_grad()\n            outputs = model(batch_x)\n            loss = criterion(outputs, batch_y)\n            loss.backward()\n            \n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            optimizer.step()\n            scheduler.step()\n            \n            train_loss += loss.item()\n            \n            preds = (torch.sigmoid(outputs) > 0.5).float()\n            train_preds_all.append(preds.cpu())\n            train_labels_all.append(batch_y.cpu())\n        \n        train_loss /= len(train_loader)\n        \n        train_preds_all = torch.cat(train_preds_all)\n        train_labels_all = torch.cat(train_labels_all)\n        train_f1 = calculate_f1(train_preds_all, train_labels_all)\n        \n        model.eval()\n        val_loss = 0\n        all_preds = []\n        all_labels = []\n        \n        with torch.no_grad():\n            for batch_x, batch_y in val_loader:\n                batch_x = batch_x.to(device)\n                batch_y = batch_y.to(device)\n                outputs = model(batch_x)\n                loss = criterion(outputs, batch_y)\n                val_loss += loss.item()\n                \n                probs = torch.sigmoid(outputs)\n                preds = (probs > 0.5).float()\n                all_preds.append(preds.cpu())\n                all_labels.append(batch_y.cpu())\n        \n        val_loss /= len(val_loader)\n        \n        all_preds = torch.cat(all_preds)\n        all_labels = torch.cat(all_labels)\n        val_f1 = calculate_f1(all_preds, all_labels)\n        \n        if verbose:\n            print(f'Epoch {epoch+1}/{num_epochs} - Loss: {train_loss:.4f}/{val_loss:.4f}, '\n                  f'F1: {train_f1:.4f}/{val_f1:.4f}')\n        \n        if val_f1 > best_val_f1:\n            best_val_f1 = val_f1\n            patience_counter = 0\n            best_state = model.state_dict().copy()\n        else:\n            patience_counter += 1\n            if patience_counter >= patience:\n                if verbose:\n                    print(f'Early stopping at epoch {epoch+1}')\n                break\n    \n    if best_state is not None:\n        model.load_state_dict(best_state)\n    \n    return model\n\ndef find_best_threshold_advanced(model, val_loader, device):\n    \"\"\"Advanced threshold finding with multiple metrics\"\"\"\n    model.eval()\n    all_probs = []\n    all_labels = []\n    \n    with torch.no_grad():\n        for batch_x, batch_y in val_loader:\n            batch_x = batch_x.to(device)\n            outputs = torch.sigmoid(model(batch_x))\n            all_probs.append(outputs.cpu().numpy())\n            all_labels.append(batch_y.cpu().numpy())\n    \n    all_probs = np.concatenate(all_probs)\n    all_labels = np.concatenate(all_labels)\n    \n    best_f1 = 0\n    best_thresh = 0.5\n    best_metrics = {}\n    \n    for thresh in np.arange(0.05, 0.95, 0.025):\n        preds = (all_probs > thresh).astype(int)\n        \n        tp = ((preds == 1) & (all_labels == 1)).sum()\n        fp = ((preds == 1) & (all_labels == 0)).sum()\n        fn = ((preds == 0) & (all_labels == 1)).sum()\n        \n        precision = tp / (tp + fp + 1e-8)\n        recall = tp / (tp + fn + 1e-8)\n        f1 = 2 * precision * recall / (precision + recall + 1e-8)\n        \n        if f1 > best_f1:\n            best_f1 = f1\n            best_thresh = thresh\n            best_metrics = {\n                'precision': precision,\n                'recall': recall,\n                'f1': f1,\n                'tp': tp,\n                'fp': fp,\n                'fn': fn\n            }\n    \n    return best_thresh, best_f1, best_metrics\n\ndef predict_with_model(model, features, window_size=60):\n    \"\"\"Generate predictions for full sequence\"\"\"\n    model.eval()\n    half_window = window_size // 2\n    \n    all_probs = []\n    \n    with torch.no_grad():\n        for i in range(half_window, len(features) - half_window):\n            window = features[i - half_window:i + half_window]\n            window_tensor = torch.FloatTensor(window).unsqueeze(0).to(device)\n            \n            output = model(window_tensor)\n            probs = torch.sigmoid(output).cpu().numpy()[0]\n            all_probs.append(probs)\n    \n    all_probs = np.array(all_probs)\n    \n    full_probs = np.zeros((len(features), all_probs.shape[1]))\n    full_probs[half_window:-half_window] = all_probs\n    full_probs[:half_window] = all_probs[0]\n    full_probs[-half_window:] = all_probs[-1]\n    \n    return full_probs\n\n# ==================== PREDICTION TO SUBMISSION ====================\n\ndef predict_multiclass_adaptive(pred, meta, action_thresholds, min_duration=3):\n    \"\"\"Convert frame predictions to submission format\"\"\"\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 PROCESSING FUNCTION ====================\n\ndef process_with_improved_cnn(body_parts_tracked_str, switch_tr, X_tr, label, meta):\n    \"\"\"Train improved CNN with residual connections and better class imbalance handling\"\"\"\n    \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    \n    X_tr_clipped = np.clip(X_tr_filled, \n                           np.percentile(X_tr_filled, 1), \n                           np.percentile(X_tr_filled, 99))\n    \n    mean = X_tr_clipped.mean(axis=0, keepdims=True)\n    std = X_tr_clipped.std(axis=0, keepdims=True) + 1e-8\n    X_tr_normalized = (X_tr_clipped - 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() >= 5:\n            pos_count = y_action.sum()\n            neg_count = (y_action == 0).sum()\n            pos_weight = neg_count / pos_count\n            imbalance_ratio = neg_count / pos_count\n            \n            if verbose:\n                print(f\"  Training Improved CNN for action: {action}\")\n                print(f\"    Positive: {int(pos_count)}, Negative: {int(neg_count)}, \"\n                      f\"Ratio: 1:{imbalance_ratio:.1f}\")\n            \n            X_action = X_tr_normalized[action_mask]\n            \n            use_focal = imbalance_ratio > 100\n            use_augmentation = pos_count < 500\n            \n            dataset = ImprovedMouseBehaviorDataset(\n                X_action, y_action.reshape(-1, 1), \n                window_size=60,\n                augment=use_augmentation\n            )\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(\n                dataset, [train_size, val_size],\n                generator=torch.Generator().manual_seed(42)\n            )\n            \n            if imbalance_ratio > 50:\n                train_indices = train_dataset.indices\n                train_weights = [dataset.sample_weights[i] for i in train_indices]\n                sampler = WeightedRandomSampler(\n                    weights=train_weights,\n                    num_samples=len(train_weights),\n                    replacement=True\n                )\n                train_loader = DataLoader(train_dataset, batch_size=32, \n                                        sampler=sampler, num_workers=0)\n            else:\n                train_loader = DataLoader(train_dataset, batch_size=32, \n                                        shuffle=True, num_workers=0)\n            \n            val_loader = DataLoader(val_dataset, batch_size=64, \n                                  shuffle=False, num_workers=0)\n            \n            model = ImprovedConv1DClassifier(\n                input_dim=X_tr_normalized.shape[1],\n                num_classes=1,\n                window_size=60,\n                dropout=0.3\n            ).to(device)\n            \n            model = train_model_improved(\n                model, train_loader, val_loader,\n                num_epochs=25 if pos_count < 200 else 20,\n                lr=5e-4 if pos_count < 100 else 1e-3,\n                patience=5,\n                use_focal_loss=use_focal,\n                pos_weight=pos_weight if not use_focal else None\n            )\n            \n            best_thresh, best_f1, metrics = find_best_threshold_advanced(\n                model, val_loader, device\n            )\n            action_thresholds[action] = best_thresh\n            \n            if verbose:\n                print(f\"    Best threshold: {best_thresh:.3f}, Val F1: {best_f1:.4f}\")\n                print(f\"    Precision: {metrics['precision']:.3f}, \"\n                      f\"Recall: {metrics['recall']:.3f}\")\n            \n            models_dict[action] = (model, mean, std)\n            \n            gc.collect()\n            torch.cuda.empty_cache()\n        \n        elif verbose:\n            print(f\"  Skipping {action}: insufficient positive samples ({int(y_action.sum())})\")\n    \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        assert switch_te == switch_tr\n        \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            X_te_clipped = np.clip(X_te_filled,\n                                  np.percentile(X_tr_filled, 1),\n                                  np.percentile(X_tr_filled, 99))\n            \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_clipped - mean) / std\n                    probs = predict_with_model(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\nfor section in range(1, 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{section}. Processing: {len(body_parts_tracked)} body parts\")\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        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_meta_list.append(meta)\n                single_label_list.append(label)\n            if len(single_list) >= 5:\n                break\n        \n        if len(single_list) > 0:\n            print(f\"  Single mouse sequences: {len(single_list)}\")\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\"  Single features shape: {X_tr.shape}\")\n            \n            sub_parts = process_with_improved_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        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_meta_list.append(meta)\n                pair_label_list.append(label)\n            if len(pair_list) >= 5:\n                break\n        \n        if len(pair_list) > 0:\n            print(f\"  Pair sequences: {len(pair_list)}\")\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\"  Pair features shape: {X_tr.shape}\")\n            \n            sub_parts = process_with_improved_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)[:100]}')\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✓ Submission created: {len(submission)} predictions\")\nprint(f\"  Actions found: {submission.action.nunique()}\")\nprint(f\"  Videos: {submission.video_id.nunique()}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}