{"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":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\n# import os\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:26.296055Z","iopub.execute_input":"2025-11-04T14:03:26.296353Z","iopub.status.idle":"2025-11-04T14:03:26.609660Z","shell.execute_reply.started":"2025-11-04T14:03:26.296329Z","shell.execute_reply":"2025-11-04T14:03:26.608737Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport cv2\nimport os\nfrom tqdm import tqdm\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom scipy.signal import medfilt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:26.610997Z","iopub.execute_input":"2025-11-04T14:03:26.611406Z","iopub.status.idle":"2025-11-04T14:03:28.773907Z","shell.execute_reply.started":"2025-11-04T14:03:26.611381Z","shell.execute_reply":"2025-11-04T14:03:28.772966Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"LEADS=['I','II','III','aVR','aVL','aVF','V1','V2','V3','V4','V5','V6']\nDEVICE=torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nIMG_HEIGHT=256\nIMG_WIDTH=512\nBATCH_SIZE=8\nEPOCHS=5\nLR=1e-4\nLAMBDA_KG=0.1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:28.775005Z","iopub.execute_input":"2025-11-04T14:03:28.775438Z","iopub.status.idle":"2025-11-04T14:03:28.852111Z","shell.execute_reply.started":"2025-11-04T14:03:28.775401Z","shell.execute_reply":"2025-11-04T14:03:28.851331Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"KG_EDGES=[\n    ('I','II'),('I','III'),('II','III'),\n    ('aVR','aVL'),('aVR','aVF'),('aVL','aVF'),\n    ('I','aVL'),('II','aVF'),('III','aVF'),\n    ('V1','V2'),('V2','V3'),('V3','V4'),('V4','V5'),('V5','V6'),\n    ('V1','V3'),('V2','V4'),('V3','V5'),('V4','V6'),\n    ('I','V6'),('aVL','V5'),('aVL','V6'),\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:28.853522Z","iopub.execute_input":"2025-11-04T14:03:28.853819Z","iopub.status.idle":"2025-11-04T14:03:28.871971Z","shell.execute_reply.started":"2025-11-04T14:03:28.853799Z","shell.execute_reply":"2025-11-04T14:03:28.871403Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def build_kg_edge_index(leads_list):\n    lead_to_idx={lead:idx for idx,lead in enumerate(leads_list)}\n    edge_pairs=[]\n    for lead1,lead2 in KG_EDGES:\n        if lead1 in lead_to_idx and lead2 in lead_to_idx:\n            idx1=lead_to_idx[lead1]\n            idx2=lead_to_idx[lead2]\n            edge_pairs.append((idx1,idx2))\n            edge_pairs.append((idx2,idx1))\n    return torch.tensor(edge_pairs,dtype=torch.long)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:28.872670Z","iopub.execute_input":"2025-11-04T14:03:28.872916Z","iopub.status.idle":"2025-11-04T14:03:28.890666Z","shell.execute_reply.started":"2025-11-04T14:03:28.872891Z","shell.execute_reply":"2025-11-04T14:03:28.890096Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_adaptive_threshold(ima):\n    mean_brightness=np.mean(ima[:,:,2])\n    if mean_brightness>200:\n        return 170\n    elif mean_brightness<150:\n        return 160","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:28.891368Z","iopub.execute_input":"2025-11-04T14:03:28.891567Z","iopub.status.idle":"2025-11-04T14:03:28.914344Z","shell.execute_reply.started":"2025-11-04T14:03:28.891545Z","shell.execute_reply":"2025-11-04T14:03:28.913575Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def preprocess_ecg_image(image_path):\n    ima=cv2.imread(image_path)\n    if ima is None:\n        return None\n    crop_top=420 if ima.shape[0] == 1700 else 400\n    threshold=get_adaptive_threshold(ima) if get_adaptive_threshold(ima) is not None else 165\n    ima_crop=ima[crop_top:,:,2]>threshold\n    ima_crop=ima_crop.astype(np.float32)\n    ima_resized=cv2.resize(ima_crop,(IMG_WIDTH,IMG_HEIGHT))\n    return ima_resized","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:28.914894Z","iopub.execute_input":"2025-11-04T14:03:28.915074Z","iopub.status.idle":"2025-11-04T14:03:28.934671Z","shell.execute_reply.started":"2025-11-04T14:03:28.915060Z","shell.execute_reply":"2025-11-04T14:03:28.934120Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ECGDataset(Dataset):\n    def __init__(self,df,data_dir,is_train=True):\n        self.df=df;\n        self.data_dir=data_dir\n        self.is_train=is_train\n    def __len__(self):\n        return len(self.df)\n    def __getitem__(self, idx):\n        row=self.df.iloc[idx]\n        ecg_id=row['id']\n        img_path=os.path.join(self.data_dir, str(ecg_id), f\"{ecg_id}-0001.png\")\n        img=preprocess_ecg_image(img_path)\n        if img is None:\n            img=np.zeros((IMG_HEIGHT, IMG_WIDTH), dtype=np.float32)\n        img_tensor=torch.from_numpy(img).unsqueeze(0).float()\n        if self.is_train:\n            csv_path=os.path.join(self.data_dir, str(ecg_id), f\"{ecg_id}.csv\")\n            df_signal=pd.read_csv(csv_path)\n            signals=[]\n            for lead in LEADS:\n                if lead in df_signal.columns:\n                    signal=df_signal[lead].dropna().values.astype(np.float32)\n                    signal=(signal - signal.mean()) / (signal.std() + 1e-8)\n                    signal_resampled = np.interp(\n                        np.linspace(0, 1, 2500),\n                        np.linspace(0, 1, len(signal)),\n                        signal\n                    )\n                else:\n                    signal_resampled=np.zeros(2500, dtype=np.float32)\n                signals.append(signal_resampled)\n            target=torch.from_numpy(np.stack(signals)).float()\n            return img_tensor, target\n        return img_tensor, ecg_id        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:28.935348Z","iopub.execute_input":"2025-11-04T14:03:28.935592Z","iopub.status.idle":"2025-11-04T14:03:28.955538Z","shell.execute_reply.started":"2025-11-04T14:03:28.935568Z","shell.execute_reply":"2025-11-04T14:03:28.954964Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    def __init__(self,in_ch,out_ch):\n        super().__init__()\n        self.conv=nn.Sequential(\n            nn.Conv2d(in_ch,out_ch,3,padding=1),\n            nn.BatchNorm2d(out_ch),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_ch,out_ch,3,padding=1), \n            nn.BatchNorm2d(out_ch),\n            nn.ReLU(inplace=True)\n        )\n    def forward(self,x):\n        return self.conv(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:28.956317Z","iopub.execute_input":"2025-11-04T14:03:28.956528Z","iopub.status.idle":"2025-11-04T14:03:28.982595Z","shell.execute_reply.started":"2025-11-04T14:03:28.956512Z","shell.execute_reply":"2025-11-04T14:03:28.981925Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ECGUNetWithKG(nn.Module):\n    def __init__(self,n_leads=12,output_length=2500,embedding_dim=128):\n        super().__init__()\n        self.n_leads=n_leads\n        self.output_length=output_length\n        self.embedding_dim=embedding_dim\n        self.enc1=DoubleConv(1,64)\n        self.pool1=nn.MaxPool2d(2)\n        self.enc2=DoubleConv(64,128)\n        self.pool2=nn.MaxPool2d(2)\n        self.enc3=DoubleConv(128,256)\n        self.pool3=nn.MaxPool2d(2)\n        self.enc4=DoubleConv(256,512)\n        self.pool4=nn.MaxPool2d(4)\n        self.bottleneck=DoubleConv(512,1024)\n        self.up1=nn.ConvTranspose2d(1024,512,4,stride=4)\n        self.dec1=DoubleConv(1024,512)\n        self.up2=nn.ConvTranspose2d(512,256,2,stride=2)\n        self.dec2=DoubleConv(512,256)\n        self.up3=nn.ConvTranspose2d(256,128,2,stride=2)\n        self.dec3=DoubleConv(256,128)\n        self.up4=nn.ConvTranspose2d(128,64,2,stride=2)\n        self.dec4=DoubleConv(128,64)\n        self.global_pool=nn.AdaptiveAvgPool2d((1,output_length//10))\n        self.lead_heads=nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(64*(output_length//10),512),\n                nn.ReLU(),\n                nn.Dropout(0,3),\n                nn.Linear(512,output_length)\n            ) for _ in range(n_leads)\n        ])\n        self.embedding_layers=nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(output_length,256),\n                nn.ReLU(),\n                nn.Dropout(0.3),\n                nn.Linear(256,embedding_dim)\n            )for _ in range(n_leads)\n        ])\n    def forward(self,x,return_embeddings=False):\n        e1=self.enc1(x)\n        e2=self.enc2(self.pool1(e1))\n        e3=self.enc3(self.pool2(e2))\n        e4=self.enc4(self.pool3(e3))\n        b=self.bottleneck(self.pool4(e4))\n        d1=self.dec1(torch.cat([self.up1(b),e4],dim=1))\n        d2=self.dec2(torch.cat([self.up2(d1),e3],dim=1))\n        d3=self.dec3(torch.cat([self.up3(d2),e2],dim=1))\n        d4=self.dec4(torch.cat([self.up4(d3),e1],dim=1))\n        features=self.global_pool(d4)\n        features=features.view(features.size(0),-1)\n        outputs=[]\n        for head in self.lead_heads:\n            lead_output=head(features)\n            outputs.append(lead_output)\n        predictions=torch.stack(outputs,dim=1)\n        if return_embeddings:\n            embeddings=[]\n            for i,emb_layer in enumerate(self.embedding_layers):\n                lead_embedding=emb_layer(predictions[:,i,:])\n                embeddings.append(lead_embedding)\n            embeddings=torch.stack(embeddings,dim=1)\n            return predictions,embeddings\n        return predictions\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:28.984785Z","iopub.execute_input":"2025-11-04T14:03:28.985067Z","iopub.status.idle":"2025-11-04T14:03:29.002713Z","shell.execute_reply.started":"2025-11-04T14:03:28.985046Z","shell.execute_reply":"2025-11-04T14:03:29.002001Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def knowledge_graph_loss(embeddings,edge_index):\n    if edge_index.size()==0:\n        return torch.tensor(0.0,device=embeddings.device)\n    batch_size=embeddings.size(0)\n    total_loss=0.0\n    for b in range(batch_size):\n        batch_embeddings=embeddings[b]\n        source_embeddings=batch_embeddings[edge_index[:,0]]\n        target_embeddings=batch_embeddings[edge_index[:,1]]\n        distances=torch.norm(source_embeddings-target_embeddings,p=2,dim=1)\n        total_loss+=distances.mean()\n    return total_loss/batch_size\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:29.003425Z","iopub.execute_input":"2025-11-04T14:03:29.003615Z","iopub.status.idle":"2025-11-04T14:03:29.025808Z","shell.execute_reply.started":"2025-11-04T14:03:29.003601Z","shell.execute_reply":"2025-11-04T14:03:29.025009Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train(model,train_loader,val_loader,edge_index,epochs=EPOCHS,lambda_kg=LAMBDA_KG):\n    optimizer=torch.optim.Adam(model.parameters(),lr=LR)\n    reconstruction_criterion=nn.MSELoss()\n    scheduler=torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,patience=3)\n    edge_index=edge_index.to(DEVICE)\n    best_val_loss=float('inf')\n    train_kg_loss=0 \n    for epoch in range(epochs):\n        model.train()\n        train_recon_loss=0\n        train_total_loss=0\n        for batch in tqdm(train_loader,desc=f\"Epoch{epoch+1}/{epochs}\"):\n          imgs,targets=batch\n          imgs=imgs.to(DEVICE)\n          targets=targets.to(DEVICE)\n          optimizer.zero_grad()\n          predictions,embeddings=model(imgs,return_embeddings=True)\n          recon_loss=reconstruction_criterion(predictions,targets)\n          kg_loss=knowledge_graph_loss(embeddings,edge_index)\n          total_loss=recon_loss+lambda_kg*kg_loss\n          total_loss.backward()\n          optimizer.step()\n          train_recon_loss+=recon_loss.item()\n          train_kg_loss+=kg_loss.item()\n          train_total_loss+=total_loss.item()\n          train_recon_loss/=len(train_loader)\n          train_kg_loss/=len(train_loader)\n          train_total_loss/=len(train_loader)\n          model.eval()\n          val_recon_loss=0\n          val_kg_loss=0\n          val_total_loss=0\n          with torch.no_grad():\n            for batch in val_loader:\n                imgs,targets=batch\n                imgs=imgs.to(DEVICE)\n                targets=targets.to(DEVICE) \n                predictions,embeddings=model(imgs,return_embeddings=True)\n                recon_loss=reconstruction_criterion(predictions,targets)\n                kg_loss=knowledge_graph_loss(embeddings,edge_index)\n                total_loss=recon_loss+lambda_kg*kg_loss\n                val_recon_loss+=recon_loss.item()\n                val_kg_loss+=kg_loss.item()\n                val_total_loss+=total_loss.item()\n          val_recon_loss /= len(val_loader)\n          val_kg_loss /= len(val_loader)\n          val_total_loss /= len(val_loader)\n          scheduler.step(val_total_loss)\n          print(f'Epoch {epoch+1}:')\n          print(f'  Train - Recon: {train_recon_loss:.4f}, KG: {train_kg_loss:.4f}, Total: {train_total_loss:.4f}')\n          print(f'  Val   - Recon: {val_recon_loss:.4f}, KG: {val_kg_loss:.4f}, Total: {val_total_loss:.4f}')\n        \n          if val_total_loss < best_val_loss:\n            best_val_loss = val_total_loss\n            torch.save(model.state_dict(), 'best_unet_kg_model.pth')\n            print(f'Model saved (val_total_loss={val_total_loss:.4f})')\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:29.026452Z","iopub.execute_input":"2025-11-04T14:03:29.026625Z","iopub.status.idle":"2025-11-04T14:03:29.045563Z","shell.execute_reply.started":"2025-11-04T14:03:29.026611Z","shell.execute_reply":"2025-11-04T14:03:29.045011Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_and_save(model,test_df,output_path):\n    model.eval()\n    if os.path.exists(output_path):\n        os.remove(output_path)\n    pd.DataFrame(columns=[\"id\",\"value\"]).to_csv(output_path,index=False)\n    old_id=None\n    cached_signals=None\n    for index,row in tqdm(test_df.iterrows(),total=len(test_df),desc=\"Processing\"):\n        current_id=row[\"id\"]\n        if current_id!=old_id:\n            old_id=current_id\n            img_path = f'/kaggle/input/physionet-ecg-image-digitization/test/{current_id}.png'\n            img = preprocess_ecg_image(img_path)\n            if img is None:\n                img=np.zeros((IMG_HEIGHT,IMG_WIDTH),dtype=np.float32)\n            img_tensor=torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float().to(DEVICE)\n            with torch.no_grad():\n                outputs=model(img_tensor,return_embeddings=False)\n                cached_signals=outputs[0].cpu().numpy()\n        lead_name=row['lead']\n        number_of_rows=row[\"number_of_rows\"]\n        if lead_name not in LEADS:\n            print(\"unknown\")\n            continue\n        lead_idx=LEADS.index(lead_name)\n        signal=cached_signals[lead_idx]\n        signal_resampled=np.interp(\n            np.linspace(0,1,number_of_rows),\n            np.linspace(0,1,len(signal)),\n            signal\n        )\n        signal_resampled=medfilt(signal_resampled,kernel_size=5)\n        mean_val=np.nanmean(signal_resampled)\n        if np.isnan(mean_val):\n            mean_val=0.0\n        signal_resampled=np.nan_to_num(signal_resampled,nan=mean_val)\n        chunk=[]\n        for t in range(number_of_rows):\n            chunk.append({\n                \"id\": f\"{current_id}_{t}_{lead_name}\",\n                \"value\": float(signal_resampled[t])\n            })\n        if chunk:\n            pd.DataFrame(chunk).to_csv(output_path, mode='a', index=False, header=False)\n    print(f\"✓ Submission saved to {output_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:29.046145Z","iopub.execute_input":"2025-11-04T14:03:29.046349Z","iopub.status.idle":"2025-11-04T14:03:29.069298Z","shell.execute_reply.started":"2025-11-04T14:03:29.046334Z","shell.execute_reply":"2025-11-04T14:03:29.068725Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == '__main__':\n    train_df = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/train.csv')\n    test_df = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/test.csv')\n    \n    print(f\"Train DF shape: {train_df.shape}\")\n    print(f\"Test DF shape: {test_df.shape}\")\n    \n    kg_edge_index = build_kg_edge_index(LEADS)\n    print(f\"Knowledge Graph: {len(KG_EDGES)} unique edges, {kg_edge_index.size(0)} directed edges\")\n    split_idx = int(0.8 * len(train_df))\n    train_data = train_df[:split_idx]\n    val_data = train_df[split_idx:]\n    train_dataset = ECGDataset(train_data, '/kaggle/input/physionet-ecg-image-digitization/train', is_train=True)\n    val_dataset = ECGDataset(val_data, '/kaggle/input/physionet-ecg-image-digitization/train', is_train=True)\n    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\n    model = ECGUNetWithKG(n_leads=12, output_length=2500, embedding_dim=128).to(DEVICE)\n    print(f'Model parameters: {sum(p.numel() for p in model.parameters()):,}')    \n    train(model, train_loader, val_loader, kg_edge_index, epochs=EPOCHS, lambda_kg=LAMBDA_KG) \n    model.load_state_dict(torch.load('best_unet_kg_model.pth'))\n    predict_and_save(model, test_df, 'submission.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T14:03:29.069921Z","iopub.execute_input":"2025-11-04T14:03:29.070160Z"}},"outputs":[],"execution_count":null}]}