{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"
\n

📊 ECG Image to Time-Series Digitization

\n \n

Introduction

\n \n\n

Key Features

\n \n\n

Reference Notebook – EDA

\n \n
\n","metadata":{}},{"cell_type":"code","source":"import os, gc\nimport cv2\nimport torch\nimport numpy as np\nimport pandas as pd\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.amp import autocast, GradScaler\nfrom scipy.signal import butter, filtfilt, resample\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm.notebook import tqdm","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Config\n\nclass Config:\n batch_size = 8\n epochs = 3\n lr = 1e-4\n num_workers = 2\n img_size = (1240, 1024)\n target_length = 5000 # For training lead II\n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nconfig = Config()\n\nLEADS = ['I','II','III','aVR','aVL','aVF','V1','V2','V3','V4','V5','V6']\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ECG Image Processor\n\nclass ECGImageProcessor:\n def __init__(self):\n # Approximate lead positions (relative)\n self.lead_positions = {\n 'I': (0.1, 0.15), 'II': (0.1, 0.3), 'III': (0.1, 0.45),\n 'aVR': (0.1, 0.6), 'aVL': (0.1, 0.75), 'aVF': (0.1, 0.9),\n 'V1': (0.55, 0.15), 'V2': (0.55, 0.3), 'V3': (0.55, 0.45),\n 'V4': (0.55, 0.6), 'V5': (0.55, 0.75), 'V6': (0.55, 0.9)\n }\n \n def preprocess_image(self, image):\n \"\"\"Convert to grayscale, enhance contrast, remove grid, binarize.\"\"\"\n gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) if len(image.shape)==3 else image\n clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))\n enhanced = clahe.apply(gray)\n kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2,2))\n opened = cv2.morphologyEx(enhanced, cv2.MORPH_OPEN, kernel)\n _, binary = cv2.threshold(opened, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)\n return binary\n \n def extract_lead_signal(self, image, lead_name, fs=500, target_length=None):\n \"\"\"Extract approximate lead signal from ROI\"\"\"\n h, w = image.shape[:2]\n x_ratio, y_ratio = self.lead_positions[lead_name]\n lead_x = int(w * x_ratio)\n lead_y = int(h * y_ratio)\n \n roi_width = int(w * 0.4)\n roi_height = int(h * 0.08)\n roi_x = max(0, lead_x - roi_width//2)\n roi_y = max(0, lead_y - roi_height//2)\n roi = image[roi_y:roi_y+roi_height, roi_x:roi_x+roi_width]\n if roi.size == 0:\n return np.zeros(target_length, dtype=np.float32)\n \n signal_y = []\n for col in range(roi.shape[1]):\n dark_pixels = np.where(roi[:,col] < 128)[0]\n signal_y.append(np.mean(dark_pixels) if len(dark_pixels)>0 else roi.shape[0]/2)\n if not signal_y:\n signal_y = [roi.shape[0]/2] * roi.shape[1]\n ecg_signal = roi_height - np.array(signal_y)\n ecg_signal = (ecg_signal - ecg_signal.mean()) / (ecg_signal.std()+1e-8)\n if target_length:\n ecg_signal = resample(ecg_signal, target_length)\n return ecg_signal.astype(np.float32)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Dataset\nclass ECGDataset(Dataset):\n def __init__(self, df, image_dir, transform=None, is_train=True):\n self.df = df\n self.image_dir = image_dir\n self.transform = transform\n self.is_train = is_train\n self.processor = ECGImageProcessor()\n \n def __len__(self):\n return len(self.df)\n \n def __getitem__(self, idx):\n row = self.df.iloc[idx]\n base_id = str(row['id'])\n # Try multiple images if missing\n img = None\n img_types = ['0001','0003','0004','0005'] if self.is_train else ['']\n for t in img_types:\n path = os.path.join(self.image_dir, base_id, f\"{base_id}-{t}.png\") if t else os.path.join(self.image_dir, f\"{base_id}.png\")\n if os.path.exists(path):\n img = cv2.imread(path)\n if img is not None:\n break\n if img is None:\n img = np.ones((config.img_size[1], config.img_size[0], 3), dtype=np.uint8)*255\n print(f\"Missing image for {base_id}\")\n \n img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n \n if self.transform:\n img_tensor = self.transform(image=img_rgb)['image']\n else:\n img_tensor = torch.from_numpy(cv2.resize(img_rgb, config.img_size).transpose(2,0,1)).float()/255.0\n \n if self.is_train:\n # Load CSV target\n csv_path = os.path.join(self.image_dir, base_id, f\"{base_id}.csv\")\n try:\n sig_df = pd.read_csv(csv_path)\n target_signal = sig_df['II'].values.astype(np.float32)\n if len(target_signal) > config.target_length:\n target_signal = target_signal[:config.target_length]\n else:\n target_signal = np.pad(target_signal, (0, config.target_length - len(target_signal)), 'constant')\n if target_signal.std()>0:\n target_signal = (target_signal - target_signal.mean()) / target_signal.std()\n except:\n target_signal = np.zeros(config.target_length, dtype=np.float32)\n # Extract input feature from image\n extracted_signal = self.processor.extract_lead_signal(self.processor.preprocess_image(img_rgb), 'II', target_length=config.target_length)\n return img_tensor, torch.FloatTensor(extracted_signal), torch.FloatTensor(target_signal), base_id\n else:\n return img_tensor, base_id\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Model\n\nclass ECGNet(nn.Module):\n def __init__(self):\n super().__init__()\n self.cnn = nn.Sequential(\n nn.Conv2d(3,64,3,padding=1), nn.BatchNorm2d(64), nn.ReLU(),\n nn.Conv2d(64,64,3,padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2),\n nn.Conv2d(64,128,3,padding=1), nn.BatchNorm2d(128), nn.ReLU(),\n nn.Conv2d(128,128,3,padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2),\n nn.Conv2d(128,256,3,padding=1), nn.BatchNorm2d(256), nn.ReLU(),\n nn.Conv2d(256,256,3,padding=1), nn.BatchNorm2d(256), nn.ReLU(),\n nn.AdaptiveAvgPool2d((4,4))\n )\n self.regressor = nn.Sequential(\n nn.Linear(256*4*4,1024), nn.ReLU(), nn.Dropout(0.5),\n nn.Linear(1024,512), nn.ReLU(), nn.Dropout(0.3),\n nn.Linear(512, config.target_length)\n )\n def forward(self,x):\n f = self.cnn(x)\n f = f.view(f.size(0),-1)\n return self.regressor(f)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Loss\n\nclass ECGLoss(nn.Module):\n def __init__(self, eps=1e-8):\n super().__init__()\n self.eps = eps\n def forward(self,pred,target):\n noise = target - pred\n signal_power = torch.sum(target**2, dim=1)\n noise_power = torch.sum(noise**2, dim=1)\n snr = signal_power/(noise_power+self.eps)\n # Smoothness regularizer\n diff2 = (pred[:,2:] - 2*pred[:,1:-1] + pred[:,:-2])**2\n smooth = torch.mean(diff2)\n return -torch.mean(torch.log(snr+self.eps)) + 0.1*smooth\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Training\ntrain_df = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/train.csv').head(500) # small for demo\ntrain_transform = A.Compose([\n A.Resize(*config.img_size),\n A.HorizontalFlip(p=0.3),\n A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=5, p=0.3),\n A.GridDistortion(p=0.1),\n A.GaussNoise(p=0.2),\n A.RandomBrightnessContrast(p=0.3),\n A.MotionBlur(p=0.1),\n A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),\n ToTensorV2()\n])\ntrain_dataset = ECGDataset(train_df, '/kaggle/input/physionet-ecg-image-digitization/train', transform=train_transform, is_train=True)\ntrain_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, num_workers=config.num_workers)\n\nmodel = ECGNet()\nif torch.cuda.device_count() > 1:\n model = nn.DataParallel(model)\nmodel.to(config.device)\n\ncriterion = ECGLoss()\noptimizer = optim.AdamW(model.parameters(), lr=config.lr, weight_decay=1e-4)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, factor=0.5)\n\nbest_loss = float('inf')\nfor epoch in range(config.epochs):\n model.train()\n running_loss = 0\n for images, extracted, targets, _ in tqdm(train_loader):\n images, targets = images.to(config.device), targets.to(config.device)\n optimizer.zero_grad()\n outputs = model(images)\n loss = criterion(outputs, targets)\n loss.backward()\n torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)\n optimizer.step()\n running_loss += loss.item()\n epoch_loss = running_loss/len(train_loader)\n scheduler.step(epoch_loss)\n print(f\"Epoch {epoch+1}/{config.epochs}, Loss={epoch_loss:.5f}\")\n if epoch_loss num_rows:\n pred = pred[:num_rows]\n elif len(pred) < num_rows:\n pred = np.pad(pred, (0,num_rows-len(pred)),'edge')\n pred = smooth_ecg(pred, fs=row['fs'])\n for i in range(num_rows):\n submission_data.append({'id':f\"{base_id}_{i}_{row['lead']}\", 'value':float(pred[i])})\n\nsubmission_df = pd.DataFrame(submission_data)\nsubmission_df.to_csv('submission.csv', index=False)\nsubmission_df.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"
\n

👍 Enjoyed this guide?

\n

If you found this guide helpful, please consider giving it an upvote! Your support helps us continue to create valuable content and improve our resources.

\n

Thank you! 😊

\n
\n","metadata":{}}]}