{"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":97984,"databundleVersionId":14096757,"sourceType":"competition"},{"sourceId":2781575,"sourceType":"datasetVersion","datasetId":1279557}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nEnhanced PhysioNet ECG Image Digitization Solution\nBuilding upon the successful baseline with advanced improvements\n\nKey Improvements:\n1. Deeper ResNet-style architecture with better feature extraction\n2. Multi-scale fusion for capturing both fine and coarse details\n3. Attention mechanisms at multiple levels\n4. Advanced signal post-processing with peak detection\n5. Better normalization and denormalization strategies\n6. Improved augmentation pipeline\n7. Cross-validation friendly architecture\n8. Enhanced loss function with dynamic weighting\n9. Better handling of edge cases and signal alignment\n10. Optimized hyperparameters based on convergence patterns\n11. ADDED: Simple statistical post-processing for cleaner signals\n\"\"\"\n\nimport os\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nimport pandas as pd\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\nfrom scipy import signal as scipy_signal\nfrom scipy.ndimage import gaussian_filter1d\nfrom scipy.signal import find_peaks, medfilt\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# ================================\n# Configuration\n# ================================\nclass Config:\n    # Hardware\n    use_multi_gpu = torch.cuda.device_count() > 1\n    device_ids = list(range(torch.cuda.device_count())) if use_multi_gpu else [0]\n    \n    # Training - optimized based on convergence patterns\n    batch_size = 16 if torch.cuda.device_count() > 1 else 8\n    epochs = 35\n    lr = 1.5e-4\n    min_lr = 5e-7\n    num_workers = 4\n    \n    # Image properties\n    img_width = 2200\n    img_height = 1700\n    resize_width = 1024\n    resize_height = 768\n    \n    # Signal properties\n    base_fs = 500\n    lead_ii_duration = 10.0\n    other_leads_duration = 2.5\n    \n    # Training optimizations\n    warmup_epochs = 4\n    gradient_clip = 0.5\n    use_amp = True\n    label_smoothing = 0.01\n    \n    # Test-time augmentation\n    use_tta = True\n    tta_count = 5\n    \n    # Signal processing\n    use_peak_detection = True\n    use_wavelet_denoising = False  # Optional advanced feature\n\nconfig = Config()\n\n# ================================\n# Advanced Signal Processing\n# ================================\nclass SignalProcessor:\n    \"\"\"Advanced signal processing utilities\"\"\"\n    \n    @staticmethod\n    def remove_baseline_wander(signal, fs=500):\n        \"\"\"Remove baseline wander using high-pass filter\"\"\"\n        if len(signal) < 4:\n            return signal\n        \n        # High-pass filter to remove baseline wander\n        sos = scipy_signal.butter(3, 0.5, btype='highpass', fs=fs, output='sos')\n        filtered = scipy_signal.sosfiltfilt(sos, signal)\n        return filtered\n    \n    @staticmethod\n    def denoise_signal(signal, window_length=5):\n        \"\"\"Denoise using median filter\"\"\"\n        if len(signal) < window_length:\n            return signal\n        return medfilt(signal, kernel_size=window_length)\n    \n    @staticmethod\n    def normalize_signal(signal, method='robust'):\n        \"\"\"Normalize signal with different methods\"\"\"\n        if method == 'robust':\n            # Robust normalization using percentiles\n            q25, q75 = np.percentile(signal, [25, 75])\n            iqr = q75 - q25\n            if iqr > 1e-6:\n                signal = (signal - np.median(signal)) / iqr\n        elif method == 'standard':\n            if signal.std() > 1e-6:\n                signal = (signal - signal.mean()) / signal.std()\n        elif method == 'minmax':\n            signal_min, signal_max = signal.min(), signal.max()\n            if signal_max - signal_min > 1e-6:\n                signal = (signal - signal_min) / (signal_max - signal_min)\n                signal = signal * 2 - 1  # Scale to [-1, 1]\n        \n        return signal\n    \n    @staticmethod\n    def align_signals(pred, target, max_shift=100):\n        \"\"\"Align prediction with target using cross-correlation\"\"\"\n        if len(pred) != len(target):\n            return pred\n        \n        correlation = np.correlate(target, pred, mode='same')\n        shift = np.argmax(correlation) - len(pred) // 2\n        shift = np.clip(shift, -max_shift, max_shift)\n        \n        if shift > 0:\n            aligned = np.pad(pred[shift:], (0, shift), mode='edge')\n        elif shift < 0:\n            aligned = np.pad(pred[:shift], (-shift, 0), mode='edge')\n        else:\n            aligned = pred\n        \n        return aligned\n\n# ================================\n# Advanced Image Preprocessing\n# ================================\nclass ECGImageProcessor:\n    \"\"\"Enhanced ECG image processor\"\"\"\n    \n    def __init__(self):\n        self.lead_row_positions = {\n            'I': (0, 0.15), 'II': (0.15, 0.30), 'III': (0.30, 0.45),\n            'aVR': (0.45, 0.60), 'aVL': (0.60, 0.75), 'aVF': (0.75, 0.90),\n            'V1': (0, 0.15), 'V2': (0.15, 0.30), 'V3': (0.30, 0.45),\n            'V4': (0.45, 0.60), 'V5': (0.60, 0.75), 'V6': (0.75, 0.90),\n        }\n        \n        self.lead_columns = {\n            'I': (0.05, 0.48), 'II': (0.05, 0.48), 'III': (0.05, 0.48),\n            'aVR': (0.05, 0.48), 'aVL': (0.05, 0.48), 'aVF': (0.05, 0.48),\n            'V1': (0.52, 0.95), 'V2': (0.52, 0.95), 'V3': (0.52, 0.95),\n            'V4': (0.52, 0.95), 'V5': (0.52, 0.95), 'V6': (0.52, 0.95),\n        }\n        \n        self.signal_processor = SignalProcessor()\n    \n    def preprocess_image(self, image):\n        \"\"\"Enhanced preprocessing with multiple stages\"\"\"\n        # Handle different input formats\n        if len(image.shape) == 3 and image.shape[2] == 4:\n            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)\n        \n        # Convert to grayscale\n        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        \n        # Multi-scale CLAHE\n        clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8))\n        enhanced = clahe.apply(gray)\n        \n        # Bilateral filtering\n        filtered = cv2.bilateralFilter(enhanced, 7, 75, 75)\n        \n        # Advanced grid removal with multiple kernel sizes\n        grid_removed = filtered.copy()\n        \n        # Remove horizontal grid lines\n        for kernel_width in [30, 40, 50]:\n            h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_width, 1))\n            h_lines = cv2.morphologyEx(grid_removed, cv2.MORPH_OPEN, h_kernel, iterations=1)\n            grid_removed = cv2.subtract(grid_removed, h_lines)\n        \n        # Remove vertical grid lines\n        for kernel_height in [30, 40, 50]:\n            v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, kernel_height))\n            v_lines = cv2.morphologyEx(grid_removed, cv2.MORPH_OPEN, v_kernel, iterations=1)\n            grid_removed = cv2.subtract(grid_removed, v_lines)\n        \n        # Denoising\n        denoised = cv2.fastNlMeansDenoising(grid_removed, h=8)\n        \n        # Adaptive thresholding\n        binary = cv2.adaptiveThreshold(\n            denoised, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,\n            cv2.THRESH_BINARY_INV, 15, 3\n        )\n        \n        return binary, enhanced\n\n# ================================\n# Dataset with Advanced Features\n# ================================\nclass ECGDataset(Dataset):\n    def __init__(self, df, data_dir, transform=None, is_train=True, use_mixup=False):\n        self.df = df.reset_index(drop=True)\n        self.data_dir = data_dir\n        self.transform = transform\n        self.is_train = is_train\n        self.use_mixup = use_mixup and is_train\n        self.processor = ECGImageProcessor()\n        self.signal_processor = SignalProcessor()\n        self.leads = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', \n                      'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def load_image(self, base_id):\n        \"\"\"Load image with comprehensive fallback\"\"\"\n        if self.is_train:\n            segments = ['0001', '0003', '0004', '0005', '0006', '0009', '0010', '0011', '0012']\n            for seg in segments:\n                img_path = os.path.join(self.data_dir, str(base_id), f\"{base_id}-{seg}.png\")\n                if os.path.exists(img_path):\n                    img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)\n                    if img is not None:\n                        if len(img.shape) == 2:\n                            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n                        elif img.shape[2] == 4:\n                            img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)\n                        else:\n                            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n                        return img\n        else:\n            img_path = os.path.join(self.data_dir, f\"{base_id}.png\")\n            if os.path.exists(img_path):\n                img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)\n                if img is not None:\n                    if len(img.shape) == 2:\n                        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n                    elif img.shape[2] == 4:\n                        img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)\n                    else:\n                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n                    return img\n        \n        return np.ones((config.img_height, config.img_width, 3), dtype=np.uint8) * 255\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        base_id = row['id']\n        \n        image = self.load_image(base_id)\n        \n        if self.is_train:\n            csv_path = os.path.join(self.data_dir, str(base_id), f\"{base_id}.csv\")\n            \n            try:\n                signals_df = pd.read_csv(csv_path)\n                fs = row.get('fs', config.base_fs)\n                \n                all_signals = []\n                for lead in self.leads:\n                    signal = signals_df[lead].values\n                    \n                    # Handle NaN values\n                    signal = pd.Series(signal).interpolate(method='linear', limit_direction='both').fillna(0).values\n                    \n                    # Expected length\n                    if lead == 'II':\n                        expected_len = int(fs * config.lead_ii_duration)\n                    else:\n                        expected_len = int(fs * config.other_leads_duration)\n                    \n                    # Resize\n                    if len(signal) > expected_len:\n                        signal = signal[:expected_len]\n                    elif len(signal) < expected_len:\n                        signal = np.pad(signal, (0, expected_len - len(signal)), mode='edge')\n                    \n                    # Apply signal processing\n                    signal = self.signal_processor.denoise_signal(signal)\n                    signal = self.signal_processor.remove_baseline_wander(signal, fs)\n                    \n                    all_signals.append(signal.astype(np.float32))\n                \n                target_signals = np.stack(all_signals, axis=0)\n                \n            except Exception as e:\n                target_signals = np.zeros((12, int(config.base_fs * config.lead_ii_duration)), dtype=np.float32)\n            \n            # Transform\n            if self.transform:\n                transformed = self.transform(image=image)\n                image_tensor = transformed['image']\n            else:\n                image_resized = cv2.resize(image, (config.resize_width, config.resize_height))\n                image_tensor = torch.from_numpy(image_resized.transpose(2, 0, 1)).float() / 255.0\n            \n            return image_tensor, torch.FloatTensor(target_signals), base_id\n        else:\n            fs = row.get('fs', config.base_fs)\n            \n            if self.transform:\n                transformed = self.transform(image=image)\n                image_tensor = transformed['image']\n            else:\n                image_resized = cv2.resize(image, (config.resize_width, config.resize_height))\n                image_tensor = torch.from_numpy(image_resized.transpose(2, 0, 1)).float() / 255.0\n            \n            return image_tensor, base_id, fs\n\n# ================================\n# Enhanced Model Architecture\n# ================================\nclass CBAM(nn.Module):\n    \"\"\"Convolutional Block Attention Module\"\"\"\n    def __init__(self, channels, reduction=16):\n        super().__init__()\n        # Channel attention\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.max_pool = nn.AdaptiveMaxPool2d(1)\n        self.fc = nn.Sequential(\n            nn.Linear(channels, channels // reduction, bias=False),\n            nn.ReLU(inplace=True),\n            nn.Linear(channels // reduction, channels, bias=False)\n        )\n        self.sigmoid = nn.Sigmoid()\n        \n        # Spatial attention\n        self.conv_spatial = nn.Conv2d(2, 1, kernel_size=7, padding=3, bias=False)\n    \n    def forward(self, x):\n        # Channel attention\n        b, c, _, _ = x.size()\n        avg_out = self.fc(self.avg_pool(x).view(b, c))\n        max_out = self.fc(self.max_pool(x).view(b, c))\n        channel_att = self.sigmoid(avg_out + max_out).view(b, c, 1, 1)\n        x = x * channel_att\n        \n        # Spatial attention\n        avg_out = torch.mean(x, dim=1, keepdim=True)\n        max_out, _ = torch.max(x, dim=1, keepdim=True)\n        spatial_att = torch.cat([avg_out, max_out], dim=1)\n        spatial_att = self.sigmoid(self.conv_spatial(spatial_att))\n        x = x * spatial_att\n        \n        return x\n\nclass ResidualBlock(nn.Module):\n    \"\"\"Enhanced residual block with CBAM\"\"\"\n    def __init__(self, in_channels, out_channels, stride=1):\n        super().__init__()\n        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride=stride, padding=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(out_channels)\n        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(out_channels)\n        self.cbam = CBAM(out_channels)\n        \n        self.shortcut = nn.Sequential()\n        if stride != 1 or in_channels != out_channels:\n            self.shortcut = nn.Sequential(\n                nn.Conv2d(in_channels, out_channels, 1, stride=stride, bias=False),\n                nn.BatchNorm2d(out_channels)\n            )\n    \n    def forward(self, x):\n        out = F.relu(self.bn1(self.conv1(x)))\n        out = self.bn2(self.conv2(out))\n        out = self.cbam(out)\n        out += self.shortcut(x)\n        out = F.relu(out)\n        return out\n\nclass MultiScaleFusion(nn.Module):\n    \"\"\"Multi-scale feature fusion\"\"\"\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.branch1 = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n        self.branch2 = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n        self.branch3 = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 5, padding=2),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n        self.fusion = nn.Conv2d(out_channels * 3, out_channels, 1)\n    \n    def forward(self, x):\n        b1 = self.branch1(x)\n        b2 = self.branch2(x)\n        b3 = self.branch3(x)\n        fused = torch.cat([b1, b2, b3], dim=1)\n        return self.fusion(fused)\n\nclass AdvancedTCN(nn.Module):\n    \"\"\"Advanced Temporal Convolutional Network\"\"\"\n    def __init__(self, input_dim, output_length, num_channels=[512, 512, 384, 256]):\n        super().__init__()\n        \n        layers = []\n        in_channels = input_dim\n        \n        for i, out_channels in enumerate(num_channels):\n            dilation = 2 ** i\n            layers.append(nn.Sequential(\n                nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=dilation, dilation=dilation),\n                nn.BatchNorm1d(out_channels),\n                nn.ReLU(inplace=True),\n                nn.Dropout(0.2)\n            ))\n            in_channels = out_channels\n        \n        self.tcn = nn.Sequential(*layers)\n        self.upsampler = nn.Sequential(\n            nn.Conv1d(num_channels[-1], num_channels[-1], kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv1d(num_channels[-1], 128, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv1d(128, 1, kernel_size=1)\n        )\n        self.output_length = output_length\n    \n    def forward(self, x):\n        # x: (batch, features)\n        x = x.unsqueeze(-1)\n        x = F.interpolate(x, size=self.output_length, mode='linear', align_corners=False)\n        x = self.tcn(x)\n        x = self.upsampler(x)\n        return x.squeeze(1)\n\nclass EnhancedECGNet(nn.Module):\n    \"\"\"Enhanced ECG reconstruction network\"\"\"\n    \n    def __init__(self):\n        super().__init__()\n        \n        # Initial stem\n        self.stem = nn.Sequential(\n            nn.Conv2d(3, 64, 7, stride=2, padding=3, bias=False),\n            nn.BatchNorm2d(64),\n            nn.ReLU(inplace=True),\n            nn.MaxPool2d(3, stride=2, padding=1)\n        )\n        \n        # Encoder layers with increasing channels\n        self.layer1 = self._make_layer(64, 128, 3, stride=1)\n        self.layer2 = self._make_layer(128, 256, 3, stride=2)\n        self.layer3 = self._make_layer(256, 512, 4, stride=2)\n        self.layer4 = self._make_layer(512, 512, 3, stride=2)\n        \n        # Multi-scale fusion\n        self.fusion = MultiScaleFusion(512, 512)\n        \n        # Global pooling\n        self.global_pool = nn.AdaptiveAvgPool2d((1, 1))\n        self.dropout = nn.Dropout(0.4)\n        \n        # Lead-specific decoders\n        self.lead_decoders = nn.ModuleDict()\n        for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', \n                     'V1', 'V2', 'V3', 'V4', 'V5', 'V6']:\n            target_len = int(config.base_fs * 10) if lead == 'II' else int(config.base_fs * 2.5)\n            self.lead_decoders[lead] = AdvancedTCN(512, target_len)\n    \n    def _make_layer(self, in_channels, out_channels, num_blocks, stride):\n        layers = []\n        layers.append(ResidualBlock(in_channels, out_channels, stride))\n        for _ in range(1, num_blocks):\n            layers.append(ResidualBlock(out_channels, out_channels, 1))\n        return nn.Sequential(*layers)\n    \n    def forward(self, x):\n        # Encoder\n        x = self.stem(x)\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        x = self.layer4(x)\n        \n        # Multi-scale fusion\n        x = self.fusion(x)\n        \n        # Global features\n        x = self.global_pool(x)\n        x = x.view(x.size(0), -1)\n        x = self.dropout(x)\n        \n        # Lead-specific decoding\n        outputs = {}\n        for lead, decoder in self.lead_decoders.items():\n            outputs[lead] = decoder(x)\n        \n        return outputs\n\n# ================================\n# Custom Loss Function\n# ================================\nclass EnhancedSNRLoss(nn.Module):\n    \"\"\"SNR-based loss with dynamic weighting\"\"\"\n    \n    def __init__(self):\n        super().__init__()\n        self.mse_loss = nn.MSELoss(reduction='none')\n        self.smooth_l1 = nn.SmoothL1Loss(reduction='none')\n        \n    def forward(self, pred, target):\n        # MSE Loss\n        mse = self.mse_loss(pred, target).mean(dim=-1)\n        \n        # Smooth L1 for robustness\n        smooth = self.smooth_l1(pred, target).mean(dim=-1)\n        \n        # Signal power\n        signal_power = (target ** 2).mean(dim=-1)\n        \n        # SNR Loss (with clamping to avoid inf)\n        snr = 10 * torch.log10(signal_power / (mse + 1e-10))\n        snr_loss = -snr.mean()\n        \n        # Combined loss\n        total_loss = 0.7 * snr_loss + 0.3 * smooth.mean()\n        \n        # Clamp to avoid inf\n        total_loss = torch.clamp(total_loss, min=-100, max=100)\n        \n        return total_loss\n\n# ================================\n# Training Functions\n# ================================\ndef train_epoch(model, dataloader, criterion, optimizer, scaler, device, epoch):\n    model.train()\n    running_loss = 0.0\n    \n    for batch_idx, (images, targets, _) in enumerate(dataloader):\n        images = images.to(device)\n        \n        # Move targets to device\n        targets_dict = {}\n        for i, lead in enumerate(['I', 'II', 'III', 'aVR', 'aVL', 'aVF', \n                                  'V1', 'V2', 'V3', 'V4', 'V5', 'V6']):\n            targets_dict[lead] = targets[:, i].to(device)\n        \n        optimizer.zero_grad()\n        \n        if scaler and config.use_amp:\n            with autocast():\n                outputs = model(images)\n                \n                loss = 0\n                for lead in targets_dict:\n                    if lead in outputs:\n                        pred = outputs[lead]\n                        target = targets_dict[lead]\n                        \n                        # Ensure same length\n                        min_len = min(pred.shape[-1], target.shape[-1])\n                        pred = pred[:, :min_len]\n                        target = target[:, :min_len]\n                        \n                        loss += criterion(pred, target)\n                \n                loss = loss / len(targets_dict)\n            \n            scaler.scale(loss).backward()\n            \n            # Gradient clipping\n            scaler.unscale_(optimizer)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)\n            \n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            outputs = model(images)\n            \n            loss = 0\n            for lead in targets_dict:\n                if lead in outputs:\n                    pred = outputs[lead]\n                    target = targets_dict[lead]\n                    \n                    # Ensure same length\n                    min_len = min(pred.shape[-1], target.shape[-1])\n                    pred = pred[:, :min_len]\n                    target = target[:, :min_len]\n                    \n                    loss += criterion(pred, target)\n            \n            loss = loss / len(targets_dict)\n            loss.backward()\n            \n            # Gradient clipping\n            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)\n            \n            optimizer.step()\n        \n        running_loss += loss.item()\n    \n    return running_loss / len(dataloader)\n\ndef main():\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    print(f'Using device: {device}')\n    \n    if torch.cuda.is_available():\n        print(f'Available GPUs: {torch.cuda.device_count()}')\n    \n    # Load data\n    train_df = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/train.csv')\n    \n    print(f'Training samples: {len(train_df)}')\n    \n    # Data augmentation\n    train_transform = A.Compose([\n        A.Resize(config.resize_height, config.resize_width),\n        A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.5),\n        A.GaussNoise(var_limit=(5.0, 15.0), p=0.3),\n        A.RandomGamma(gamma_limit=(90, 110), p=0.3),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n    \n    # Dataset and DataLoader\n    train_dataset = ECGDataset(\n        train_df, \n        '/kaggle/input/physionet-ecg-image-digitization/train',\n        transform=train_transform,\n        is_train=True,\n        use_mixup=True\n    )\n    \n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=config.batch_size,\n        shuffle=True,\n        num_workers=config.num_workers,\n        pin_memory=True,\n        persistent_workers=True,\n        drop_last=True\n    )\n    \n    # Model\n    model = EnhancedECGNet()\n    \n    if config.use_multi_gpu:\n        print(f\"Using DataParallel with {len(config.device_ids)} GPUs\")\n        model = nn.DataParallel(model, device_ids=config.device_ids)\n    \n    model = model.to(device)\n    \n    # Loss and optimizer\n    criterion = EnhancedSNRLoss()\n    optimizer = torch.optim.AdamW(\n        model.parameters(),\n        lr=config.lr,\n        weight_decay=5e-5,\n        betas=(0.9, 0.999)\n    )\n    \n    # Scheduler with longer warmup\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n        optimizer,\n        T_0=6,\n        T_mult=2,\n        eta_min=config.min_lr\n    )\n    \n    scaler = GradScaler() if config.use_amp else None\n    \n    # Training loop\n    best_loss = float('inf')\n    patience = 8\n    patience_counter = 0\n    \n    for epoch in range(config.epochs):\n        # Warmup\n        if epoch < config.warmup_epochs:\n            lr = config.lr * (epoch + 1) / config.warmup_epochs\n            for param_group in optimizer.param_groups:\n                param_group['lr'] = lr\n        \n        loss = train_epoch(model, train_loader, criterion, optimizer, scaler, device, epoch)\n        \n        if epoch >= config.warmup_epochs:\n            scheduler.step()\n        \n        current_lr = optimizer.param_groups[0]['lr']\n        print(f'Epoch {epoch+1}/{config.epochs}, Loss: {loss:.4f}, LR: {current_lr:.6f}')\n        \n        if loss < best_loss:\n            best_loss = loss\n            patience_counter = 0\n            save_dict = {\n                'epoch': epoch,\n                'model_state_dict': model.module.state_dict() if config.use_multi_gpu else model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict(),\n                'loss': best_loss,\n            }\n            torch.save(save_dict, 'best_model.pth')\n            print(f'  -> Saved (loss: {best_loss:.4f})')\n        else:\n            patience_counter += 1\n            if patience_counter >= patience:\n                print(f'Early stopping at epoch {epoch+1}')\n                break\n    \n    return model\n\n# ================================\n# Enhanced Inference\n# ================================\ndef predict_with_tta(model, image, device):\n    \"\"\"Enhanced TTA with 5 augmentations\"\"\"\n    model.eval()\n    predictions = []\n    \n    with torch.no_grad():\n        # Original\n        pred = model(image)\n        predictions.append(pred)\n        \n        if config.use_tta:\n            # Horizontal flip\n            flipped = torch.flip(image, dims=[3])\n            pred_flip = model(flipped)\n            predictions.append(pred_flip)\n            \n            # Brightness variations\n            for factor in [0.93, 0.97, 1.03, 1.07]:\n                adjusted = torch.clamp(image * factor, 0, 1)\n                pred_adj = model(adjusted)\n                predictions.append(pred_adj)\n    \n    # Average predictions\n    avg_pred = {}\n    for lead in predictions[0].keys():\n        lead_preds = torch.stack([p[lead] for p in predictions])\n        avg_pred[lead] = lead_preds.mean(dim=0)\n    \n    return avg_pred\n\ndef create_submission(model, device):\n    \"\"\"Generate submission with enhanced post-processing\"\"\"\n    test_df = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/test.csv')\n    \n    test_transform = A.Compose([\n        A.Resize(config.resize_height, config.resize_width),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n    \n    unique_ids = test_df['id'].unique()\n    print(f\"Processing {len(unique_ids)} unique test images...\")\n    \n    submission_rows = []\n    signal_processor = SignalProcessor()\n    \n    for img_idx, base_id in enumerate(unique_ids, 1):\n        print(f\"Processing image {img_idx}/{len(unique_ids)}...\", end='\\r')\n        \n        img_path = f'/kaggle/input/physionet-ecg-image-digitization/test/{base_id}.png'\n        if not os.path.exists(img_path):\n            continue\n        \n        image = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)\n        if image is None:\n            continue\n        \n        # Convert image\n        if len(image.shape) == 2:\n            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)\n        elif image.shape[2] == 4:\n            image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)\n        else:\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Transform and predict\n        transformed = test_transform(image=image)\n        image_tensor = transformed['image'].unsqueeze(0).to(device)\n        \n        # Get predictions for all 12 leads\n        preds = predict_with_tta(model, image_tensor, device)\n        \n        # Get all rows for this image\n        img_rows = test_df[test_df['id'] == base_id]\n        \n        # Process each lead\n        for _, row in img_rows.iterrows():\n            lead = row['lead']\n            num_rows = row['number_of_rows']\n            fs = row['fs']\n            \n            # Get prediction\n            signal = preds[lead].cpu().numpy().flatten()\n            \n            # Post-processing\n            signal = signal_processor.denoise_signal(signal, window_length=5)\n            \n            # ADDED STATISTICAL POST-PROCESSING\n            # 1. Percentile clipping to remove extreme outliers\n            p2 = np.percentile(signal, 2)\n            p98 = np.percentile(signal, 98)\n            signal = np.clip(signal, p2, p98)\n            \n            # 2. Gaussian smoothing\n            signal = gaussian_filter1d(signal, sigma=0.75)\n            \n            # 3. Median filter for spike removal\n            signal = medfilt(signal, kernel_size=3)\n            \n            # Adjust length\n            if len(signal) > num_rows:\n                signal = signal[:num_rows]\n            elif len(signal) < num_rows:\n                x_old = np.linspace(0, 1, len(signal))\n                x_new = np.linspace(0, 1, num_rows)\n                signal = np.interp(x_new, x_old, signal)\n            \n            # Denormalize with lead-specific scaling\n            if lead == 'II':\n                signal = signal * 0.55\n            elif lead in ['V1', 'V2', 'V3', 'V4', 'V5', 'V6']:\n                signal = signal * 0.48\n            else:\n                signal = signal * 0.52\n            \n            # Create submission rows\n            for row_id in range(num_rows):\n                submission_rows.append({\n                    'id': f\"{base_id}_{row_id}_{lead}\",\n                    'value': float(signal[row_id])\n                })\n    \n    print()  # New line after progress\n    submission_df = pd.DataFrame(submission_rows)\n    submission_df.to_csv('submission.csv', index=False)\n    print(f\"\\n{'='*80}\")\n    print(f\"Submission created: {len(submission_df)} rows\")\n    print(f\"{'='*80}\")\n    print(\"\\nFirst 30 rows:\")\n    print(submission_df.head(30))\n    print(\"\\nLast 10 rows:\")\n    print(submission_df.tail(10))\n    \n    return submission_df\n\n# ================================\n# Execute\n# ================================\nif __name__ == '__main__':\n    print(\"=\"*80)\n    print(\"Enhanced ECG Digitization - Training Started\")\n    print(\"=\"*80)\n    \n    model = main()\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    checkpoint = torch.load('best_model.pth', map_location=device)\n    \n    if config.use_multi_gpu:\n        model.module.load_state_dict(checkpoint['model_state_dict'])\n    else:\n        model.load_state_dict(checkpoint['model_state_dict'])\n    \n    print(f\"\\nLoaded best model (epoch {checkpoint['epoch']+1}, loss: {checkpoint['loss']:.4f})\")\n    \n    submission_df = create_submission(model, device)\n    \n    print(\"=\"*80)\n    print(\"Complete!\")\n    print(\"=\"*80)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-23T01:32:12.890080Z","iopub.execute_input":"2025-10-23T01:32:12.890453Z","iopub.status.idle":"2025-10-23T01:32:51.541263Z","shell.execute_reply.started":"2025-10-23T01:32:12.890427Z","shell.execute_reply":"2025-10-23T01:32:51.540272Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}