{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# PhysioNet v2 - Digitization of ECG Images","metadata":{}},{"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\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))\n        return nn.Sequential(*layers)\n    \n    def forward(self, x):\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        x = self.fusion(x)\n        \n        features = self.global_pool(x)\n        features = features.view(features.size(0), -1)\n        features = self.dropout(features)\n        \n        outputs = {}\n        for lead_name, decoder in self.lead_decoders.items():\n            outputs[lead_name] = decoder(features)\n        \n        return outputs\n\n# ================================\n# Enhanced Loss Function\n# ================================\nclass EnhancedSNRLoss(nn.Module):\n    \"\"\"Advanced loss function with dynamic weighting\"\"\"\n    \n    def __init__(self):\n        super().__init__()\n        self.mse = nn.MSELoss()\n        self.l1 = nn.L1Loss()\n        self.smooth_l1 = nn.SmoothL1Loss(beta=0.5)\n        self.signal_processor = SignalProcessor()\n    \n    def pearson_correlation_loss(self, pred, target):\n        \"\"\"Pearson correlation coefficient loss\"\"\"\n        pred_mean = pred.mean(dim=1, keepdim=True)\n        target_mean = target.mean(dim=1, keepdim=True)\n        \n        pred_centered = pred - pred_mean\n        target_centered = target - target_mean\n        \n        numerator = (pred_centered * target_centered).sum(dim=1)\n        denominator = torch.sqrt((pred_centered ** 2).sum(dim=1) * (target_centered ** 2).sum(dim=1))\n        \n        correlation = numerator / (denominator + 1e-6)\n        return (1 - correlation).mean()\n    \n    def forward(self, pred_dict, target_signals):\n        total_loss = 0\n        lead_names = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', \n                      'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n        \n        for idx, lead_name in enumerate(lead_names):\n            pred = pred_dict[lead_name]\n            target = target_signals[:, idx, :pred.size(1)]\n            \n            # Core losses\n            mse_loss = self.mse(pred, target)\n            smooth_loss = self.smooth_l1(pred, target)\n            \n            # Correlation loss\n            corr_loss = self.pearson_correlation_loss(pred, target)\n            \n            # Frequency domain loss\n            try:\n                pred_fft = torch.fft.rfft(pred, dim=1)\n                target_fft = torch.fft.rfft(target, dim=1)\n                \n                # Magnitude and phase losses\n                mag_loss = F.l1_loss(torch.abs(pred_fft), torch.abs(target_fft))\n                phase_loss = F.l1_loss(torch.angle(pred_fft), torch.angle(target_fft))\n                freq_loss = 0.7 * mag_loss + 0.3 * phase_loss\n            except:\n                freq_loss = 0\n            \n            # Gradient loss\n            pred_grad = pred[:, 1:] - pred[:, :-1]\n            target_grad = target[:, 1:] - target[:, :-1]\n            grad_loss = F.l1_loss(pred_grad, target_grad)\n            \n            # Second derivative loss (for smoothness)\n            pred_grad2 = pred_grad[:, 1:] - pred_grad[:, :-1]\n            target_grad2 = target_grad[:, 1:] - target_grad[:, :-1]\n            smooth_loss2 = F.l1_loss(pred_grad2, target_grad2)\n            \n            # Combine with adaptive weights\n            lead_loss = (0.2 * mse_loss + \n                        0.2 * smooth_loss + \n                        0.2 * corr_loss + \n                        0.15 * freq_loss +\n                        0.15 * grad_loss +\n                        0.1 * smooth_loss2)\n            \n            # Dynamic weighting for lead II\n            weight = 2.5 if lead_name == 'II' else 1.0\n            total_loss += weight * lead_loss\n        \n        return total_loss / (len(lead_names) + 1.5)\n\n# ================================\n# Training Functions\n# ================================\ndef train_epoch(model, loader, criterion, optimizer, scaler, device, epoch):\n    model.train()\n    total_loss = 0\n    \n    for batch_idx, (images, targets, _) in enumerate(loader):\n        images = images.to(device, non_blocking=True)\n        targets = targets.to(device, non_blocking=True)\n        \n        optimizer.zero_grad()\n        \n        if config.use_amp:\n            with autocast():\n                outputs = model(images)\n                loss = criterion(outputs, targets)\n            \n            scaler.scale(loss).backward()\n            scaler.unscale_(optimizer)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            outputs = model(images)\n            loss = criterion(outputs, targets)\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)\n            optimizer.step()\n        \n        total_loss += loss.item()\n    \n    return total_loss / len(loader)\n\n# ================================\n# Main Training\n# ================================\ndef main():\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    print(f\"Using device: {device}\")\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    print(f\"Training samples: {len(train_df)}\")\n    \n    # Enhanced augmentation\n    train_transform = A.Compose([\n        A.Resize(config.resize_height, config.resize_width),\n        A.OneOf([\n            A.GaussNoise(var_limit=(5, 40), p=1.0),\n            A.ISONoise(color_shift=(0.01, 0.04), intensity=(0.05, 0.4), p=1.0),\n            A.MultiplicativeNoise(multiplier=(0.95, 1.05), p=1.0),\n        ], p=0.35),\n        A.OneOf([\n            A.MotionBlur(blur_limit=3, p=1.0),\n            A.GaussianBlur(blur_limit=(3, 5), p=1.0),\n            A.MedianBlur(blur_limit=3, p=1.0),\n        ], p=0.25),\n        A.ShiftScaleRotate(\n            shift_limit=0.02,\n            scale_limit=0.03,\n            rotate_limit=1,\n            border_mode=cv2.BORDER_CONSTANT,\n            value=255,\n            p=0.3\n        ),\n        A.RandomBrightnessContrast(\n            brightness_limit=0.12,\n            contrast_limit=0.12,\n            p=0.35\n        ),\n        A.OneOf([\n            A.GridDistortion(num_steps=5, distort_limit=0.05, p=1.0),\n            A.ElasticTransform(alpha=1, sigma=10, p=1.0),\n        ], p=0.15),\n        A.CoarseDropout(\n            max_holes=6,\n            max_height=10,\n            max_width=10,\n            fill_value=255,\n            p=0.2\n        ),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n    \n    # Dataset and loader\n    train_dataset = ECGDataset(\n        train_df,\n        '/kaggle/input/physionet-ecg-image-digitization/train',\n        transform=train_transform,\n        is_train=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            # 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},"outputs":[],"execution_count":null}]}