{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"},{"sourceId":623160,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":468838,"modelId":484689},{"sourceId":623170,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":468846,"modelId":484698}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/open-ecg-digitizer/pytorch/default/1')\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn.functional as F\nfrom scipy.signal import medfilt\nfrom tqdm import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\nfrom torchvision.io import read_image\nfrom src.model.unet import UNet\nfrom src.model.perspective_detector import PerspectiveDetector\nfrom src.model.cropper import Cropper\nfrom src.model.pixel_size_finder import PixelSizeFinder\nfrom src.model.signal_extractor import SignalExtractor\nfrom src.model.lead_identifier import LeadIdentifier\nimport yaml\n\n# ================================================================================\n# CONFIGURATION\n# ================================================================================\n\nLEADS = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nTRAIN_DIR = '/kaggle/input/physionet-ecg-image-digitization/train/'\nTEST_DIR = '/kaggle/input/physionet-ecg-image-digitization/test/'\nMODEL_DIR = '/kaggle/input/open-ecg-digitizer-weights/pytorch/default/1'\n\nprint(f\"🔧 Device: {DEVICE}\")\n\n# ================================================================================\n# LOAD PRE-TRAINED MODELS\n# ================================================================================\n\ndef load_unet_model(weights_path, num_in_channels=3, num_out_channels=4, \n                    dims=[32, 64, 128, 256, 320, 320, 320, 320], depth=2):\n    \"\"\"Load U-Net model with pre-trained weights\"\"\"\n    model = UNet(\n        num_in_channels=num_in_channels,\n        num_out_channels=num_out_channels,\n        dims=dims,\n        depth=depth,\n    )\n    \n    state_dict = torch.load(weights_path, map_location=DEVICE)\n    state_dict = {k.replace(\"_orig_mod.\", \"\"): v for k, v in state_dict.items()}\n    model.load_state_dict(state_dict)\n    model.eval().to(DEVICE)\n    \n    return model\n\nprint(\"🔧 Loading models...\")\nsegmentation_unet = load_unet_model(\n    weights_path=f\"{MODEL_DIR}/unet_weights_07072025.pt\",\n    num_in_channels=3,\n    num_out_channels=4,\n    dims=[32, 64, 128, 256, 320, 320, 320, 320],\n    depth=2\n)\n\nlead_identifier_unet = load_unet_model(\n    weights_path=f\"{MODEL_DIR}/lead_name_unet_weights_07072025.pt\",\n    num_in_channels=1,\n    num_out_channels=13,\n    dims=[32, 64, 128, 256, 256],\n    depth=2\n)\n\nlayouts_path = '/kaggle/input/open-ecg-digitizer/pytorch/default/1/src/config/lead_layouts_george-moody-2024.yml'\nlayouts = yaml.safe_load(open(layouts_path, 'r'))\nprint(\"✅ Models loaded successfully\")\n\n# ================================================================================\n# HELPER FUNCTIONS\n# ================================================================================\n\ndef add_noise_to_image(input_img, sigma=1.8, opacity=0.18):\n    \"\"\"Add synthetic noise - improves generalization\"\"\"\n    noise = torch.sigmoid(torch.randn_like(input_img) * sigma)\n    input_img = (1 - opacity) * input_img + opacity * noise\n    return input_img\n\ndef load_png_file(path):\n    \"\"\"Load and preprocess PNG image\"\"\"\n    img = read_image(path)\n    img = img.float() / 255.0\n    img = img.unsqueeze(0)\n    if img.shape[1] > 3:\n        img = img[:, :3, :, :]\n    return img\n\ndef resample_image(image, resample_size=3200):\n    \"\"\"Resample image to target size\"\"\"\n    height, width = image.shape[2], image.shape[3]\n    max_dim = max(height, width)\n    \n    if max_dim > resample_size:\n        scale = resample_size / max_dim\n        new_size = (int(height * scale), int(width * scale))\n        return F.interpolate(image, size=new_size, mode=\"bilinear\", \n                           align_corners=False, antialias=True)\n    return image\n\ndef _crop_y(image, signal_prob, grid_prob, text_prob):\n    \"\"\"Crop image vertically based on signal content\"\"\"\n    def get_bounds(tensor):\n        prob = torch.clamp(\n            tensor.squeeze().sum(dim=tensor.dim() - 3) - \n            tensor.squeeze().sum(dim=tensor.dim() - 3).mean(),\n            min=0,\n        )\n        non_zero = (prob > 0).nonzero(as_tuple=True)[0]\n        if non_zero.numel() == 0:\n            return 0, tensor.shape[2] - 1\n        return int(non_zero[0].item()), int(non_zero[-1].item())\n    \n    y1, y2 = get_bounds(signal_prob + grid_prob)\n    slices = (slice(None), slice(None), slice(y1, y2 + 1), slice(None))\n    \n    return (image[slices], signal_prob[slices], grid_prob[slices], text_prob[slices])\n\ndef _align_feature_maps(cropper, image, signal_prob, grid_prob, text_prob, source_points):\n    \"\"\"Apply perspective transformation to align image\"\"\"\n    aligned_signal = cropper.apply_perspective(signal_prob, source_points, fill_value=0)\n    aligned_image = cropper.apply_perspective(image, source_points, fill_value=0)\n    aligned_grid = cropper.apply_perspective(grid_prob, source_points, fill_value=0)\n    aligned_text = cropper.apply_perspective(text_prob, source_points, fill_value=0)\n    \n    return _crop_y(aligned_image, aligned_signal, aligned_grid, aligned_text)\n\ndef crop_and_align_image(image, probs):\n    \"\"\"Detect perspective and crop image\"\"\"\n    perspective_detector = PerspectiveDetector(num_thetas=200)\n    cropper = Cropper(percentiles=(0.03, 0.97), alpha=0.99)\n    \n    alignment_params = perspective_detector(probs[0, 0])\n    source_points = cropper(probs[0, 1], alignment_params)\n    \n    signal_prob = probs[:, [2]]\n    grid_prob = probs[:, [0]]\n    text_prob = probs[:, [1]]\n    \n    return _align_feature_maps(cropper, image, signal_prob, grid_prob, \n                              text_prob, source_points)\n\ndef extract_ecg_signals(aligned_signal, aligned_grid, aligned_text, target_num_samples):\n    \"\"\"Extract ECG signals from aligned feature maps\"\"\"\n    pixel_size_finder = PixelSizeFinder(\n        min_number_of_grid_lines=30,\n        max_number_of_grid_lines=100,\n        lower_grid_line_factor=0.1,\n    )\n    \n    signal_extractor = SignalExtractor()\n    \n    identifier = LeadIdentifier(\n        layouts=layouts,\n        unet=lead_identifier_unet,\n        device=DEVICE,\n        possibly_flipped=False,\n        target_num_samples=target_num_samples,\n        required_valid_samples=2,\n    )\n    \n    mm_per_pixel_x, mm_per_pixel_y = pixel_size_finder(aligned_grid)\n    avg_pixel_per_mm = (1 / mm_per_pixel_x + 1 / mm_per_pixel_y) / 2\n    \n    signals = signal_extractor(aligned_signal.squeeze())\n    signals = identifier(signals, aligned_text, avg_pixel_per_mm=avg_pixel_per_mm)\n    \n    return signals\n\ndef apply_gentle_baseline_correction(signal, fs=500):\n    \"\"\"GENTLE baseline correction - preserves QRS complexes\"\"\"\n    # Remove DC offset\n    signal = signal - np.median(signal)\n    \n    # Very gentle median filter (2.5 second window)\n    if len(signal) > 400:\n        window_size = min(int(fs * 2.5), len(signal))\n        if window_size % 2 == 0:\n            window_size += 1\n        window_size = max(3, window_size)\n        \n        try:\n            baseline = medfilt(signal, kernel_size=window_size)\n            signal = signal - baseline\n        except:\n            pass\n    \n    return signal\n\ndef apply_physics_constraints(signals_dict):\n    \"\"\"Apply Einthoven's and Goldberger's laws\"\"\"\n    # Einthoven's Law: I + III = II\n    if 'I' in signals_dict and 'III' in signals_dict and 'II' in signals_dict:\n        min_len = min(len(signals_dict['I']), len(signals_dict['III']), len(signals_dict['II']))\n        residual = (signals_dict['I'][:min_len] + signals_dict['III'][:min_len] - \n                   signals_dict['II'][:min_len])\n        correction = residual / 3\n        signals_dict['I'][:min_len] -= correction\n        signals_dict['III'][:min_len] -= correction\n        signals_dict['II'][:min_len] += correction\n    \n    # Goldberger's Law: aVR + aVL + aVF = 0\n    if 'aVR' in signals_dict and 'aVL' in signals_dict and 'aVF' in signals_dict:\n        min_len = min(len(signals_dict['aVR']), len(signals_dict['aVL']), len(signals_dict['aVF']))\n        residual = (signals_dict['aVR'][:min_len] + signals_dict['aVL'][:min_len] + \n                   signals_dict['aVF'][:min_len])\n        correction = residual / 3\n        signals_dict['aVR'][:min_len] -= correction\n        signals_dict['aVL'][:min_len] -= correction\n        signals_dict['aVF'][:min_len] -= correction\n    \n    return signals_dict\n\ndef get_lead_slice(lead_name, number_of_rows):\n    \"\"\"Get slice for extracting specific lead\"\"\"\n    if lead_name == \"II\":\n        return slice(0, number_of_rows)\n    elif lead_name in (\"I\", \"III\"):\n        return slice(0, number_of_rows)\n    elif lead_name in (\"aVR\", \"aVF\", \"aVL\"):\n        return slice(1 * number_of_rows, 2 * number_of_rows)\n    elif lead_name in (\"V1\", \"V2\", \"V3\"):\n        return slice(2 * number_of_rows, 3 * number_of_rows)\n    elif lead_name in (\"V4\", \"V5\", \"V6\"):\n        return slice(3 * number_of_rows, 4 * number_of_rows)\n    else:\n        return slice(0, number_of_rows)\n\n# ================================================================================\n# PROCESS TEST IMAGES\n# ================================================================================\n\nprint(\"\\n\" + \"=\" * 80)\nprint(\"PROCESSING TEST IMAGES\")\nprint(\"=\" * 80)\n\ntest = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/test.csv')\nsubmission_data = []\nold_id = None\nextracted_signals = None\n\nfor idx, row in tqdm(test.iterrows(), total=len(test), desc=\"Processing\"):\n    \n    if row['id'] != old_id:\n        old_id = row['id']\n        \n        img_path = f\"{TEST_DIR}{row['id']}.png\"\n        target_num_samples = row['fs'] * 10\n        fs = row['fs']\n        \n        try:\n            # Load and preprocess image\n            input_img = load_png_file(img_path)\n            input_img = add_noise_to_image(input_img, sigma=1.8, opacity=0.18)\n            input_img = resample_image(input_img, resample_size=3200)\n            \n            # Segmentation\n            with torch.no_grad():\n                logits = segmentation_unet(input_img.to(DEVICE))\n                output_probs = torch.softmax(logits, dim=1)\n            \n            # Alignment and cropping\n            aligned_image, aligned_signal, aligned_grid, aligned_text = \\\n                crop_and_align_image(input_img, output_probs)\n            \n            # Signal extraction\n            signals = extract_ecg_signals(\n                aligned_signal, aligned_grid, aligned_text, \n                target_num_samples=target_num_samples\n            )\n            \n            # CRITICAL FIX: Convert tensor to numpy BEFORE using .copy()\n            # This was the bug in File 1 causing all zeros\n            extracted_signals = signals[\"canonical_lines\"].cpu().numpy() * 1e-3  # Convert to mV\n            \n            # Post-processing\n            signals_dict = {}\n            for idx_lead, lead_name in enumerate(LEADS):\n                if idx_lead < len(extracted_signals):\n                    signal = extracted_signals[idx_lead].copy()  # Now this works!\n                    \n                    # Fix NaNs\n                    if np.any(np.isnan(signal)):\n                        nan_mask = np.isnan(signal)\n                        valid_indices = np.where(~nan_mask)[0]\n                        if len(valid_indices) > 0:\n                            signal[nan_mask] = np.interp(\n                                np.where(nan_mask)[0],\n                                valid_indices,\n                                signal[valid_indices]\n                            )\n                        else:\n                            signal = np.zeros_like(signal)\n                    \n                    # Gentle baseline correction\n                    signal = apply_gentle_baseline_correction(signal, fs)\n                    signals_dict[lead_name] = signal\n            \n            # Apply physics constraints\n            signals_dict = apply_physics_constraints(signals_dict)\n            extracted_signals = np.array([signals_dict[lead] for lead in LEADS])\n            \n        except Exception as e:\n            print(f\"⚠️  Error processing {row['id']}: {str(e)}\")\n            extracted_signals = None\n    \n    # Extract lead data\n    file_id = row['id']\n    lead_name = row['lead']\n    number_of_rows = row['number_of_rows']\n    \n    if extracted_signals is not None:\n        try:\n            lead_index = LEADS.index(lead_name)\n            lead_data = extracted_signals[lead_index]\n            lead_data = lead_data[get_lead_slice(lead_name, number_of_rows)]\n            \n            # Final NaN check\n            if np.any(np.isnan(lead_data)):\n                lead_data = np.nan_to_num(lead_data, nan=0.0)\n            \n        except Exception as e:\n            lead_data = np.zeros(number_of_rows)\n    else:\n        lead_data = np.zeros(number_of_rows)\n    \n    # Resample if needed\n    if len(lead_data) != number_of_rows:\n        lead_data = np.interp(\n            np.linspace(0, 1, number_of_rows),\n            np.linspace(0, 1, len(lead_data)),\n            lead_data\n        )\n    \n    # Add to submission\n    for t in range(number_of_rows):\n        submission_data.append({\n            'id': f\"{file_id}_{t}_{lead_name}\",\n            'value': float(lead_data[t])\n        })\n\n# ================================================================================\n# SAVE SUBMISSION\n# ================================================================================\n\nsubmission = pd.DataFrame(submission_data)\nsubmission.to_csv('submission.csv', index=False)\n\nprint(f\"\\n✅ Submission saved: {len(submission)} predictions\")\nprint(f\"   NaNs: {submission['value'].isna().sum()}\")\nprint(f\"   Zeros: {(submission['value'] == 0).sum()}\")\nprint(f\"   Value range: [{submission['value'].min():.3f}, {submission['value'].max():.3f}]\")\n\nprint(\"\\n\" + \"=\" * 80)\nprint(\"✅ OPTIMIZED SOLUTION COMPLETE!\")\nprint(\"Expected Score: 9-10 dB\")\nprint(\"=\" * 80)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}