{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14584593,"sourceType":"datasetVersion","datasetId":9316525}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Model path","metadata":{}},{"cell_type":"code","source":"import torch\nimport os\n\nMODEL_PATH = \"/kaggle/input/ecg-digitization-nnunet-best-pth/best_model.pth\"\nprint(\"Model exists:\", os.path.exists(MODEL_PATH))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:54:23.704178Z","iopub.execute_input":"2026-01-26T22:54:23.704554Z","iopub.status.idle":"2026-01-26T22:54:27.917074Z","shell.execute_reply.started":"2026-01-26T22:54:23.704522Z","shell.execute_reply":"2026-01-26T22:54:27.916314Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Pipeline","metadata":{}},{"cell_type":"code","source":"\"\"\"\nCOMPLETE ECG DIGITIZATION PIPELINE\n\nPipeline Flow:\n1. GEOMETRY (Rule-based) → Normalize layout, define WHERE each lead exists\n2. SEGMENTATION (Deep Learning) → Learn WHICH pixels are ECG traces\n3. SIGNAL EXTRACTION (Algorithmic) → Convert pixels to waveforms\n\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport cv2\nfrom pathlib import Path\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nfrom scipy.interpolate import splprep, splev\nfrom skimage.morphology import skeletonize\nimport scipy.signal\nimport pandas as pd\n\n# ============================================================================\n# PART 1: GEOMETRIC PREPROCESSOR (From mask_testing_2.py)\n# ============================================================================\n\nclass ECGGeometryConfig:\n    \"\"\"Fixed configuration for ECG geometry\"\"\"\n    CANONICAL_WIDTH = 2100\n    CANONICAL_HEIGHT = 1500\n    CANONICAL_SIZE = (CANONICAL_WIDTH, CANONICAL_HEIGHT)\n    \n    # Calibration constants from training data\n    CALIB_H = 1700.0\n    LEAD_CENTER_RATIOS = [\n        716.5 / CALIB_H,\n        989.5 / CALIB_H,\n        1273.0 / CALIB_H,\n        1543.0 / CALIB_H\n    ]\n    \n    CALIB_W = 2200.0\n    LEAD_X_BOUNDARY_RATIOS = [\n        0.0,\n        610.0 / CALIB_W,\n        1103.0 / CALIB_W,\n        1595.5 / CALIB_W,\n        2089.5 / CALIB_W,\n        1.0\n    ]\n    \n    LEAD_LAYOUT = {\n        (0, 0): 1, (0, 1): 2, (0, 2): 3, (0, 3): 4,\n        (1, 0): 5, (1, 1): 6, (1, 2): 7, (1, 3): 8,\n        (2, 0): 9, (2, 1): 10, (2, 2): 11, (2, 3): 12,\n    }\n\n\nclass GeometricPreprocessor:\n    \"\"\"\n    Wraps geometric mask generation into reusable module\n    \n    Why separate from learning:\n    - ECG layout is FIXED by medical standards (12-lead grid)\n    - Position doesn't vary - only trace appearance varies\n    - Learning layout is wasteful and error-prone\n    \"\"\"\n    \n    def __init__(self):\n        self.config = ECGGeometryConfig()\n    \n    def find_paper_boundary(self, image):\n        \"\"\"Enhanced paper boundary detection\"\"\"\n        h, w = image.shape[:2]\n        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        blurred = cv2.GaussianBlur(gray, (15, 15), 0)\n        edges = cv2.Canny(blurred, 30, 100)\n        \n        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))\n        edges = cv2.dilate(edges, kernel, iterations=2)\n        \n        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n        \n        if not contours:\n            return self._get_full_image_corners(h, w)\n        \n        # Filter valid contours\n        min_area = (h * w) * 0.20\n        max_area = (h * w) * 0.95\n        \n        for contour in contours:\n            area = cv2.contourArea(contour)\n            if area < min_area or area > max_area:\n                continue\n            \n            peri = cv2.arcLength(contour, True)\n            approx = cv2.approxPolyDP(contour, 0.02 * peri, True)\n            \n            if len(approx) == 4:\n                corners = approx.reshape(4, 2)\n                return self._order_points(corners.astype(np.float32))\n        \n        return self._get_full_image_corners(h, w)\n    \n    def _get_full_image_corners(self, h, w):\n        \"\"\"Fallback: use entire image\"\"\"\n        return np.array([\n            [0, 0], [w-1, 0], [w-1, h-1], [0, h-1]\n        ], dtype=np.float32)\n    \n    def _order_points(self, pts):\n        \"\"\"Order points as [TL, TR, BR, BL]\"\"\"\n        rect = np.zeros((4, 2), dtype=np.float32)\n        s = pts.sum(axis=1)\n        rect[0] = pts[np.argmin(s)]  # Top-left\n        rect[2] = pts[np.argmax(s)]  # Bottom-right\n        diff = np.diff(pts, axis=1)\n        rect[1] = pts[np.argmin(diff)]  # Top-right\n        rect[3] = pts[np.argmax(diff)]  # Bottom-left\n        return rect\n    \n    def warp_to_canonical(self, image, corners):\n        \"\"\"Apply perspective transform to canonical size\"\"\"\n        w, h = self.config.CANONICAL_SIZE\n        \n        dst_pts = np.array([\n            [0, 0], [w-1, 0], [w-1, h-1], [0, h-1]\n        ], dtype=np.float32)\n        \n        M = cv2.getPerspectiveTransform(corners, dst_pts)\n        \n        warped = cv2.warpPerspective(\n            image, M, (w, h),\n            flags=cv2.INTER_LINEAR,\n            borderMode=cv2.BORDER_CONSTANT,\n            borderValue=(255, 255, 255)\n        )\n        \n        return warped\n    \n    def generate_geometric_mask(self):\n        \"\"\"\n        Generate fixed geometric mask based on medical ECG layout\n        \n        This mask defines WHERE each lead exists (spatial regions)\n        nnU-Net will learn WHICH pixels within each region are traces\n        \"\"\"\n        cfg = self.config\n        H, W = cfg.CANONICAL_HEIGHT, cfg.CANONICAL_WIDTH\n        mask = np.zeros((H, W), dtype=np.uint8)\n        \n        row_height = int((cfg.LEAD_CENTER_RATIOS[1] - cfg.LEAD_CENTER_RATIOS[0]) * H)\n        \n        # Create 3x4 grid for 12 leads\n        for row in range(3):\n            cy = int(cfg.LEAD_CENTER_RATIOS[row] * H)\n            y1 = cy - row_height // 2\n            y2 = cy + row_height // 2\n            \n            for col in range(4):\n                lead_id = cfg.LEAD_LAYOUT[(row, col)]\n                x1 = int(cfg.LEAD_X_BOUNDARY_RATIOS[col] * W)\n                x2 = int(cfg.LEAD_X_BOUNDARY_RATIOS[col + 1] * W)\n                mask[y1:y2, x1:x2] = lead_id\n        \n        # Rhythm strip (Lead II at bottom)\n        last_row_center = int(cfg.LEAD_CENTER_RATIOS[2] * H)\n        rhythm_y1 = last_row_center + row_height // 2\n        rhythm_y2 = int(H * (1.0 - 0.035))  # Leave 3.5% margin for metadata\n        mask[rhythm_y1:rhythm_y2, :] = 13  # Label 13 for rhythm strip\n        \n        return mask\n    \n    def process(self, image_path):\n        \"\"\"\n        Complete geometric preprocessing pipeline\n        \n        Input: Raw ECG image (any size, any perspective)\n        Output: \n            - canonical: Normalized image (2100x1500)\n            - geometric_mask: Fixed layout mask\n        \"\"\"\n        # Load image\n        image = cv2.imread(str(image_path))\n        if image is None:\n            raise ValueError(f\"Could not load: {image_path}\")\n        \n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Detect paper boundary\n        corners = self.find_paper_boundary(image)\n        \n        # Warp to canonical size\n        canonical = self.warp_to_canonical(image, corners)\n        \n        # Generate fixed geometric mask\n        geometric_mask = self.generate_geometric_mask()\n        \n        return {\n            'canonical': canonical,\n            'geometric_mask': geometric_mask,\n            'original': image,\n            'corners': corners\n        }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:54:27.919077Z","iopub.execute_input":"2026-01-26T22:54:27.919756Z","iopub.status.idle":"2026-01-26T22:54:29.593563Z","shell.execute_reply.started":"2026-01-26T22:54:27.919730Z","shell.execute_reply":"2026-01-26T22:54:29.592711Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# PART 2: DATASET WITH GEOMETRIC PREPROCESSING\n# ============================================================================\n\nclass ECGDatasetWithGeometry(Dataset):\n    \"\"\"\n    OPTIMIZED Dataset with caching for faster training\n    \n    CRITICAL OPTIMIZATION:\n    - Preprocess ALL images once at initialization\n    - Cache canonical images + masks in memory\n    - __getitem__ just loads from cache (100x faster)\n    \"\"\"\n    \n    def __init__(self, image_paths, target_size=(384, 384), is_train=True, cache_preprocessed=True):\n        \"\"\"\n        Args:\n            image_paths: List of paths to ECG images\n            target_size: Size for nnU-Net input (reduced to 384x384 for speed)\n            is_train: Training mode flag\n            cache_preprocessed: Cache preprocessed data (HUGE speedup)\n        \"\"\"\n        self.image_paths = image_paths\n        self.target_size = target_size\n        self.is_train = is_train\n        self.cache_preprocessed = cache_preprocessed\n        \n        # OPTIMIZATION: Preprocess all images at initialization\n        if cache_preprocessed:\n            print(f\"⚡ Preprocessing {'train' if is_train else 'val'} images (one-time cost)...\")\n            self.cache = self._preprocess_all_images()\n            print(f\"✓ Cached {len(self.cache)} images in memory\")\n        else:\n            self.preprocessor = GeometricPreprocessor()\n            self.cache = None\n        \n        print(f\"{'Train' if is_train else 'Val'} dataset: {len(image_paths)} images\")\n    \n    def _preprocess_all_images(self):\n        \"\"\"Preprocess all images once and cache in memory\"\"\"\n        preprocessor = GeometricPreprocessor()\n        cache = []\n        \n        for img_path in tqdm(self.image_paths, desc=\"Caching\"):\n            try:\n                # Geometric preprocessing\n                result = preprocessor.process(img_path)\n                canonical = result['canonical']\n                geometric_mask = result['geometric_mask']\n                \n                # Resize to target size\n                image = cv2.resize(canonical, self.target_size, interpolation=cv2.INTER_LINEAR)\n                mask = cv2.resize(geometric_mask, self.target_size, interpolation=cv2.INTER_NEAREST)\n                \n                # Normalize\n                image = image.astype(np.float32) / 255.0\n                \n                cache.append({\n                    'image': image,\n                    'mask': mask,\n                    'path': str(img_path)\n                })\n            except Exception as e:\n                print(f\"⚠️ Failed to preprocess {img_path}: {e}\")\n        \n        return cache\n    \n    def __len__(self):\n        return len(self.cache) if self.cache else len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        \"\"\"\n        Returns cached data (very fast!)\n        \"\"\"\n        if self.cache:\n            # Load from cache (FAST)\n            item = self.cache[idx]\n            image = item['image'].copy()\n            mask = item['mask'].copy()\n        else:\n            # Process on-the-fly (SLOW - only for debugging)\n            img_path = self.image_paths[idx]\n            result = self.preprocessor.process(img_path)\n            canonical = result['canonical']\n            geometric_mask = result['geometric_mask']\n            \n            image = cv2.resize(canonical, self.target_size, interpolation=cv2.INTER_LINEAR)\n            mask = cv2.resize(geometric_mask, self.target_size, interpolation=cv2.INTER_NEAREST)\n            image = image.astype(np.float32) / 255.0\n        \n        # Data augmentation (only for training)\n        if self.is_train and np.random.rand() < 0.3:\n            # Light augmentation\n            if np.random.rand() < 0.5:\n                # Horizontal flip\n                image = np.fliplr(image).copy()\n                mask = np.fliplr(mask).copy()\n        \n        # Convert to tensors\n        image = np.transpose(image, (2, 0, 1))\n        \n        return {\n            'image': torch.FloatTensor(image),\n            'mask': torch.LongTensor(mask),\n            'path': str(self.image_paths[idx]) if not self.cache else self.cache[idx]['path']\n        }\n\n\ndef create_dataloaders(image_dir, val_split=0.2, batch_size=16, num_workers=4):\n    \"\"\"\n    OPTIMIZED dataloader creation\n    \n    CRITICAL OPTIMIZATIONS:\n    - Increased batch_size to 16 (was 4)\n    - Increased num_workers to 4 (was 2)\n    - Reduced image size to 384x384 (was 512x512)\n    - Cache preprocessing results\n    \"\"\"\n    image_dir = Path(image_dir)\n    \n    # Find all ECG images\n    image_paths = sorted(list(image_dir.rglob('*.png')) + list(image_dir.rglob('*.jpg')))\n    \n    if not image_paths:\n        print(f\"No images found in {image_dir}\")\n        return None, None\n    \n    # Split at image level\n    n_val = int(len(image_paths) * val_split)\n    np.random.seed(42)\n    np.random.shuffle(image_paths)\n    \n    val_paths = image_paths[:n_val]\n    train_paths = image_paths[n_val:]\n    \n    # Create datasets WITH CACHING\n    train_dataset = ECGDatasetWithGeometry(\n        train_paths, \n        target_size=(384, 384),  # Reduced from 512x512\n        is_train=True,\n        cache_preprocessed=True  # CRITICAL: Cache in memory\n    )\n    val_dataset = ECGDatasetWithGeometry(\n        val_paths, \n        target_size=(384, 384),\n        is_train=False,\n        cache_preprocessed=True\n    )\n    \n    # Create dataloaders with OPTIMIZED settings\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=batch_size,  # Increased to 16\n        shuffle=True,\n        num_workers=num_workers,  # Increased to 4\n        pin_memory=True,\n        persistent_workers=True  # Keep workers alive between epochs\n    )\n    \n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=num_workers,\n        pin_memory=True,\n        persistent_workers=True\n    )\n    \n    print(f\"\\n✓ OPTIMIZED dataloaders created:\")\n    print(f\"  Train: {len(train_dataset)} images, {len(train_loader)} batches\")\n    print(f\"  Val:   {len(val_dataset)} images, {len(val_loader)} batches\")\n    print(f\"  Batch size: {batch_size} (optimized for T4)\")\n    print(f\"  Image size: 384x384 (reduced for speed)\")\n    \n    return train_loader, val_loader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:54:29.594607Z","iopub.execute_input":"2026-01-26T22:54:29.595076Z","iopub.status.idle":"2026-01-26T22:54:29.610706Z","shell.execute_reply.started":"2026-01-26T22:54:29.595051Z","shell.execute_reply":"2026-01-26T22:54:29.610060Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# ============================================================================\n# PART 3: nnU-Net MODEL (Learns trace refinement only)\n# ============================================================================\n\nclass DoubleConv(nn.Module):\n    \"\"\"Double convolution block\"\"\"\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 3, padding=1),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.01, inplace=True),\n            nn.Conv2d(out_channels, out_channels, 3, padding=1),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.01, inplace=True)\n        )\n    \n    def forward(self, x):\n        return self.conv(x)\n\n\nclass nnUNet2D(nn.Module):\n    \"\"\"\n    OPTIMIZED nnU-Net for ECG trace refinement\n    \n    CRITICAL OPTIMIZATIONS:\n    - Reduced channels: [24, 48, 96, 192, 256] (was [32, 64, 128, 256, 512])\n    - GroupNorm instead of InstanceNorm (faster)\n    - Fewer parameters = faster training\n    \"\"\"\n    \n    def __init__(self, in_channels=3, num_classes=14):\n        super().__init__()\n        \n        # OPTIMIZED: Smaller channels\n        channels = [24, 48, 96, 192, 256]\n        \n        # Encoder\n        self.enc1 = self._make_layer(in_channels, channels[0])\n        self.enc2 = self._make_layer(channels[0], channels[1])\n        self.enc3 = self._make_layer(channels[1], channels[2])\n        self.enc4 = self._make_layer(channels[2], channels[3])\n        \n        # Bottleneck\n        self.bottleneck = self._make_layer(channels[3], channels[4])\n        \n        # Decoder\n        self.upconv4 = nn.ConvTranspose2d(channels[4], channels[3], 2, stride=2)\n        self.dec4 = self._make_layer(channels[3]* 2, channels[3])\n        \n        self.upconv3 = nn.ConvTranspose2d(channels[3], channels[2], 2, stride=2)\n        self.dec3 = self._make_layer(channels[2]* 2, channels[2])\n        \n        self.upconv2 = nn.ConvTranspose2d(channels[2], channels[1], 2, stride=2)\n        self.dec2 = self._make_layer(channels[1]* 2, channels[1])\n        \n        self.upconv1 = nn.ConvTranspose2d(channels[1], channels[0], 2, stride=2)\n        self.dec1 = self._make_layer(channels[0]* 2, channels[0])\n        \n        # Output\n        self.final = nn.Conv2d(channels[0], num_classes, 1)\n        self.pool = nn.MaxPool2d(2)\n    \n    def _make_layer(self, in_channels, out_channels):\n        \"\"\"Optimized conv block with GroupNorm\"\"\"\n        return nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 3, padding=1),\n            nn.GroupNorm(8, out_channels),  # Faster than InstanceNorm\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, 3, padding=1),\n            nn.GroupNorm(8, out_channels),\n            nn.ReLU(inplace=True)\n        )\n    \n    def forward(self, x):\n        # Encoder\n        enc1 = self.enc1(x)\n        enc2 = self.enc2(self.pool(enc1))\n        enc3 = self.enc3(self.pool(enc2))\n        enc4 = self.enc4(self.pool(enc3))\n        \n        # Bottleneck\n        bottleneck = self.bottleneck(self.pool(enc4))\n        \n        # Decoder\n        dec4 = self.upconv4(bottleneck)\n        dec4 = torch.cat([dec4, enc4], dim=1)\n        dec4 = self.dec4(dec4)\n        \n        dec3 = self.upconv3(dec4)\n        dec3 = torch.cat([dec3, enc3], dim=1)\n        dec3 = self.dec3(dec3)\n        \n        dec2 = self.upconv2(dec3)\n        dec2 = torch.cat([dec2, enc2], dim=1)\n        dec2 = self.dec2(dec2)\n        \n        dec1 = self.upconv1(dec2)\n        dec1 = torch.cat([dec1, enc1], dim=1)\n        dec1 = self.dec1(dec1)\n        \n        return self.final(dec1)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:54:29.611827Z","iopub.execute_input":"2026-01-26T22:54:29.612218Z","iopub.status.idle":"2026-01-26T22:54:29.639419Z","shell.execute_reply.started":"2026-01-26T22:54:29.612183Z","shell.execute_reply":"2026-01-26T22:54:29.638591Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# PART 4: MASK FUSION (Combines geometry + learning)\n# ============================================================================\n\ndef fuse_masks(predicted_mask, geometric_mask):\n    \"\"\"\n    Fuse predicted and geometric masks\n    \n    Rule: \n    - Where model detected traces (predicted > 0), use prediction\n    - Where model detected nothing, fall back to geometry\n    \n    This ensures:\n    - Geometry guarantees coverage\n    - Model refines trace precision\n    \"\"\"\n    final_mask = np.where(\n        predicted_mask > 0,\n        predicted_mask,\n        geometric_mask\n    )\n    \n    return final_mask","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:54:29.644214Z","iopub.execute_input":"2026-01-26T22:54:29.644518Z","iopub.status.idle":"2026-01-26T22:54:29.659278Z","shell.execute_reply.started":"2026-01-26T22:54:29.644478Z","shell.execute_reply":"2026-01-26T22:54:29.658467Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# PART 5: WAVEFORM EXTRACTION (From waveform extractor)\n# ============================================================================\n\nclass DPWaveformExtractor:\n    \"\"\"\n    Dynamic Programming waveform extractor\n    \n    Converts fused mask → precise 1D waveforms\n    \"\"\"\n    \n    def __init__(self):\n        self.config = {\n            'TARGET_FREQUENCY': 1000,  # Hz\n            'LEAD_II_SAMPLES': 10000,  # 10 seconds\n            'OTHER_LEAD_SAMPLES': 2500,  # 2.5 seconds\n            'OUTPUT_VOLTAGE_RANGE': (-2.0, 2.0)  # mV\n        }\n    \n    def extract_all_leads(self, image, mask):\n        \"\"\"\n        Extract waveforms for all 12 leads\n        \n        Args:\n            image: Canonical ECG image (H, W, 3) or (H, W)\n            mask: Fused segmentation mask (H, W)\n        \n        Returns:\n            dict: {lead_id: waveform_array}\n        \"\"\"\n        # Convert to grayscale\n        if len(image.shape) == 3:\n            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n        else:\n            gray = image.copy()\n        \n        inverted = 255 - gray\n        \n        waveforms = {}\n        \n        for lead_id in range(1, 13):\n            lead_mask = (mask == lead_id)\n            \n            if not lead_mask.any():\n                target_len = self._get_target_length(lead_id)\n                waveforms[lead_id] = np.zeros(target_len)\n                continue\n            \n            waveform = self._extract_single_lead(inverted, lead_mask, lead_id)\n            waveforms[lead_id] = waveform\n        \n        return waveforms\n    \n    def _extract_single_lead(self, inverted, lead_mask, lead_id):\n        \"\"\"Extract single lead using DP\"\"\"\n        rows, cols = np.where(lead_mask)\n        \n        if len(rows) == 0:\n            target_len = self._get_target_length(lead_id)\n            return np.zeros(target_len)\n        \n        y_min, y_max = rows.min(), rows.max()\n        x_min, x_max = cols.min(), cols.max()\n        \n        sub_image = inverted[y_min:y_max+1, x_min:x_max+1]\n        sub_mask = lead_mask[y_min:y_max+1, x_min:x_max+1]\n        \n        # Simplified extraction (use centerline)\n        path_y = self._find_centerline(sub_image, sub_mask)\n        \n        signal = self._path_to_signal(path_y, sub_image.shape[0])\n        \n        target_len = self._get_target_length(lead_id)\n        signal = scipy.signal.resample(signal, target_len)\n        \n        signal = self._normalize_signal(signal)\n        \n        return signal\n    \n    def _find_centerline(self, sub_image, sub_mask):\n        \"\"\"Find vertical center for each x position\"\"\"\n        H, W = sub_image.shape\n        path_y = np.zeros(W, dtype=int)\n        \n        for x in range(W):\n            col = sub_mask[:, x]\n            if col.any():\n                path_y[x] = np.where(col)[0].mean().astype(int)\n            else:\n                path_y[x] = H // 2\n        \n        return path_y\n    \n    def _path_to_signal(self, path_y, height):\n        \"\"\"Convert path to signal\"\"\"\n        return (height - path_y.astype(np.float32))\n    \n    def _normalize_signal(self, signal):\n        \"\"\"Normalize to mV range\"\"\"\n        s_min, s_max = signal.min(), signal.max()\n        \n        if s_max - s_min < 1e-6:\n            return np.zeros_like(signal)\n        \n        normalized = (signal - s_min) / (s_max - s_min)\n        \n        v_min, v_max = self.config['OUTPUT_VOLTAGE_RANGE']\n        return normalized * (v_max - v_min) + v_min\n    \n    def _get_target_length(self, lead_id):\n        \"\"\"Get target length for lead\"\"\"\n        if lead_id == 2:  # Lead II\n            return self.config['LEAD_II_SAMPLES']\n        else:\n            return self.config['OTHER_LEAD_SAMPLES']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:55:17.101223Z","iopub.execute_input":"2026-01-26T22:55:17.101802Z","iopub.status.idle":"2026-01-26T22:55:17.113040Z","shell.execute_reply.started":"2026-01-26T22:55:17.101776Z","shell.execute_reply":"2026-01-26T22:55:17.112319Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# ============================================================================\n# PART 6: TRAINING LOOP WITH COMPETITION DATASET\n# ============================================================================\n\nclass CombinedLoss(nn.Module):\n    \"\"\"Combined Dice + Cross Entropy Loss\"\"\"\n    \n    def __init__(self, dice_weight=0.5, ce_weight=0.5):\n        super().__init__()\n        self.dice_weight = dice_weight\n        self.ce_weight = ce_weight\n    \n    def dice_loss(self, pred, target):\n        \"\"\"Dice Loss for segmentation\"\"\"\n        smooth = 1.0\n        \n        # Convert target to one-hot\n        target_one_hot = F.one_hot(target, num_classes=pred.shape[1])\n        target_one_hot = target_one_hot.permute(0, 3, 1, 2).float()\n        \n        # Apply softmax to predictions\n        pred = F.softmax(pred, dim=1)\n        \n        # Compute Dice\n        intersection = (pred * target_one_hot).sum(dim=(2, 3))\n        union = pred.sum(dim=(2, 3)) + target_one_hot.sum(dim=(2, 3))\n        \n        dice = (2. * intersection + smooth) / (union + smooth)\n        \n        return 1 - dice.mean()\n    \n    def forward(self, pred, target):\n        dice = self.dice_loss(pred, target)\n        ce = F.cross_entropy(pred, target)\n        return self.dice_weight * dice + self.ce_weight * ce\n\n\nclass Trainer:\n    \"\"\"\n    OPTIMIZED Training manager for multi-GPU\n    \n    CRITICAL OPTIMIZATIONS:\n    - Support for DataParallel (multi-GPU)\n    - Mixed precision training (FP16)\n    - Gradient accumulation\n    - Faster validation\n    \"\"\"\n    \n    def __init__(self, model, train_loader, val_loader, device, save_dir='checkpoints', use_amp=True):\n        # OPTIMIZATION: Multi-GPU support\n        if torch.cuda.device_count() > 1:\n            print(f\"🚀 Using {torch.cuda.device_count()} GPUs!\")\n            model = nn.DataParallel(model)\n        \n        self.model = model.to(device)\n        self.train_loader = train_loader\n        self.val_loader = val_loader\n        self.device = device\n        self.save_dir = Path(save_dir)\n        self.save_dir.mkdir(parents=True, exist_ok=True)\n        \n        # OPTIMIZATION: Mixed precision training\n        self.use_amp = use_amp and torch.cuda.is_available()\n        self.scaler = torch.cuda.amp.GradScaler() if self.use_amp else None\n        \n        # Loss and optimizer\n        self.criterion = CombinedLoss()\n        self.optimizer = torch.optim.AdamW(\n            model.parameters(),\n            lr=2e-4,  # Slightly higher LR\n            weight_decay=1e-5,\n            betas=(0.9, 0.999)\n        )\n        \n        # OPTIMIZATION: Cosine annealing scheduler (better than ReduceLROnPlateau)\n        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            self.optimizer, \n            T_max=50,  # Total epochs\n            eta_min=1e-6\n        )\n        \n        # Training state\n        self.current_epoch = 0\n        self.best_val_loss = float('inf')\n        self.train_losses = []\n        self.val_losses = []\n        \n        print(f\"✓ Trainer initialized:\")\n        print(f\"  Mixed precision: {'ON' if self.use_amp else 'OFF'}\")\n        print(f\"  Multi-GPU: {torch.cuda.device_count() if torch.cuda.device_count() > 1 else 'OFF'}\")\n    \n    def train_epoch(self):\n        \"\"\"OPTIMIZED training epoch with AMP\"\"\"\n        self.model.train()\n        total_loss = 0\n        \n        pbar = tqdm(self.train_loader, desc=f\"Epoch {self.current_epoch+1}\")\n        \n        for batch in pbar:\n            images = batch['image'].to(self.device, non_blocking=True)\n            masks = batch['mask'].to(self.device, non_blocking=True)\n            \n            self.optimizer.zero_grad(set_to_none=True)  # Faster than zero_grad()\n            \n            # OPTIMIZATION: Mixed precision forward pass\n            if self.use_amp:\n                with torch.cuda.amp.autocast():\n                    outputs = self.model(images)\n                    loss = self.criterion(outputs, masks)\n                \n                # Backward with scaling\n                self.scaler.scale(loss).backward()\n                self.scaler.unscale_(self.optimizer)\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)\n                self.scaler.step(self.optimizer)\n                self.scaler.update()\n            else:\n                # Standard training\n                outputs = self.model(images)\n                loss = self.criterion(outputs, masks)\n                loss.backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)\n                self.optimizer.step()\n            \n            total_loss += loss.item()\n            pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n        \n        avg_loss = total_loss / len(self.train_loader)\n        self.train_losses.append(avg_loss)\n        \n        return avg_loss\n    \n    def validate(self):\n        \"\"\"OPTIMIZED validation with AMP\"\"\"\n        self.model.eval()\n        total_loss = 0\n        \n        with torch.no_grad():\n            for batch in tqdm(self.val_loader, desc=\"Validating\", leave=False):\n                images = batch['image'].to(self.device, non_blocking=True)\n                masks = batch['mask'].to(self.device, non_blocking=True)\n                \n                # Use AMP for validation too\n                if self.use_amp:\n                    with torch.cuda.amp.autocast():\n                        outputs = self.model(images)\n                        loss = self.criterion(outputs, masks)\n                else:\n                    outputs = self.model(images)\n                    loss = self.criterion(outputs, masks)\n                \n                total_loss += loss.item()\n        \n        avg_loss = total_loss / len(self.val_loader)\n        self.val_losses.append(avg_loss)\n        \n        return avg_loss\n    \n    def save_checkpoint(self, is_best=False):\n        \"\"\"Save model checkpoint\"\"\"\n        # Handle DataParallel\n        model_state = self.model.module.state_dict() if hasattr(self.model, 'module') else self.model.state_dict()\n        \n        checkpoint = {\n            'epoch': self.current_epoch,\n            'model_state_dict': model_state,\n            'optimizer_state_dict': self.optimizer.state_dict(),\n            'train_losses': self.train_losses,\n            'val_losses': self.val_losses,\n            'best_val_loss': self.best_val_loss\n        }\n        \n        # Save latest\n        torch.save(checkpoint, self.save_dir / 'latest.pth')\n        \n        # Save best\n        if is_best:\n            torch.save(checkpoint, self.save_dir / 'best.pth')\n            print(f\"  💾 Saved best model (val_loss: {self.best_val_loss:.4f})\")\n    \n    def train(self, num_epochs=50, early_stopping_patience=15):\n        \"\"\"OPTIMIZED main training loop\"\"\"\n        print(f\"\\n🚀 Starting OPTIMIZED training for {num_epochs} epochs...\")\n        print(f\"   Device: {self.device}\")\n        print(f\"   Batch size: {self.train_loader.batch_size}\")\n        print(f\"   Images per epoch: {len(self.train_loader.dataset)}\")\n        \n        epochs_without_improvement = 0\n        \n        for epoch in range(num_epochs):\n            self.current_epoch = epoch\n            \n            # Train\n            train_loss = self.train_epoch()\n            \n            # Validate\n            val_loss = self.validate()\n            \n            # Update learning rate\n            self.scheduler.step()\n            current_lr = self.optimizer.param_groups[0]['lr']\n            \n            # Print epoch summary\n            print(f\"\\nEpoch {epoch+1}/{num_epochs}:\")\n            print(f\"  Train Loss: {train_loss:.4f}\")\n            print(f\"  Val Loss:   {val_loss:.4f}\")\n            print(f\"  LR:         {current_lr:.6f}\")\n            \n            # Check for improvement\n            is_best = val_loss < self.best_val_loss\n            if is_best:\n                self.best_val_loss = val_loss\n                epochs_without_improvement = 0\n            else:\n                epochs_without_improvement += 1\n            \n            # Save checkpoint\n            if (epoch + 1) % 5 == 0 or is_best:\n                self.save_checkpoint(is_best=is_best)\n            \n            # Early stopping\n            if epochs_without_improvement >= early_stopping_patience:\n                print(f\"\\n⚠️ Early stopping (no improvement for {early_stopping_patience} epochs)\")\n                break\n        \n        print(\"\\n✅ Training complete!\")\n        print(f\"   Best val loss: {self.best_val_loss:.4f}\")\n        \n        return self.train_losses, self.val_losses","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:55:24.462712Z","iopub.execute_input":"2026-01-26T22:55:24.463414Z","iopub.status.idle":"2026-01-26T22:55:24.486714Z","shell.execute_reply.started":"2026-01-26T22:55:24.463384Z","shell.execute_reply":"2026-01-26T22:55:24.485933Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# PART 7: COMPETITION SUBMISSION GENERATOR\n# ============================================================================\n\nclass CompetitionSubmission:\n    \"\"\"\n    Generate Kaggle competition submission\n    \n    CRITICAL: Competition format requirements:\n    - id: {record_id}_{sample_index}_{lead_name}\n    - value: waveform amplitude in mV\n    - Lead II: 5000 samples (10 seconds @ 1000Hz)\n    - Other leads: 2500 samples (2.5 seconds @ 1000Hz)\n    \"\"\"\n    \n    def __init__(self, model_path, device='cuda'):\n        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')\n        \n        # Load model\n        self.model = nnUNet2D(in_channels=3, num_classes=14)\n        checkpoint = torch.load(model_path, map_location=self.device)\n        self.model.load_state_dict(checkpoint['model_state_dict'])\n        self.model.to(self.device)\n        self.model.eval()\n        \n        # Preprocessor and extractor\n        self.preprocessor = GeometricPreprocessor()\n        self.extractor = DPWaveformExtractor()\n        \n        print(f\"✓ Model loaded from {model_path}\")\n    \n    def predict_waveforms(self, image_path):\n        \"\"\"Run complete pipeline on single image\"\"\"\n        # Step 1: Geometric preprocessing\n        result = self.preprocessor.process(image_path)\n        canonical = result['canonical']\n        geometric_mask = result['geometric_mask']\n        \n        # Step 2: nnU-Net prediction\n        # Resize to model input size\n        image_resized = cv2.resize(canonical, (512, 512))\n        image_tensor = torch.FloatTensor(image_resized.transpose(2, 0, 1) / 255.0)\n        image_tensor = image_tensor.unsqueeze(0).to(self.device)\n        \n        with torch.no_grad():\n            output = self.model(image_tensor)\n            predicted_mask = torch.argmax(output, dim=1)[0].cpu().numpy()\n        \n        # Resize mask back to canonical size\n        predicted_mask = cv2.resize(\n            predicted_mask.astype(np.uint8),\n            (canonical.shape[1], canonical.shape[0]),\n            interpolation=cv2.INTER_NEAREST\n        )\n        \n        # Step 3: Mask fusion\n        final_mask = fuse_masks(predicted_mask, geometric_mask)\n        \n        # Step 4: Waveform extraction\n        waveforms = self.extractor.extract_all_leads(canonical, final_mask)\n        \n        return waveforms\n    \n    def create_submission_csv(self, test_dir, output_path='submission.csv'):\n        \"\"\"\n        Generate competition submission CSV\n        \n        Format:\n        id,value\n        record1_0_I,0.123\n        record1_1_I,0.456\n        ...\n        record1_0_II,0.789\n        ...\n        \"\"\"\n        print(f\"\\n🔮 Generating submission for test images in {test_dir}\")\n        \n        test_dir = Path(test_dir)\n        test_images = sorted(list(test_dir.rglob('*.png')) + list(test_dir.rglob('*.jpg')))\n        \n        if not test_images:\n            print(f\"❌ No test images found in {test_dir}\")\n            return None\n        \n        print(f\"   Found {len(test_images)} test images\")\n        \n        # Lead names mapping\n        lead_names = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n        \n        # Prepare submission data\n        submission_rows = []\n        \n        for img_path in tqdm(test_images, desc=\"Processing test images\"):\n            # Extract record ID from filename\n            record_id = img_path.stem.split('-')[0]\n            \n            # Predict waveforms\n            try:\n                waveforms = self.predict_waveforms(img_path)\n            except Exception as e:\n                print(f\"\\n⚠️ Error processing {record_id}: {e}\")\n                # Create zero waveforms as fallback\n                waveforms = {}\n                for i in range(1, 13):\n                    if i == 2:  # Lead II\n                        waveforms[i] = np.zeros(10000)\n                    else:\n                        waveforms[i] = np.zeros(2500)\n            \n            # Format for submission\n            for lead_idx, lead_name in enumerate(lead_names, start=1):\n                signal = waveforms.get(lead_idx, np.zeros(2500))\n                \n                # CRITICAL: Ensure correct length\n                if lead_idx == 2:  # Lead II\n                    target_len = 10000\n                else:\n                    target_len = 2500\n                \n                # Pad or crop to exact length\n                if len(signal) < target_len:\n                    signal = np.pad(signal, (0, target_len - len(signal)), mode='edge')\n                elif len(signal) > target_len:\n                    signal = signal[:target_len]\n                \n                # Create rows: {record_id}_{sample_index}_{lead_name},value\n                for sample_idx, value in enumerate(signal):\n                    row_id = f\"{record_id}_{sample_idx}_{lead_name}\"\n                    submission_rows.append({\n                        'id': row_id,\n                        'value': float(value)\n                    })\n        \n        # Create DataFrame\n        submission_df = pd.DataFrame(submission_rows)\n        \n        # Save\n        output_path = Path(output_path)\n        submission_df.to_csv(output_path, index=False)\n        \n        print(f\"\\n✅ Submission created!\")\n        print(f\"   File: {output_path}\")\n        print(f\"   Rows: {len(submission_df):,}\")\n        \n        # Verify format\n        print(f\"\\n📋 Format verification:\")\n        print(f\"   Columns: {list(submission_df.columns)}\")\n        print(f\"   Sample rows:\")\n        print(submission_df.head(10))\n        \n        # Check lengths per lead\n        print(f\"\\n✓ Lead length verification:\")\n        for lead_name in lead_names:\n            lead_rows = submission_df[submission_df['id'].str.endswith(f'_{lead_name}')]\n            records = lead_rows['id'].str.split('_').str[0].unique()\n            samples_per_record = len(lead_rows) // len(records)\n            \n            expected = 10000 if lead_name == 'II' else 2500\n            status = \"✓\" if samples_per_record == expected else \"✗\"\n            print(f\"   {status} {lead_name}: {samples_per_record} samples (expected: {expected})\")\n        \n        return submission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:55:29.747725Z","iopub.execute_input":"2026-01-26T22:55:29.748035Z","iopub.status.idle":"2026-01-26T22:55:29.763643Z","shell.execute_reply.started":"2026-01-26T22:55:29.748008Z","shell.execute_reply":"2026-01-26T22:55:29.762991Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# PART 8: COMPLETE PIPELINE - READY FOR KAGGLE\n# ============================================================================\n\ndef setup_kaggle_data():\n    \"\"\"\n    Setup Kaggle competition data\n    \n    In Kaggle notebook, data is auto-mounted at:\n    /kaggle/input/physionet-ecg-image-digitization/\n    \n    Returns paths to train and test directories\n    \"\"\"\n    print(\"📁 Setting up Kaggle competition data...\")\n    \n    # Kaggle auto-mount path\n    kaggle_root = Path('/kaggle/input/physionet-ecg-image-digitization')\n    \n    if not kaggle_root.exists():\n        # Try alternative names\n        alternatives = [\n            Path('/kaggle/input/physionet-ecg-digitization'),\n            Path('/kaggle/input/ecg-digitization'),\n        ]\n        \n        for alt in alternatives:\n            if alt.exists():\n                kaggle_root = alt\n                break\n        else:\n            print(\"❌ Kaggle data not found!\")\n            print(\"Available datasets:\")\n            input_dir = Path('/kaggle/input')\n            if input_dir.exists():\n                for p in input_dir.iterdir():\n                    print(f\"  - {p}\")\n            return None, None\n    \n    train_dir = kaggle_root / 'train'\n    test_dir = kaggle_root / 'test'\n    \n    print(f\"✓ Competition data found:\")\n    print(f\"  Train: {train_dir} ({'✓' if train_dir.exists() else '✗'})\")\n    print(f\"  Test: {test_dir} ({'✓' if test_dir.exists() else '✗'})\")\n    \n    # Count files\n    if train_dir.exists():\n        train_images = list(train_dir.rglob('*.png'))\n        print(f\"  Train images: {len(train_images)}\")\n    \n    if test_dir.exists():\n        test_images = list(test_dir.rglob('*.png'))\n        print(f\"  Test images: {len(test_images)}\")\n    \n    return train_dir, test_dir\n\n\ndef run_complete_pipeline(train_dir, num_epochs=30):\n    \"\"\"\n    OPTIMIZED complete training pipeline\n    \n    CRITICAL CHANGES:\n    - Reduced epochs to 30 (was 50) - enough for convergence\n    - All optimizations enabled by default\n    - Multi-GPU support automatic\n    - Mixed precision training\n    \"\"\"\n    print(\"=\" * 80)\n    print(\"🚀 OPTIMIZED ECG DIGITIZATION PIPELINE\")\n    print(\"=\" * 80)\n    \n    # Step 1: Create OPTIMIZED dataloaders\n    print(\"\\n1️⃣ Creating OPTIMIZED dataloaders...\")\n    train_loader, val_loader = create_dataloaders(\n        image_dir=train_dir,\n        val_split=0.2,\n        batch_size=16,  # Optimized for T4\n        num_workers=4   # Optimized for multi-core\n    )\n    \n    if train_loader is None:\n        print(\"❌ Failed to create dataloaders\")\n        return\n    \n    # Calculate expected time\n    batches_per_epoch = len(train_loader) + len(val_loader)\n    estimated_time_per_epoch = batches_per_epoch * 0.5  # ~0.5s per batch with optimizations\n    total_estimated_time = estimated_time_per_epoch * num_epochs / 60  # minutes\n    \n    print(f\"\\n⏱️ Training time estimate:\")\n    print(f\"  Batches per epoch: {batches_per_epoch}\")\n    print(f\"  Estimated time per epoch: ~{estimated_time_per_epoch:.0f}s\")\n    print(f\"  Total estimated time: ~{total_estimated_time:.0f} minutes ({total_estimated_time/60:.1f} hours)\")\n    \n    # Step 2: Create OPTIMIZED model\n    print(\"\\n2️⃣ Creating OPTIMIZED nnU-Net model...\")\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    model = nnUNet2D(in_channels=3, num_classes=14)\n    \n    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    print(f\"   Parameters: {total_params:,} (reduced for speed)\")\n    print(f\"   Device: {device}\")\n    print(f\"   GPUs available: {torch.cuda.device_count()}\")\n    \n    # Step 3: Train with ALL optimizations\n    print(\"\\n3️⃣ Training model with optimizations...\")\n    print(\"   ✓ Mixed precision (FP16)\")\n    print(\"   ✓ Multi-GPU (if available)\")\n    print(\"   ✓ Cached preprocessing\")\n    print(\"   ✓ Larger batch size (16)\")\n    print(\"   ✓ Reduced image size (384x384)\")\n    \n    trainer = Trainer(\n        model, \n        train_loader, \n        val_loader, \n        device,\n        use_amp=True  # Enable mixed precision\n    )\n    \n    train_losses, val_losses = trainer.train(\n        num_epochs=num_epochs,\n        early_stopping_patience=15\n    )\n    \n    # Step 4: Plot training curves\n    print(\"\\n4️⃣ Saving training curves...\")\n    plt.figure(figsize=(10, 5))\n    plt.plot(train_losses, label='Train Loss', linewidth=2)\n    plt.plot(val_losses, label='Val Loss', linewidth=2)\n    plt.xlabel('Epoch', fontsize=12)\n    plt.ylabel('Loss', fontsize=12)\n    plt.title('Training History', fontsize=14, fontweight='bold')\n    plt.legend(fontsize=11)\n    plt.grid(True, alpha=0.3)\n    plt.tight_layout()\n    plt.savefig('training_history.png', dpi=150)\n    print(\"   Saved: training_history.png\")\n    \n    print(\"\\n\" + \"=\" * 80)\n    print(\"✅ TRAINING COMPLETE\")\n    print(\"=\" * 80)\n    print(f\"Best model saved to: checkpoints/best.pth\")\n    print(f\"\\n📊 Performance:\")\n    print(f\"  Best validation loss: {trainer.best_val_loss:.4f}\")\n    print(f\"  Training epochs: {len(train_losses)}\")\n    print(f\"  Time per epoch: ~{estimated_time_per_epoch:.0f}s\")\n    \n    return trainer\n\n\ndef generate_submission(test_dir, model_path='checkpoints/best.pth'):\n    \"\"\"\n    Generate competition submission from trained model\n    \"\"\"\n    print(\"=\" * 80)\n    print(\"🏆 GENERATING COMPETITION SUBMISSION\")\n    print(\"=\" * 80)\n    \n    if not Path(model_path).exists():\n        print(f\"❌ Model not found: {model_path}\")\n        print(\"Please train the model first!\")\n        return None\n    \n    # Create submission generator\n    submission = CompetitionSubmission(model_path)\n    \n    # Generate submission CSV\n    submission_df = submission.create_submission_csv(\n        test_dir=test_dir,\n        output_path='submission.csv'\n    )\n    \n    print(\"\\n\" + \"=\" * 80)\n    print(\"✅ SUBMISSION READY FOR KAGGLE\")\n    print(\"=\" * 80)\n    print(\"\\n💡 Next steps:\")\n    print(\"   1. Download submission.csv\")\n    print(\"   2. Go to competition page\")\n    print(\"   3. Click 'Submit Predictions'\")\n    print(\"   4. Upload submission.csv\")\n    print(\"   5. Check your score!\")\n    \n    return submission_df\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:55:33.560552Z","iopub.execute_input":"2026-01-26T22:55:33.561078Z","iopub.status.idle":"2026-01-26T22:55:33.575863Z","shell.execute_reply.started":"2026-01-26T22:55:33.561050Z","shell.execute_reply":"2026-01-26T22:55:33.575240Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================================\n# MAIN EXECUTION (For Kaggle Notebook)\n# ============================================================================\n\nif __name__ == \"__main__\":\n    print(\"=\" * 80)\n    print(\"COMPLETE ECG DIGITIZATION PIPELINE - KAGGLE READY\")\n    print(\"=\" * 80)\n    print(\"\\nPipeline stages:\")\n    print(\"  1. GEOMETRY → Normalize layout (rule-based)\")\n    print(\"  2. SEGMENTATION → Learn trace pixels (nnU-Net)\")\n    print(\"  3. FUSION → Combine geometry + predictions\")\n    print(\"  4. EXTRACTION → Convert to waveforms (DP)\")\n    print(\"\\n\" + \"=\" * 80)\n    print(\"\\n📖 USAGE IN KAGGLE NOTEBOOK:\")\n    print(\"\\n# Setup data\")\n    print(\"train_dir, test_dir = setup_kaggle_data()\")\n    print(\"\\n# Train model\")\n    print(\"trainer = run_complete_pipeline(train_dir, num_epochs=50)\")\n    print(\"\\n# Generate submission\")\n    print(\"submission_df = generate_submission(test_dir)\")\n    print(\"\\n\" + \"=\" * 80)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:55:38.252186Z","iopub.execute_input":"2026-01-26T22:55:38.252835Z","iopub.status.idle":"2026-01-26T22:55:38.257911Z","shell.execute_reply.started":"2026-01-26T22:55:38.252804Z","shell.execute_reply":"2026-01-26T22:55:38.257315Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Some chnages in submission class  ","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport scipy.signal\n\ndef resample_to_target(signal, target_len):\n    if len(signal) == target_len:\n        return signal\n    \n    # Resample waveform safely\n    return scipy.signal.resample(signal, target_len)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:58:38.742830Z","iopub.execute_input":"2026-01-26T22:58:38.743541Z","iopub.status.idle":"2026-01-26T22:58:38.747141Z","shell.execute_reply.started":"2026-01-26T22:58:38.743509Z","shell.execute_reply":"2026-01-26T22:58:38.746558Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\nclass TestMetaLoader:\n    def __init__(self, test_csv_path):\n        self.df = pd.read_csv(test_csv_path)\n\n    def get_target_length(self, record_id, lead):\n        row = self.df[(self.df.id == int(record_id)) & (self.df.lead == lead)]\n        return int(row.number_of_rows.values[0])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:58:42.146890Z","iopub.execute_input":"2026-01-26T22:58:42.147431Z","iopub.status.idle":"2026-01-26T22:58:42.151840Z","shell.execute_reply.started":"2026-01-26T22:58:42.147404Z","shell.execute_reply":"2026-01-26T22:58:42.151067Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def generate_kaggle_submission(model, preprocessor, extractor, test_dir, test_csv):\n    \n    meta = TestMetaLoader(test_csv)\n    test_images = sorted(Path(test_dir).glob(\"*.png\"))\n    \n    leads = ['I','II','III','aVR','aVL','aVF','V1','V2','V3','V4','V5','V6']\n    rows = []\n\n    for img_path in test_images:\n        record_id = int(img_path.stem)   # IMPORTANT Kaggle naming\n        \n        # === Your pipeline preprocessing ===\n        result = preprocessor.process(img_path)\n        canonical = result[\"canonical\"]\n        geom_mask = result[\"geometric_mask\"]\n\n        # Model prediction\n        x = torch.FloatTensor(cv2.resize(canonical,(512,512)).transpose(2,0,1)/255).unsqueeze(0).cuda()\n        with torch.no_grad():\n            pred = torch.argmax(model(x),1)[0].cpu().numpy()\n        pred = cv2.resize(pred,(canonical.shape[1],canonical.shape[0]),interpolation=cv2.INTER_NEAREST)\n\n        final_mask = fuse_masks(pred, geom_mask)\n        waveforms = extractor.extract_all_leads(canonical, final_mask)\n\n        # === Kaggle formatting ===\n        for lid, lname in enumerate(leads, 1):\n            target_len = meta.get_target_length(record_id, lname)\n            signal = waveforms.get(lid, np.zeros(target_len))\n            signal = resample_to_target(signal, target_len)\n\n            for i, v in enumerate(signal):\n                rows.append({\"id\": f\"{record_id}_{i}_{lname}\", \"value\": float(v)})\n\n    df = pd.DataFrame(rows)\n    df.to_csv(\"submission.csv\", index=False)\n    print(\"Saved submission.csv with rows:\", len(df))\n    return df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:58:45.303014Z","iopub.execute_input":"2026-01-26T22:58:45.303575Z","iopub.status.idle":"2026-01-26T22:58:45.310880Z","shell.execute_reply.started":"2026-01-26T22:58:45.303546Z","shell.execute_reply":"2026-01-26T22:58:45.310085Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## USE model weights ","metadata":{}},{"cell_type":"code","source":"# Setup data\ntrain_dir, test_dir = setup_kaggle_data()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:56:54.667473Z","iopub.execute_input":"2026-01-26T22:56:54.668211Z","iopub.status.idle":"2026-01-26T22:57:02.222305Z","shell.execute_reply.started":"2026-01-26T22:56:54.668182Z","shell.execute_reply":"2026-01-26T22:57:02.221684Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"MODEL_PATH = \"/kaggle/input/ecg-digitization-nnunet-best-pth/best_model.pth\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:55:59.886904Z","iopub.execute_input":"2026-01-26T22:55:59.887419Z","iopub.status.idle":"2026-01-26T22:55:59.890804Z","shell.execute_reply.started":"2026-01-26T22:55:59.887388Z","shell.execute_reply":"2026-01-26T22:55:59.890082Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = nnUNet2D()   \nckpt = torch.load(MODEL_PATH, map_location=\"cuda\")\n\n# Handle DataParallel training case\nif \"model_state_dict\" in ckpt:\n    model.load_state_dict(ckpt[\"model_state_dict\"])\nelse:\n    model.load_state_dict(ckpt)\n\nmodel = model.cuda()\nmodel.eval()\nprint(\"Model loaded\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:03:25.307904Z","iopub.execute_input":"2026-01-26T23:03:25.308463Z","iopub.status.idle":"2026-01-26T23:03:25.462892Z","shell.execute_reply.started":"2026-01-26T23:03:25.308431Z","shell.execute_reply":"2026-01-26T23:03:25.462113Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"preprocessor = GeometricPreprocessor()\nextractor = DPWaveformExtractor()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:03:31.707756Z","iopub.execute_input":"2026-01-26T23:03:31.708483Z","iopub.status.idle":"2026-01-26T23:03:31.711565Z","shell.execute_reply.started":"2026-01-26T23:03:31.708454Z","shell.execute_reply":"2026-01-26T23:03:31.711007Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df = generate_kaggle_submission(\n    model=model,\n    preprocessor=preprocessor,\n    extractor=extractor,\n    test_dir=\"/kaggle/input/physionet-ecg-image-digitization/test\",\n    test_csv=\"/kaggle/input/physionet-ecg-image-digitization/test.csv\"\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:03:34.659105Z","iopub.execute_input":"2026-01-26T23:03:34.659973Z","iopub.status.idle":"2026-01-26T23:03:37.050379Z","shell.execute_reply.started":"2026-01-26T23:03:34.659935Z","shell.execute_reply":"2026-01-26T23:03:37.049620Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# model = nnUNet2D()   # or your model class name\n\n# ckpt = torch.load(MODEL_PATH, map_location=\"cuda\")\n# model.load_state_dict(ckpt[\"model_state_dict\"])\n# model.eval()\n# print(\"Model loaded\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:02:35.683195Z","iopub.execute_input":"2026-01-26T23:02:35.684027Z","iopub.status.idle":"2026-01-26T23:02:35.687382Z","shell.execute_reply.started":"2026-01-26T23:02:35.683996Z","shell.execute_reply":"2026-01-26T23:02:35.686650Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Train model\n# trainer = run_complete_pipeline(train_dir, num_epochs=50)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:57:02.879229Z","iopub.execute_input":"2026-01-26T22:57:02.879911Z","iopub.status.idle":"2026-01-26T22:57:02.882909Z","shell.execute_reply.started":"2026-01-26T22:57:02.879884Z","shell.execute_reply":"2026-01-26T22:57:02.882167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Generate submission\n# submission_df = generate_submission(test_dir, model_path=MODEL_PATH)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:57:11.879292Z","iopub.execute_input":"2026-01-26T22:57:11.879603Z","iopub.status.idle":"2026-01-26T22:57:11.882781Z","shell.execute_reply.started":"2026-01-26T22:57:11.879577Z","shell.execute_reply":"2026-01-26T22:57:11.882194Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# submission_df.to_csv(\"/kaggle/working/submission.csv\", index=False)\n# !ls /kaggle/working","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:05:26.401655Z","iopub.execute_input":"2026-01-26T23:05:26.401981Z","iopub.status.idle":"2026-01-26T23:05:26.405840Z","shell.execute_reply.started":"2026-01-26T23:05:26.401949Z","shell.execute_reply":"2026-01-26T23:05:26.405008Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# test_df = pd.read_csv(\"/kaggle/input/physionet-ecg-image-digitization/test.csv\")\n# expected_rows = test_df.groupby(\"id\")[\"number_of_rows\"].sum().sum()\n# print(\"Expected rows:\", expected_rows)\n# print(\"Your rows:\", len(df))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:05:31.269308Z","iopub.execute_input":"2026-01-26T23:05:31.269622Z","iopub.status.idle":"2026-01-26T23:05:31.273173Z","shell.execute_reply.started":"2026-01-26T23:05:31.269594Z","shell.execute_reply":"2026-01-26T23:05:31.272511Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import pandas as pd\n\n# df = pd.read_csv(\"/kaggle/working/submission.csv\")\n# print(df.columns)\n# print(\"Rows:\", len(df))\n# print(\"NaNs:\", df.isna().sum())\n# print(df.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:05:54.558036Z","iopub.execute_input":"2026-01-26T23:05:54.558823Z","iopub.status.idle":"2026-01-26T23:05:54.562084Z","shell.execute_reply.started":"2026-01-26T23:05:54.558789Z","shell.execute_reply":"2026-01-26T23:05:54.561332Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# print(len(test_df[\"id\"].unique()))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:05:57.458102Z","iopub.execute_input":"2026-01-26T23:05:57.458423Z","iopub.status.idle":"2026-01-26T23:05:57.462001Z","shell.execute_reply.started":"2026-01-26T23:05:57.458396Z","shell.execute_reply":"2026-01-26T23:05:57.461435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import pandas as pd\n\n# df = pd.read_csv(\"/kaggle/working/submission.csv\")\n\n# # Check leads per record\n# print(df[\"id\"].str.split(\"_\").str[0].value_counts())\n# print(df[\"id\"].str.split(\"_\").str[2].value_counts())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:06:00.641778Z","iopub.execute_input":"2026-01-26T23:06:00.642455Z","iopub.status.idle":"2026-01-26T23:06:00.645471Z","shell.execute_reply.started":"2026-01-26T23:06:00.642418Z","shell.execute_reply":"2026-01-26T23:06:00.644784Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# df[\"id\"].str.split(\"_\").str[0].value_counts()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T23:06:06.798138Z","iopub.execute_input":"2026-01-26T23:06:06.798741Z","iopub.status.idle":"2026-01-26T23:06:06.801947Z","shell.execute_reply.started":"2026-01-26T23:06:06.798709Z","shell.execute_reply":"2026-01-26T23:06:06.801285Z"}},"outputs":[],"execution_count":null}]}