{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import cv2\nimport numpy as np\nfrom skimage.morphology import skeletonize\nfrom scipy.signal import medfilt\nimport pandas as pd\nimport torch\nfrom torch import nn\nfrom torchvision import transforms\nfrom PIL import Image\n\n# Dummy U-Net model for demonstration\n# In a real-world scenario, you would define, train, and load a proper U-Net model here.\nclass UNet(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Simplified for demonstration. A full U-Net has multiple encoder/decoder blocks.\n        self.conv1 = nn.Conv2d(1, 64, 3, padding=1)\n        self.conv2 = nn.Conv2d(64, 1, 3, padding=1)\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.conv2(x)\n        return torch.sigmoid(x)\n\ndef load_unet_model(model_path, device):\n    \"\"\"\n    Loads a trained U-Net model from a file.\n    Args:\n        model_path (str): The path to the saved model state.\n        device (torch.device): The device to load the model onto.\n    Returns:\n        UNet: The loaded and prepped U-Net model.\n    \"\"\"\n    model = UNet()\n    model.load_state_dict(torch.load(model_path, map_location=device))\n    model.to(device)\n    model.eval()\n    return model\n\n# 1. Pre-process the ECG images\ndef preprocess_image(image_path):\n    \"\"\"\n    Performs all necessary preprocessing steps on an ECG image.\n    Args:\n        image_path (str): Path to the input image.\n    Returns:\n        np.ndarray: The preprocessed image ready for segmentation.\n    \"\"\"\n    # Load image in grayscale\n    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)\n    if img is None:\n        raise FileNotFoundError(f\"Image not found at {image_path}\")\n\n    # Invert colors if needed (assuming black signal on white paper)\n    img = cv2.bitwise_not(img)\n\n    # Correct skew and rotation (using Hough lines or other techniques)\n    # This is a complex task. For simplicity, assume manual alignment or a basic rotation.\n    # Here, we use a simple horizontal line detection approach.\n    edges = cv2.Canny(img, 50, 150, apertureSize=3)\n    lines = cv2.HoughLines(edges, 1, np.pi/180, 200)\n    angle = 0\n    if lines is not None:\n        for rho, theta in lines[0]:\n            angle = theta\n            break\n    rotation_angle = np.degrees(angle)\n    M = cv2.getRotationMatrix2D((img.shape[1]/2, img.shape[0]/2), rotation_angle - 90, 1)\n    img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]), borderMode=cv2.BORDER_CONSTANT, borderValue=255)\n\n    # Basic denoising using a median filter\n    img = medfilt(img, kernel_size=3)\n    \n    return img\n\n# 2. Segment the ECG waveform (using a trained U-Net model)\ndef segment_waveform(preprocessed_img, unet_model, device):\n    \"\"\"\n    Segments the ECG waveform from the preprocessed image using a U-Net model.\n    Args:\n        preprocessed_img (np.ndarray): The preprocessed input image.\n        unet_model (UNet): The trained U-Net model.\n        device (torch.device): The computation device.\n    Returns:\n        np.ndarray: A binary mask of the segmented waveform.\n    \"\"\"\n    # Resize and convert to tensor\n    transform = transforms.Compose([\n        transforms.Resize((256, 256)),  # U-Net typically works with fixed-size images\n        transforms.ToTensor(),\n    ])\n    \n    pil_img = Image.fromarray(preprocessed_img).convert('L')\n    tensor_img = transform(pil_img).unsqueeze(0).to(device)\n\n    # Get segmentation mask from U-Net\n    with torch.no_grad():\n        mask_tensor = unet_model(tensor_img)\n    \n    # Post-process mask: resize to original size and convert to binary\n    mask_np = mask_tensor.squeeze().cpu().numpy()\n    mask_resized = cv2.resize(mask_np, (preprocessed_img.shape[1], preprocessed_img.shape[0]))\n    \n    return (mask_resized > 0.5).astype(np.uint8) * 255\n\n# 3. Extract and refine the pixel coordinates\ndef extract_coordinates(mask):\n    \"\"\"\n    Extracts x, y pixel coordinates from the binary mask.\n    Args:\n        mask (np.ndarray): The binary mask of the ECG waveform.\n    Returns:\n        tuple: Arrays for x-coordinates (time) and y-coordinates (voltage).\n    \"\"\"\n    # Thin the waveform to a single pixel line using skeletonization\n    skeleton = skeletonize(mask > 0)\n    \n    # Extract coordinates\n    coords = np.argwhere(skeleton)\n    \n    # Sort coordinates by x-axis (time)\n    coords = coords[np.argsort(coords[:, 1])]\n    \n    # Extract median y-coordinate for each x-coordinate\n    time_pixels = []\n    voltage_pixels = []\n    \n    current_x = -1\n    y_vals = []\n    \n    for y, x in coords:\n        if x != current_x and y_vals:\n            time_pixels.append(current_x)\n            voltage_pixels.append(np.median(y_vals))\n            y_vals = []\n        current_x = x\n        y_vals.append(y)\n    \n    if y_vals:\n        time_pixels.append(current_x)\n        voltage_pixels.append(np.median(y_vals))\n\n    # Path refinement with a simple moving average filter\n    voltage_pixels_refined = medfilt(voltage_pixels, kernel_size=5)\n    \n    return np.array(time_pixels), np.array(voltage_pixels_refined)\n\n# 4. Calibrate pixels to time-series data\ndef calibrate_to_signal(image, time_pixels, voltage_pixels):\n    \"\"\"\n    Calibrates pixel coordinates to time and voltage values.\n    Args:\n        image (np.ndarray): The original image to detect grid.\n        time_pixels (np.ndarray): Array of extracted time pixel coordinates.\n        voltage_pixels (np.ndarray): Array of extracted voltage pixel coordinates.\n    Returns:\n        pd.DataFrame: A DataFrame with calibrated time and voltage data.\n    \"\"\"\n    # Detect grid scale\n    # This is a complex task and simplified here by hardcoding standard ECG paper properties.\n    # Advanced versions would use template matching or Hough lines to detect grid lines.\n    pixel_per_mm = 10 # Assuming 10 pixels per mm for this example\n\n    # Standard ECG paper properties\n    mm_per_mv = 10\n    mm_per_sec = 25\n\n    # Calculate scaling factors\n    pixel_per_mv = pixel_per_mm * mm_per_mv\n    pixel_per_sec = pixel_per_mm * mm_per_sec\n\n    # Invert voltage axis as image y-axis is inverted\n    max_y = image.shape[0]\n    voltage_calibrated = ((max_y - voltage_pixels) / pixel_per_mv) * 1 # Assuming 10 mm per mV\n\n    # Calculate time\n    time_calibrated = (time_pixels - time_pixels[0]) / pixel_per_sec\n\n    return pd.DataFrame({'time_s': time_calibrated, 'voltage_mV': voltage_calibrated})\n\n# Main execution block\ndef main(image_path, model_path, output_csv):\n    \"\"\"\n    Main function to run the entire digitization pipeline.\n    Args:\n        image_path (str): Path to the ECG image.\n        model_path (str): Path to the trained U-Net model.\n        output_csv (str): Path for the output CSV file.\n    \"\"\"\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # Step 1: Pre-process the ECG image\n    print(\"Preprocessing image...\")\n    preprocessed_img = preprocess_image(image_path)\n    \n    # Step 2: Segment the ECG waveform\n    print(\"Segmenting waveform with U-Net...\")\n    unet_model = load_unet_model(model_path, device)\n    mask = segment_waveform(preprocessed_img, unet_model, device)\n    \n    # Step 3: Extract and refine pixel coordinates\n    print(\"Extracting pixel coordinates...\")\n    time_pixels, voltage_pixels = extract_coordinates(mask)\n    \n    # Step 4: Calibrate and save the time-series data\n    print(\"Calibrating and saving data...\")\n    time_series_df = calibrate_to_signal(preprocessed_img, time_pixels, voltage_pixels)\n    time_series_df.to_csv(output_csv, index=False)\n    \n    print(f\"Digitization complete. Output saved to {output_csv}\")\n\nif __name__ == \"__main__\":\n    # Example usage:\n    # Set up dummy image and model file paths for demonstration.\n    # In a real project, replace with actual paths.\n    DUMMY_IMAGE_PATH = 'dummy_ecg_image.png'\n    DUMMY_MODEL_PATH = 'dummy_unet_model.pth'\n    OUTPUT_CSV_PATH = 'digitized_ecg.csv'\n\n    # Create dummy files for demonstration purposes\n    dummy_img = np.zeros((500, 1000), dtype=np.uint8)\n    cv2.line(dummy_img, (0, 250), (1000, 250), 255, 2)\n    cv2.imwrite(DUMMY_IMAGE_PATH, dummy_img)\n    torch.save(UNet().state_dict(), DUMMY_MODEL_PATH)\n    \n    main(DUMMY_IMAGE_PATH, DUMMY_MODEL_PATH, OUTPUT_CSV_PATH)\n    ","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}