digitization.py

"""
digitization.py
===============
Core ECG image-to-signal conversion module.

The public interface is the :class:`ECGImage` class.  Instantiate it with
four pre-loaded YOLO models and the path to an ECG image, then call
:meth:`ECGImage.run_full_pipeline` followed by
:meth:`ECGImage.save_signals_as_csv` (or :meth:`ECGImage.save_signals_as_wfdb`).

Module-level helpers
--------------------
plot_image        — Quick matplotlib preview of a grayscale array.
shadow_removal    — Morphological background subtraction.
line_length       — Euclidean distance between two 2-D points.
parse_layout_from_folder — Parse rows/cols/layout flags from a folder name.
"""

import os
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"]      = "1"
os.environ["MKL_NUM_THREADS"]      = "1"

import numpy as np
import cv2
cv2.setNumThreads(0)   # disable OpenCV's internal thread pool

import matplotlib.pyplot as plt
from ultralytics import YOLO
from PIL import Image
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score, mean_squared_error
from scipy.interpolate import interp1d
import wfdb
from scipy import signal
from skimage import morphology, segmentation
from scipy.signal import savgol_filter, find_peaks
from scipy.stats import pearsonr
from skimage.filters import threshold_multiotsu
from concurrent.futures import ThreadPoolExecutor
import re
import pandas as pd
import torch
import time


from patched_yolo_infer import (
    MakeCropsDetectThem,
    CombineDetections,
    visualize_results,
)


def plot_image(img, title="Image Plot", size=(12, 12), show_axis=False):
    """Display a grayscale image with matplotlib.

    Parameters
    ----------
    img : np.ndarray
        Grayscale image array (H × W).
    title : str, optional
        Figure title (default ``"Image Plot"``).
    size : tuple[int, int], optional
        Figure size in inches ``(width, height)`` (default ``(12, 12)``).
    show_axis : bool, optional
        Whether to draw axis ticks (default ``False``).
    """
    plt.figure(figsize=size)
    plt.imshow(img, cmap='gray')
    plt.title(title)
    if not show_axis:
        plt.axis('off')
    plt.show()


def shadow_removal(img):
    """Remove uneven illumination / shadow from a grayscale image.

    Uses morphological dilation followed by median blur to estimate the
    background, then subtracts it from the original to yield a
    normalised, shadow-free image.

    Parameters
    ----------
    img : np.ndarray
        Grayscale uint8 image.

    Returns
    -------
    np.ndarray
        Shadow-corrected uint8 image, intensity range [0, 255].
    """
    dilated_img = cv2.dilate(img, np.ones((7, 7), np.uint8))
    bg_img = cv2.medianBlur(dilated_img, 15)
    diff_img = 255 - cv2.absdiff(img, bg_img)
    norm_img = cv2.normalize(diff_img, None, alpha=0, beta=255,
                             norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)
    return norm_img


def line_length(x1, y1, x2, y2):
    """Return the Euclidean length of a line segment.

    Parameters
    ----------
    x1, y1 : float
        Coordinates of the first endpoint.
    x2, y2 : float
        Coordinates of the second endpoint.

    Returns
    -------
    float
        Length of the segment in the same units as the inputs.
    """
    return np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)


def parse_layout_from_folder(folder_path):
    """Parse ECG layout metadata encoded in a folder name.

    Expected naming convention::

        ecg_signals_<rows>x<cols>_<Rythm|None>_<Cabrera|Normal>

    Parameters
    ----------
    folder_path : str
        Absolute or relative path whose final component encodes the layout.

    Returns
    -------
    layout_key : tuple[int, int, bool] or None
        ``(rows, cols, is_cabrera)``.  ``None`` if the name does not match.
    calibration : bool or None
        ``True`` when a rhythm (calibration) strip is present, ``None`` on
        parse failure.
    """
    base_name = folder_path.split('/')[-1]
    match = re.search(r'ecg_signals_(\d+)x(\d+)_(None|Rythm)_(Cabrera|Normal)', base_name)
    if match:
        rows = int(match.group(1))
        cols = int(match.group(2))
        calibration = (match.group(3) == 'Rythm')
        cabrera_flag = (match.group(4) == 'Cabrera')
        layout_key = (rows, cols, cabrera_flag)
        return layout_key, calibration
    else:
        return None, None


class ECGImage:
    """End-to-end digitizer for a single 12-lead ECG image.

    The class encapsulates every processing stage — from raw pixel
    loading to per-lead voltage time-series export.  All intermediate
    results (masks, bounding boxes, calibration constants, signal grid)
    are stored as instance attributes so that individual stages can be
    inspected or visualised after the pipeline runs.

    Typical usage
    -------------
    >>> ecg = ECGImage(box_model, seg_model, name_model, pulse_model, "ecg.png")
    >>> ecg.run_full_pipeline()
    >>> ecg.save_signals_as_csv("record_001", directory="./output")

    Parameters
    ----------
    box_model : ultralytics.YOLO
        Pre-loaded YOLO model that detects lead bounding boxes.
    segmentation_model : str
        Path to the YOLO segmentation model checkpoint.  Loaded
        internally by ``patched_yolo_infer.MakeCropsDetectThem`` so
        that it can operate on image patches.
    lead_name_model : ultralytics.YOLO
        Pre-loaded YOLO model that classifies lead name labels
        (I, II, III, aVR, aVL, aVF, V1–V6).
    pulse_model : ultralytics.YOLO
        Pre-loaded YOLO model that detects calibration pulse boxes.
    image_path : str
        Path to the ECG image (.png, .jpg, or .jpeg).
    wfdb_path : str, optional
        Path to a WFDB record (without extension) used for ground-truth
        comparison in :meth:`calculate_metrics_ptb`.  Not required for
        normal digitization.

    Attributes set by the pipeline
    --------------------------------
    image : np.ndarray
        Loaded and padded grayscale image.
    processed_image : np.ndarray
        Shadow-removed, blurred version used for all model inference.
    lead_segmentation : list
        ``CombineDetections`` objects from the three segmentation passes.
    mask_image : np.ndarray
        Binary mask (H × W, uint8) combining all segmented lead polygons.
    row_centers : np.ndarray
        Y-pixel positions of detected lead row centres.
    roi : tuple[float, float]
        ``(min_y, max_y)`` bounding the active ECG area.
    lead_bboxes : list[list[float]]
        ``[x1, y1, x2, y2]`` boxes from ``box_model``.
    lead_name_bboxes : list[dict]
        ``{'bbox': [...], 'class_name': str}`` from ``lead_name_model``.
    reference_pulses : list[dict]
        ``{'bbox': [...], 'image': np.ndarray}`` from ``pulse_model``.
    volt_per_pixel : float
        Calibration constant: millivolts per pixel (vertical axis).
    time_per_pixel : float
        Calibration constant: seconds per pixel (horizontal axis).
    is_cabrera : bool
        Whether the ECG uses Cabrera lead ordering.
    has_calibration_pulse : bool
        Whether a rhythm/calibration strip row is present.
    layout : list[list[str]]
        2-D grid of lead names matching the detected row/column layout.
    grid : list[list[dict]]
        Mask grid — each cell contains ``'lead'`` (str) and ``'signal'``
        (np.ndarray of the binary mask slice).
    signal_grid : list[list[dict]]
        After :meth:`extract_signals`: cells additionally contain
        ``'signal'`` as a list of float mV values.
    """

    def __init__(self, box_model, segmentation_model, lead_name_model, pulse_model,
                 image_path, wfdb_path=""):
        self.image_path = image_path          # stored for retry in run_full_pipeline
        self.load_image(image_path)
        self.wfdb_path = wfdb_path

        self.box_model = box_model
        self.segmentation_model = segmentation_model
        self.lead_name_model = lead_name_model
        self.pulse_model = pulse_model

        self.is_cabrera = None
        self.has_calibration_pulse = None

        self.standard_layouts = {
            # Standard layouts
            (12, 1, False): [['I'], ['II'], ['III'], ['aVR'], ['aVL'], ['aVF'],
                             ['V1'], ['V2'], ['V3'], ['V4'], ['V5'], ['V6']],
            (6, 2, False):  [['I', 'V1'], ['II', 'V2'], ['III', 'V3'],
                             ['aVR', 'V4'], ['aVL', 'V5'], ['aVF', 'V6']],
            (4, 3, False):  [['I', 'II', 'III'], ['aVR', 'aVL', 'aVF'],
                             ['V1', 'V2', 'V3'], ['V4', 'V5', 'V6']],
            (3, 4, False):  [['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'],
                             ['III', 'aVF', 'V3', 'V6']],
            # Cabrera layouts
            (12, 1, True):  [['aVL'], ['I'], ['aVR'], ['II'], ['aVF'], ['III'],
                             ['V1'], ['V2'], ['V3'], ['V4'], ['V5'], ['V6']],
            (6, 2, True):   [['aVL', 'V1'], ['I', 'V2'], ['aVR', 'V3'],
                             ['II', 'V4'], ['aVF', 'V5'], ['III', 'V6']],
            (4, 3, True):   [['aVL', 'I', 'aVR'], ['II', 'aVF', 'III'],
                             ['V1', 'V2', 'V3'], ['V4', 'V5', 'V6']],
            (3, 4, True):   [['aVL', 'II', 'V1', 'V4'], ['I', 'aVF', 'V2', 'V5'],
                             ['aVR', 'III', 'V3', 'V6']],
        }

    # ------------------------------------------------------------------
    # Image loading & preprocessing
    # ------------------------------------------------------------------

    def load_image(self, path, target_size=2100):
        """Load, resize, and pad an ECG image.

        The image is scaled so that its height equals *target_size* (cubic
        interpolation) and then padded with a 20-pixel white border on all
        sides.  The result is stored in ``self.image``.

        Parameters
        ----------
        path : str
            Path to the image file.
        target_size : int, optional
            Target image height in pixels (default ``2100``).

        Raises
        ------
        FileNotFoundError
            If OpenCV cannot open the file at *path*.
        """
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            raise FileNotFoundError(f"Cannot read image: {path}")
        resample_factor = target_size / img.shape[0]
        img = cv2.resize(img,
                         (int(img.shape[1] * resample_factor),
                          int(img.shape[0] * resample_factor)),
                         interpolation=cv2.INTER_CUBIC)
        self.image = cv2.copyMakeBorder(img, 20, 20, 20, 20,
                                        cv2.BORDER_CONSTANT, value=[255, 255, 255])

    def preprocess_image(self):
        """Apply shadow removal and Gaussian smoothing to ``self.image``.

        The result is stored in ``self.processed_image`` and is used as
        input for all subsequent YOLO inference calls.
        """
        rem = shadow_removal(self.image)
        self.processed_image = cv2.GaussianBlur(rem, (3, 3), 0)

    # ------------------------------------------------------------------
    # Segmentation
    # ------------------------------------------------------------------

    def segment_leads(self):
        """Run patched YOLO instance segmentation at three crop scales.

        The segmentation model is applied to overlapping image patches at
        crop factors of 4, 4.5, and 5 (relative to image height).
        Detections from each pass are combined with NMS
        (``nms_threshold=0.5``) and stored in ``self.lead_segmentation``
        as a list of three ``CombineDetections`` objects.
        """
        segmentations = []
        for shape in [4, 4.5, 5]:
            element_crops = MakeCropsDetectThem(
                image=cv2.cvtColor(self.processed_image, cv2.COLOR_GRAY2BGR),
                model_path=self.segmentation_model,
                segment=True,
                show_crops=False,
                shape_x=int(self.processed_image.shape[0] // shape),
                shape_y=int(self.processed_image.shape[0] // shape),
                overlap_x=50,
                overlap_y=50,
                conf=0.8,
                iou=0.7,
                classes_list=[0],
            )
            segmentations.append(CombineDetections(element_crops, nms_threshold=0.5))
        self.lead_segmentation = segmentations

    def make_segmentation_mask(self):
        """Rasterise all segmented polygons into a single binary mask.

        Polygons from all three segmentation passes in
        ``self.lead_segmentation`` are filled and merged into a combined
        mask.  A morphological opening (5 × 5 kernel) removes small
        spurious regions.  The result is stored in ``self.mask_image``
        (uint8, values 0 or 255).
        """
        height, width = self.image.shape[:2]
        combined_mask = np.zeros((height, width), dtype=np.uint8)

        for segmentation in self.lead_segmentation:
            for poly in segmentation.filtered_polygons:
                pts = np.array(poly, dtype=np.int32)
                if pts.ndim == 2:
                    pts = [pts]
                cv2.fillPoly(combined_mask, pts, color=255)

        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
        self.mask_image = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)

    # ------------------------------------------------------------------
    # Row detection
    # ------------------------------------------------------------------

    def find_row_centers(self):
        """Detect the vertical centre of each ECG lead row.

        Projects the binary mask onto the vertical axis (row sums) and
        applies ``scipy.signal.find_peaks`` twice: once to identify all
        peaks, and a second time with an adaptive minimum distance
        (⅔ × mean inter-peak spacing) to consolidate merged rows.

        Sets
        ----
        self.row_centers : np.ndarray
            Y-pixel indices of each row centre.
        self.first_peak_start : int
            Y-pixel where the topmost row's signal begins.
        self.last_peak_end : int
            Y-pixel where the bottommost row's signal ends.
        """
        image_height, image_width = self.mask_image.shape[:2]
        proj = np.sum(self.mask_image, 1)

        height   = (image_width // 10) * 255
        distance = image_height // 30

        peaks, _ = find_peaks(proj, height=height, distance=distance)
        row_centers, _ = find_peaks(
            proj, height=height,
            distance=int(np.mean(np.diff(peaks)) * (2 / 3))
        )
        self.row_centers = row_centers

        if len(peaks) > 0:
            first_peak  = peaks[0]
            start_index = first_peak
            zero_gap    = 0
            for i in range(first_peak - 1, -1, -1):
                if proj[i] == 0:
                    zero_gap += 1
                    if zero_gap > 2:
                        break
                else:
                    zero_gap    = 0
                    start_index = i
            self.first_peak_start = start_index

            last_peak = peaks[-1]
            end_index = last_peak
            zero_gap  = 0
            for i in range(last_peak + 1, image_height):
                if proj[i] == 0:
                    zero_gap += 1
                    if zero_gap > 2:
                        break
                else:
                    zero_gap  = 0
                    end_index = i
            self.last_peak_end = end_index

    def get_roi(self):
        """Compute the vertical region-of-interest (ROI) spanning all lead rows.

        Extends ⅔ of the mean row spacing above the first row centre and
        below the last.  Stores the result as ``self.roi = (min_y, max_y)``.
        """
        spacing = 2 / 3 * np.mean(np.diff(self.row_centers))
        min_y = max(0, self.row_centers[0] - spacing)
        max_y = min(self.image.shape[0], self.row_centers[-1] + spacing)
        self.roi = (min_y, max_y)

    # ------------------------------------------------------------------
    # YOLO inference helpers — GPU-safe wrappers
    # ------------------------------------------------------------------

    def _predict_safe(self, model, image_bgr, **kwargs):
        """Run a YOLO model prediction in a thread-safe manner.

        On CUDA: wraps the call in a dedicated ``torch.cuda.Stream`` and
        synchronises before returning, preventing contention on the
        default stream when multiple threads call this simultaneously.
        On CPU: wraps with ``torch.no_grad()`` only.

        Parameters
        ----------
        model : ultralytics.YOLO
            The detection/classification model to run.
        image_bgr : np.ndarray
            BGR image array as expected by Ultralytics.
        **kwargs
            Forwarded to ``model.predict()`` (e.g. ``conf``, ``iou``).

        Returns
        -------
        list
            Ultralytics ``Results`` objects.
        """
        if torch.cuda.is_available():
            stream = torch.cuda.Stream()
            with torch.no_grad():
                with torch.cuda.stream(stream):
                    results = model.predict(image_bgr, verbose=False, **kwargs)
                stream.synchronize()
        else:
            with torch.no_grad():
                results = model.predict(image_bgr, verbose=False, **kwargs)
        return results

    # ------------------------------------------------------------------
    # Timing helper
    # ------------------------------------------------------------------

    @staticmethod
    def _timed(label, fn, *args, **kwargs):
        """Call *fn* with *args*/*kwargs*, print its wall-clock duration, and return its result.

        Parameters
        ----------
        label : str
            Human-readable stage name printed alongside the elapsed time.
        fn : callable
            The stage function to time.
        *args, **kwargs
            Forwarded to *fn*.

        Returns
        -------
        Any
            Whatever *fn* returns.
        """
        t = time.time()
        result = fn(*args, **kwargs)
        print(f"    ⏱ {label}: {time.time() - t:.1f}s", flush=True)
        return result

    # ------------------------------------------------------------------
    # Box & name extraction
    # ------------------------------------------------------------------

    def extract_lead_boxes(self):
        """Detect lead bounding boxes with ``box_model`` and filter to the ROI.

        Runs ``_predict_safe`` on the preprocessed BGR image at confidence
        0.8, then discards boxes whose ``y1``/``y2`` coordinates fall
        outside ``self.roi``.  Sets ``self.lead_bboxes`` as a list of
        ``[x1, y1, x2, y2]`` float lists.
        """
        image_bgr = cv2.cvtColor(self.processed_image, cv2.COLOR_GRAY2BGR)
        results   = self._predict_safe(self.box_model, image_bgr, conf=0.8)

        min_y, max_y = self.roi
        lead_boxes   = []
        for r in results:
            for box in r.boxes:
                x1, y1, x2, y2 = box.xyxy.cpu().numpy()[0].tolist()
                if y1 >= min_y and y2 <= max_y:
                    lead_boxes.append([x1, y1, x2, y2])
        self.lead_bboxes = lead_boxes

    def extract_lead_name_boxes(self):
        """Detect and classify lead name labels with ``lead_name_model``.

        Runs at confidence 0.8, filters to the ROI, and stores results in
        ``self.lead_name_bboxes`` as a list of
        ``{'bbox': [x1, y1, x2, y2], 'class_name': str}`` dicts.
        """
        image_bgr = cv2.cvtColor(self.processed_image, cv2.COLOR_GRAY2BGR)
        results   = self._predict_safe(self.lead_name_model, image_bgr, conf=0.8)

        min_y, max_y = self.roi
        name_boxes   = []
        for r in results:
            for box in r.boxes:
                x1, y1, x2, y2 = box.xyxy.cpu().numpy()[0].tolist()
                if y1 >= min_y and y2 <= max_y:
                    cls_id   = int(box.cls.cpu().numpy()[0])
                    cls_name = self.lead_name_model.names[cls_id]
                    name_boxes.append({'bbox': [x1, y1, x2, y2], 'class_name': cls_name})
        self.lead_name_bboxes = name_boxes

    def extract_reference_pulses(self):
        """Detect calibration pulse boxes with ``pulse_model``.

        Runs at confidence 0.7 (lower than lead detection to improve
        recall on small pulses).  For each detection, stores the
        corresponding image crop alongside the bounding-box coordinates.
        Sets ``self.reference_pulses`` as a list of
        ``{'bbox': [x1, y1, x2, y2], 'image': np.ndarray}`` dicts.
        """
        image_bgr  = cv2.cvtColor(self.processed_image, cv2.COLOR_GRAY2BGR)
        results    = self._predict_safe(self.pulse_model, image_bgr, conf=0.7)

        pulse_boxes = []
        for r in results:
            for box in r.boxes:
                coord = box.xyxy.cpu().numpy()[0].tolist()
                pulse_boxes.append({
                    'image': self.image[int(coord[1]) - 5:int(coord[3]) + 5,
                                        int(coord[0]) - 5:int(coord[2]) + 5],
                    'bbox': coord,
                })
        self.reference_pulses = pulse_boxes

    # ------------------------------------------------------------------
    # Visualisation helpers
    # ------------------------------------------------------------------

    def visualize_boxes(self, task='Lead name', show_axis=False):
        """Overlay detection bounding boxes on the original image.

        Parameters
        ----------
        task : {'Lead name', 'Lead box', 'Reference pulse'}
            Which set of boxes to draw.
        show_axis : bool, optional
            Whether to display axis ticks (default ``False``).
        """
        img_copy = cv2.cvtColor(self.image.copy(), cv2.COLOR_GRAY2BGR)

        if task == 'Lead name':
            if not self.lead_name_bboxes:
                print("Lead name boxes not extracted.")
                return
            for box in self.lead_name_bboxes:
                x1, y1, x2, y2 = map(int, box['bbox'])
                cv2.rectangle(img_copy, (x1, y1), (x2, y2), (255, 0, 0), 2)
                cv2.putText(img_copy, box['class_name'], (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 3)

        elif task == 'Lead box':
            if not self.lead_bboxes:
                print("Lead boxes not extracted.")
                return
            for bbox in self.lead_bboxes:
                x1, y1, x2, y2 = map(int, bbox)
                cv2.rectangle(img_copy, (x1, y1), (x2, y2), (255, 0, 0), 2)

        elif task == 'Reference pulse':
            if not self.reference_pulses:
                print("Reference pulses not extracted.")
                return
            for bbox in self.reference_pulses:
                x1, y1, x2, y2 = map(int, bbox['bbox'])
                cv2.rectangle(img_copy, (x1, y1), (x2, y2), (255, 0, 0), 2)

        else:
            print(f"Unknown task: {task}")
            return

        plt.figure(figsize=(12, 10))
        plt.imshow(cv2.cvtColor(img_copy, cv2.COLOR_BGR2RGB))
        if not show_axis:
            plt.axis('off')
        plt.show()

    def visualize_segmentation(self, show_boxes=False, show_axis=False,
                               fill_mask=True, thickness=1):
        """Render all segmented lead polygons on the original image.

        Aggregates results from all three segmentation passes in
        ``self.lead_segmentation`` and passes them to
        ``patched_yolo_infer.visualize_results``.

        Parameters
        ----------
        show_boxes : bool, optional
            Also draw bounding boxes around each polygon (default ``False``).
        show_axis : bool, optional
            Display axis ticks (default ``False``).
        fill_mask : bool, optional
            Fill polygon interiors with a translucent colour (default ``True``).
        thickness : int, optional
            Polygon outline thickness in pixels (default ``1``).
        """
        all_confidences, all_boxes, all_polygons = [], [], []
        all_classes_ids, all_classes_names       = [], []

        for seg in self.lead_segmentation:
            all_confidences.extend(seg.filtered_confidences)
            all_boxes.extend(seg.filtered_boxes)
            all_polygons.extend(seg.filtered_polygons)
            all_classes_ids.extend(seg.filtered_classes_id)
            all_classes_names.extend(seg.filtered_classes_names)

        visualize_results(
            img=cv2.cvtColor(self.image, cv2.COLOR_GRAY2RGB),
            confidences=all_confidences,
            boxes=all_boxes,
            polygons=all_polygons,
            classes_ids=all_classes_ids,
            classes_names=all_classes_names,
            segment=True,
            thickness=thickness,
            fill_mask=fill_mask,
            show_boxes=show_boxes,
            show_class=False,
            axis_off=(not show_axis),
        )

    # ------------------------------------------------------------------
    # Calibration (volt/pixel & time/pixel)
    # ------------------------------------------------------------------

    def get_reference_scale(self):
        """Derive pixel-to-physical calibration constants from the calibration pulses.

        For each pulse in ``self.reference_pulses``, the method:

        1. Enhances contrast with CLAHE and binarises with multi-Otsu thresholding.
        2. Detects horizontal and vertical lines via morphological filtering and
           Hough transform to measure the pulse amplitude (voltage) and width (time).
        3. Refines the time calibration using the Line Segment Detector (LSD).

        Sets
        ----
        self.volt_per_pixel : float
            Millivolts per pixel (vertical axis).  Calibrated to the standard
            1 mV / 10 mm calibration pulse.
        self.time_per_pixel : float
            Seconds per pixel (horizontal axis).  Calibrated to 0.2 s between
            the two inner vertical edges of the calibration pulse.

        Raises
        ------
        RuntimeError
            If no valid pulse yields voltage or time measurements.
        """
        voltages, times, dist = [], [], []

        def _line_length(x1, y1, x2, y2):
            return np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)

        for pulse_idx, pulse in enumerate(self.reference_pulses):
            try:
                img = pulse['image']
                if img is None or img.size == 0:
                    continue
                h, w = img.shape
                if h < 10 or w < 10:
                    continue

                clahe    = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
                enhanced = clahe.apply(img)
                blurred  = cv2.GaussianBlur(enhanced, (3, 3), 0)

                thresholds = threshold_multiotsu(blurred, classes=2)
                regions    = np.digitize(blurred, bins=thresholds)
                binary     = (regions == 0).astype(np.uint8) * 255

                h_kernel    = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 1))
                v_kernel    = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 25))
                h_lines_img = cv2.morphologyEx(binary, cv2.MORPH_OPEN, h_kernel)
                v_lines_img = cv2.morphologyEx(binary, cv2.MORPH_OPEN, v_kernel)

                h_lines_raw = cv2.HoughLinesP(
                    h_lines_img, rho=1, theta=np.pi / 180,
                    threshold=w // 4, minLineLength=w // 2, maxLineGap=1
                )
                v_lines_raw = cv2.HoughLinesP(
                    v_lines_img, rho=1, theta=np.pi / 180,
                    threshold=h // 4, minLineLength=h // 2, maxLineGap=1
                )

                combined_lines = []
                if h_lines_raw is not None:
                    combined_lines.extend(h_lines_raw)
                if v_lines_raw is not None:
                    combined_lines.extend(v_lines_raw)

                horizontal_lines, vertical_lines = [], []
                for line in combined_lines:
                    x1, y1, x2, y2 = line[0]
                    angle = np.arctan2(abs(y2 - y1), abs(x2 - x1)) * 180 / np.pi
                    if angle < 10 or angle > 170:
                        horizontal_lines.append((x1, y1, x2, y2))
                    elif 80 < angle < 100:
                        vertical_lines.append((x1, y1, x2, y2))

                if len(vertical_lines) >= 2:
                    longest_v = sorted(vertical_lines,
                                       key=lambda l: _line_length(*l), reverse=True)[:2]
                    xc1 = (longest_v[0][0] + longest_v[0][2]) / 2
                    xc2 = (longest_v[1][0] + longest_v[1][2]) / 2
                    vertical_spacing = abs(xc1 - xc2)

                    volt_vals = sorted([_line_length(*l) for l in vertical_lines],
                                       reverse=True)[:2]
                    if volt_vals:
                        voltages.append(volt_vals)
                    if vertical_spacing > 5:
                        times.append(vertical_spacing)

                # LSD for more precise time calibration
                lsd = cv2.createLineSegmentDetector(refine=2)
                lsd_lines, _, _, _ = lsd.detect(v_lines_img)

                min_length    = h / 3
                angle_tol_rad = np.deg2rad(5)
                filtered      = []

                if lsd_lines is not None:
                    for line in lsd_lines:
                        x1, y1, x2, y2 = line[0]
                        length = np.hypot(x2 - x1, y2 - y1)
                        if length >= min_length:
                            angle = np.arctan2(y2 - y1, x2 - x1)
                            if np.abs(np.abs(angle) - np.pi / 2) <= angle_tol_rad:
                                filtered.append([[x1, y1, x2, y2]])

                if len(filtered) >= 4:
                    filtered.sort(
                        key=lambda l: np.hypot(l[0][2] - l[0][0], l[0][3] - l[0][1]),
                        reverse=True
                    )
                    top4 = filtered[:4]
                    top4.sort(key=lambda l: (l[0][0] + l[0][2]) / 2)

                    midpoints = []
                    for l1, l2 in [(top4[0], top4[1]), (top4[2], top4[3])]:
                        xc1 = (l1[0][0] + l1[0][2]) / 2
                        xc2 = (l2[0][0] + l2[0][2]) / 2
                        midpoints.append((xc1 + xc2) / 2)

                    dist.append(midpoints[1] - midpoints[0])

            except Exception as e:
                print(f"    ⚠ Skipping pulse {pulse_idx}: {e}", flush=True)
                continue

        if not voltages:
            raise RuntimeError("No valid reference pulses found for voltage calibration.")
        if not dist:
            raise RuntimeError("No valid reference pulses found for time calibration.")

        self.volt_per_pixel = 1 / np.mean(voltages)
        self.time_per_pixel = 0.2 / np.mean(dist)

    # ------------------------------------------------------------------
    # Layout helpers
    # ------------------------------------------------------------------

    def make_bounding_box_features(self, box, axis):
        """Extract three positional features from a bounding box along one axis.

        Parameters
        ----------
        box : list[float]
            Bounding box as ``[x1, y1, x2, y2]``.
        axis : {'x', 'y'}
            Axis along which to compute features.

        Returns
        -------
        list[float]
            ``[axis_min, axis_center, axis_max]``.
        """
        if axis == 'y':
            axis_min, axis_max = box[1], box[3]
        else:
            axis_min, axis_max = box[0], box[2]
        axis_center = (axis_min + axis_max) / 2
        return [axis_min, axis_center, axis_max]

    def bounding_boxes_kmeans(self, bounding_boxes, axis='y',
                              k_min=1, k_max=13, return_model=True):
        """Cluster bounding boxes along one axis using K-Means.

        Tries all values of *k* in ``[k_min, k_max]`` and selects the one
        that maximises the silhouette score.  Cluster labels are sorted by
        mean position along *axis* (ascending).

        Parameters
        ----------
        bounding_boxes : list[list[float]]
            List of ``[x1, y1, x2, y2]`` boxes.
        axis : {'x', 'y'}, optional
            Axis along which to cluster (default ``'y'``).
        k_min : int, optional
            Minimum number of clusters (default ``1``).
        k_max : int, optional
            Maximum number of clusters (default ``13``).
        return_model : bool, optional
            If ``True``, also return the fitted ``KMeans`` model and the
            label remapping dict (default ``True``).

        Returns
        -------
        sorted_labels : np.ndarray of int
            Cluster index (sorted) for each input box.
        best_k : int
            Optimal number of clusters.
        sorted_centers : np.ndarray
            Mean positional value of each cluster, sorted ascending.
        label_map : dict
            Mapping from original K-Means label → sorted label.
        best_model : sklearn.cluster.KMeans
            Fitted model (only when ``return_model=True``).

        Raises
        ------
        ValueError
            If *axis* is not ``'x'`` or ``'y'``, or if there are fewer
            boxes than *k_min*.
        """
        if axis not in ('x', 'y'):
            raise ValueError("Axis must be 'x' or 'y'")
        if len(bounding_boxes) < k_min:
            raise ValueError("Not enough bounding boxes to cluster")

        features = np.array([self.make_bounding_box_features(b, axis)
                             for b in bounding_boxes])

        best_score, best_k = -1, k_min
        best_labels = best_centers = best_model = None

        for k in range(k_min, min(k_max + 1, len(bounding_boxes))):
            kmeans = KMeans(n_clusters=k, random_state=42, n_init="auto")
            labels = kmeans.fit_predict(features)
            score  = silhouette_score(features, labels)
            if score > best_score:
                best_score   = score
                best_k       = k
                best_labels  = labels
                best_centers = kmeans.cluster_centers_
                best_model   = kmeans

        cluster_avgs   = best_centers.mean(axis=1)
        sorted_indices = np.argsort(cluster_avgs)
        label_map      = {old: new for new, old in enumerate(sorted_indices)}
        sorted_labels  = np.array([label_map[l] for l in best_labels])
        sorted_centers = cluster_avgs[sorted_indices]

        if return_model:
            return sorted_labels, best_k, sorted_centers, label_map, best_model
        return sorted_labels, best_k, sorted_centers, label_map

    def check_cabrera(self, num_rows, num_cols):
        """Determine whether the ECG uses Cabrera lead ordering.

        The heuristic varies by layout:

        * **12- or 6-row layouts:** compares the vertical spacing of
          augmented-limb leads (aVR/aVL/aVF) against precordial leads.
          Cabrera re-orders the limb leads so their spacing differs.
        * **4-row, 3-col / 5-row layouts:** checks the vertical spread
          of augmented leads (high std → interleaved → Cabrera).
        * **3-row / 4-col layouts:** checks the horizontal spread of
          augmented leads.

        Sets ``self.is_cabrera`` as a side-effect.

        Parameters
        ----------
        num_rows : int
            Number of detected lead rows.
        num_cols : int
            Number of detected lead columns.

        Returns
        -------
        bool
            ``True`` if Cabrera ordering is detected.
        """
        if num_rows in [13, 12, 7, 6]:
            av_leads = [b for b in self.lead_name_bboxes
                        if b['class_name'] in {'aVR', 'aVL', 'aVF'}]
            v_leads  = [b for b in self.lead_name_bboxes
                        if b['class_name'] in {'V1', 'V2', 'V3', 'V4', 'V5', 'V6'}]

            if not av_leads or not v_leads:
                return False

            y_v  = sorted([(b['bbox'][1] + b['bbox'][3]) / 2 for b in v_leads])
            y_av = sorted([(b['bbox'][1] + b['bbox'][3]) / 2 for b in av_leads])

            threshold = 30
            diff_v    = np.diff(y_v)
            diff_av   = np.diff(y_av)
            diff_v    = diff_v[np.abs(diff_v) > threshold]
            diff_av   = diff_av[np.abs(diff_av) > threshold]

            if len(diff_v) == 0 or len(diff_av) == 0:
                return False

            if abs(np.min(diff_v) - np.min(diff_av)) > (0.25 * np.min(diff_v)):
                self.is_cabrera = True
                return True
            self.is_cabrera = False
            return False

        elif (num_rows == 4 and num_cols == 3) or num_rows == 5:
            av_leads = [b for b in self.lead_name_bboxes
                        if b['class_name'] in {'aVR', 'aVL', 'aVF'}]
            if not av_leads:
                return False
            y_coords = [(b['bbox'][1] + b['bbox'][3]) / 2 for b in av_leads]
            self.is_cabrera = np.std(y_coords) > 25
            return self.is_cabrera

        elif (num_rows == 4 and num_cols == 4) or num_rows == 3:
            av_leads = [b for b in self.lead_name_bboxes
                        if b['class_name'] in {'aVR', 'aVL', 'aVF'}]
            if not av_leads:
                return False
            x_coords = [(b['bbox'][0] + b['bbox'][2]) / 2 for b in av_leads]
            self.is_cabrera = np.std(x_coords) > 25
            return self.is_cabrera

    def get_layout(self, num_rows):
        """Map the detected row count to a standard ECG layout.

        Calls :meth:`check_cabrera` to determine lead ordering, sets
        ``self.layout`` (2-D list of lead name strings) and
        ``self.has_calibration_pulse``, and returns the number of columns.

        Supported row counts
        --------------------
        * 13 → 12×1 with calibration pulse
        * 12 → 12×1 without calibration pulse
        * 7  → 6×2 with calibration pulse
        * 6  → 6×2 without calibration pulse
        * 5  → 4×3 with calibration pulse
        * 4  → 3×4 or 4×3 (disambiguated by V-lead spatial distribution)
        * 3  → 3×4 without calibration pulse

        Parameters
        ----------
        num_rows : int
            Number of detected lead rows.

        Returns
        -------
        int
            Number of lead columns in the detected layout.
        """
        if num_rows == 13:
            num_cols = 1
            self.has_calibration_pulse = True
            cabrera = self.check_cabrera(num_rows, num_cols)
            self.layout = self.standard_layouts[(12, 1, cabrera)]

        elif num_rows == 12:
            num_cols = 1
            self.has_calibration_pulse = False
            cabrera = self.check_cabrera(num_rows, num_cols)
            self.layout = self.standard_layouts[(12, 1, cabrera)]

        elif num_rows == 7:
            num_cols = 2
            self.has_calibration_pulse = True
            cabrera = self.check_cabrera(num_rows, num_cols)
            self.layout = self.standard_layouts[(6, 2, cabrera)]

        elif num_rows == 6:
            num_cols = 2
            self.has_calibration_pulse = False
            cabrera = self.check_cabrera(num_rows, num_cols)
            self.layout = self.standard_layouts[(6, 2, cabrera)]

        elif num_rows == 5:
            num_cols = 3
            self.has_calibration_pulse = True
            cabrera = self.check_cabrera(num_rows, num_cols)
            self.layout = self.standard_layouts[(4, 3, cabrera)]

        elif num_rows == 4:
            v_leads1 = [b for b in self.lead_name_bboxes
                        if b['class_name'] in {'V1', 'V2', 'V3'}]
            v_leads2 = [b for b in self.lead_name_bboxes
                        if b['class_name'] in {'V4', 'V5', 'V6'}]

            def centers(boxes, axis):
                if axis == 'x':
                    return [(b['bbox'][0] + b['bbox'][2]) / 2 for b in boxes]
                return [(b['bbox'][1] + b['bbox'][3]) / 2 for b in boxes]

            x_std1 = np.std(centers(v_leads1, 'x'))
            y_std1 = np.std(centers(v_leads1, 'y'))
            x_std2 = np.std(centers(v_leads2, 'x'))
            y_std2 = np.std(centers(v_leads2, 'y'))

            if x_std1 < y_std1 and x_std2 < y_std2:
                num_cols = 4
                self.has_calibration_pulse = True
            elif y_std1 < x_std1 and y_std2 < x_std2:
                num_cols = 3
                self.has_calibration_pulse = False
            else:
                num_cols = 4
                self.has_calibration_pulse = True

            cabrera = self.check_cabrera(num_rows, num_cols)
            key = (3, 4, cabrera) if num_cols == 4 else (4, 3, cabrera)
            self.layout = self.standard_layouts[key]

        elif num_rows == 3:
            num_cols = 4
            self.has_calibration_pulse = False
            cabrera = self.check_cabrera(num_rows, num_cols)
            self.layout = self.standard_layouts[(3, 4, cabrera)]

        return num_cols

    def clean_lead_mask(self, img, baseline_y):
        """Remove leaked segments of neighboring leads from a single lead cell binary mask.

        Parameters
        ----------
        img : np.ndarray
            Binary mask for a single lead cell (H × W).
        baseline_y : float
            Y-coordinate of the row baseline relative to the cropped row mask.

        Returns
        -------
        np.ndarray
            Cleaned binary mask.
        """
        height, width = img.shape
        if height < 10 or width < 10:
            return img

        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img)
        if num_labels <= 2:
            return img

        # Find the primary component (largest horizontal span)
        primary_c = -1
        max_w = -1
        for c in range(1, num_labels):
            w = stats[c, cv2.CC_STAT_WIDTH]
            if w > max_w:
                max_w = w
                primary_c = c

        if primary_c == -1:
            return img

        primary_area = stats[primary_c, cv2.CC_STAT_AREA]
        cleaned_img = img.copy()

        for c in range(1, num_labels):
            if c == primary_c:
                continue

            c_top = stats[c, cv2.CC_STAT_TOP]
            c_height = stats[c, cv2.CC_STAT_HEIGHT]
            c_bottom = c_top + c_height
            c_width = stats[c, cv2.CC_STAT_WIDTH]
            c_area = stats[c, cv2.CC_STAT_AREA]
            c_centroid_y = centroids[c][1]

            # Heuristics to identify a leak:
            # 1. It is relatively small (either narrow or small area)
            is_small = (c_width < 0.35 * width) or (c_area < 0.25 * primary_area)

            # 2. It is located near the vertical boundaries of the cell
            near_boundary = (c_top < 0.15 * height) or (c_bottom > 0.85 * height)

            # 3. It doesn't cross the baseline (with safety margin)
            baseline_margin = max(5, int(0.08 * height))
            crosses_baseline = (c_top - baseline_margin <= baseline_y <= c_bottom + baseline_margin)

            # 4. It is far from the baseline
            far_from_baseline = abs(c_centroid_y - baseline_y) > 0.2 * height

            if is_small and near_boundary and not crosses_baseline and far_from_baseline:
                cleaned_img[labels == c] = 0

        return cleaned_img

    # ------------------------------------------------------------------
    # Grid construction
    # ------------------------------------------------------------------

    def make_grid(self, padding=0):
        """Build the lead mask grid from segmentation polygons and lead boxes.

        Steps:

        1. Assigns each segmented polygon to the nearest row centre by
           polygon centroid Y-coordinate.
        2. Crops each row mask to its vertical extent.
        3. Clusters lead bounding boxes into *num_cols* columns with K-Means.
        4. Slices each (row, col) cell from the corresponding row mask.
        5. Applies a morphological opening to each cell slice.

        Sets
        ----
        self.grid : list[list[dict]]
            ``grid[row][col]`` contains ``{'lead': str, 'signal': np.ndarray}``,
            where ``'signal'`` is currently a binary mask slice.
        self.baseline : list[float]
            Relative Y-position of the isoelectric baseline within each row
            (pixels from the top of the cropped row mask).
        self.row : dict[int, np.ndarray]
            Cropped binary mask for each row index.

        Parameters
        ----------
        padding : int, optional
            Extra pixel padding (currently unused, reserved for future use).
        """
        image_height, image_width = self.image.shape[:2]
        num_rows = len(self.row_centers)
        num_cols = self.get_layout(num_rows)

        row_masks    = {i: np.zeros((image_height, image_width), dtype=np.uint8)
                        for i in range(num_rows)}
        row_polygons = {i: [] for i in range(num_rows)}
        row_limits   = {i: [] for i in range(num_rows)}

        for seg in self.lead_segmentation:
            for box, polygon in zip(seg.filtered_boxes, seg.filtered_polygons):
                poly_np   = np.array(polygon, dtype=np.int32).reshape((-1, 1, 2))
                temp_mask = np.zeros((image_height, image_width), dtype=np.uint8)
                cv2.fillPoly(temp_mask, [poly_np], color=1)

                temp_proj      = np.sum(temp_mask, axis=1)
                total_weight_y = np.sum(temp_proj)
                if total_weight_y == 0:
                    continue

                centroid_y = int(np.sum(np.arange(temp_proj.shape[0]) * temp_proj)
                                 / total_weight_y)

                y_vals      = [pt[1] for pt in polygon]
                min_y, max_y = min(y_vals), max(y_vals)

                if max_y < self.first_peak_start or min_y > self.last_peak_end:
                    continue

                closest_idx = int(np.argmin(np.abs(self.row_centers - centroid_y)))

                if self.has_calibration_pulse and closest_idx == num_rows - 1:
                    continue

                row_polygons[closest_idx].append(polygon)
                cv2.fillPoly(row_masks[closest_idx], [poly_np], color=1)

                if not row_limits[closest_idx]:
                    row_limits[closest_idx] = [min_y, max_y]
                else:
                    row_limits[closest_idx][0] = min(row_limits[closest_idx][0], min_y)
                    row_limits[closest_idx][1] = max(row_limits[closest_idx][1], max_y)

        cropped_row_masks = {}
        for i in range(num_rows):
            if self.has_calibration_pulse and i == num_rows - 1:
                continue
            if not row_limits[i]:
                continue
            min_y, max_y = row_limits[i]
            cropped_row_masks[i] = row_masks[i][min_y:max_y + 1, :]

        lead_boxes = [
            box for box in self.lead_bboxes
            if not (
                self.has_calibration_pulse and
                int(np.argmin(np.abs(self.row_centers - (box[1] + box[3]) / 2)))
                == num_rows - 1
            )
        ] if self.has_calibration_pulse else self.lead_bboxes

        self.row = cropped_row_masks

        if num_cols != 1:
            labels_cols, _, _, _ = self.bounding_boxes_kmeans(
                lead_boxes, axis='x', k_min=num_cols, k_max=num_cols, return_model=False
            )
        else:
            labels_cols = np.zeros(len(lead_boxes), dtype=int)

        boxes_arr = np.array(lead_boxes)
        min_x_per_col, max_x_per_col = [], []
        for col_label in range(num_cols):
            col_boxes = boxes_arr[labels_cols == col_label]
            min_x_per_col.append(col_boxes[:, 0].min())
            max_x_per_col.append(col_boxes[:, 2].max())

        mask_grid, relative_baselines = [], []
        for row_idx, row_slice in cropped_row_masks.items():
            row_cells = []
            for col_idx in range(num_cols):
                x_min      = max(0, int(min_x_per_col[col_idx]))
                x_max      = min(image_width, int(max_x_per_col[col_idx]))
                cell_slice = row_slice[:, x_min:x_max]

                kernel     = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
                cell_slice = cv2.morphologyEx(cell_slice, cv2.MORPH_OPEN, kernel)

                # Post-processing: Clean up leaked segments from neighboring leads
                baseline_y = self.row_centers[row_idx] - row_limits[row_idx][0]
                cell_slice = self.clean_lead_mask(cell_slice, baseline_y)

                row_cells.append({
                    'lead':   self.layout[row_idx][col_idx],
                    'signal': cell_slice,
                })

            relative_baselines.append(
                self.row_centers[row_idx] - row_limits[row_idx][0]
            )
            mask_grid.append(row_cells)

        self.grid     = mask_grid
        self.baseline = relative_baselines

    def visualize_grid(self, figsize=(15, 10)):
        """Display the lead mask grid — one subplot per (row, col) cell.

        Parameters
        ----------
        figsize : tuple[int, int], optional
            Overall figure size in inches (default ``(15, 10)``).

        Raises
        ------
        ValueError
            If :meth:`make_grid` has not been called yet.
        """
        if not hasattr(self, 'grid') or not self.grid:
            raise ValueError("Grid not generated. Call make_grid() first.")

        num_rows = len(self.grid)
        num_cols = len(self.grid[0]) if isinstance(self.grid[0], list) else 1

        fig, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
        if num_cols == 1:
            axes = np.atleast_2d(axes).T
        elif num_rows == 1:
            axes = np.atleast_2d(axes)

        for row_idx in range(num_rows):
            row = self.grid[row_idx] if isinstance(self.grid[row_idx], list) \
                  else [self.grid[row_idx]]
            for col_idx, cell in enumerate(row):
                ax = axes[row_idx][col_idx]
                ax.imshow(cell['signal'], cmap='gray', aspect='auto')
                ax.set_title(cell['lead'], fontsize=10)
                ax.axis('off')

        plt.tight_layout()
        plt.show()

    # ------------------------------------------------------------------
    # Signal extraction
    # ------------------------------------------------------------------

    def binarize_signal(self, img, window_length=11, polyorder=2):
        """Extract a 1-D signal trace from a binary lead mask column by column.

        For each image column, computes the intensity-weighted centroid of
        non-zero pixels.  Uses the second derivative of the initial centroid
        trace (Savitzky-Golay filtered) to disambiguate ambiguous columns:
        upward concavity → take the topmost pixels; downward concavity →
        take the bottommost pixels; otherwise use all non-zero pixels.

        Parameters
        ----------
        img : np.ndarray
            Binary uint8 mask for a single lead cell (H × W).
        window_length : int, optional
            Savitzky-Golay filter window (must be odd, default ``11``).
        polyorder : int, optional
            Savitzky-Golay polynomial order (default ``2``).

        Returns
        -------
        x_coords : np.ndarray of int
            Column indices where a signal was detected.
        final_signal : np.ndarray of float
            Weighted-centroid Y-coordinate for each column in *x_coords*.
        """
        height, width = img.shape
        x_coords = []
        initial_signal = []

        for col in range(width):
            column = img[:, col]
            if np.sum(column) == 0:
                continue

            y_indices = np.arange(height)
            weights = column.astype(float)
            centroid = np.average(y_indices, weights=weights)
            initial_signal.append(centroid)
            x_coords.append(col)

        x_coords = np.array(x_coords)
        initial_signal = np.array(initial_signal)

        second_deriv = np.gradient(np.gradient(savgol_filter(initial_signal, window_length, polyorder)))
        second_deriv = savgol_filter(second_deriv, 11, 3)

        x_out, final_signal = [], []
        for idx, col in enumerate(x_coords):
            column = img[:, col]
            nz_idx = np.where(column > 0)[0]
            if nz_idx.size == 0:
                continue

            if second_deriv[idx] > 0.5:
                sel_idx = nz_idx[:5]
            elif second_deriv[idx] < -0.5:
                sel_idx = nz_idx[-5:]
            else:
                sel_idx = nz_idx

            weights  = column[sel_idx].astype(float)
            centroid = np.average(sel_idx, weights=weights)
            x_out.append(col)
            final_signal.append(centroid)

        return np.array(x_out), np.array(final_signal)

    def fill_gaps(self, x_coords, y_coords, method='linear'):
        """Interpolate missing columns to produce a dense, evenly-spaced signal.

        Parameters
        ----------
        x_coords : np.ndarray of int
            Sparse column indices from :meth:`binarize_signal`.
        y_coords : np.ndarray of float
            Corresponding centroid Y-values.
        method : str, optional
            Interpolation kind passed to ``scipy.interpolate.interp1d``
            (default ``'linear'``).

        Returns
        -------
        x_full : np.ndarray of int
            Dense array ``[x_coords[0], …, x_coords[-1]]``.
        y_interp : np.ndarray of float
            Interpolated Y-values at every integer column.
        """
        x_full = np.arange(x_coords[0], x_coords[-1] + 1)
        interpolator = interp1d(x_coords, y_coords, kind=method, fill_value="extrapolate")
        y_interp = interpolator(x_full)
        return x_full, y_interp

    def smooth_signal(self, sig, window_length=7, polyorder=4):
        """Apply a Savitzky-Golay low-pass filter to a 1-D signal.

        Parameters
        ----------
        sig : np.ndarray
            Input signal array.
        window_length : int, optional
            Filter window size in samples (must be odd, default ``7``).
        polyorder : int, optional
            Polynomial order (default ``4``).

        Returns
        -------
        np.ndarray
            Smoothed signal of the same length as *sig*.
        """
        return savgol_filter(sig, window_length, polyorder)

    def extract_signals(self):
        """Convert mask grid cells into calibrated voltage time-series.

        For each cell in ``self.grid``:

        1. Calls :meth:`binarize_signal` to get column-centroid Y-values.
        2. Trims 5 edge samples, fills gaps, and smooths with Savitzky-Golay.
        3. Converts pixel Y-coordinates to millivolts using
           ``self.baseline`` and ``self.volt_per_pixel``.
        4. Converts column indices to seconds using ``self.time_per_pixel``.
        5. Resamples uniformly at 500 Hz via linear interpolation.

        Sets ``self.signal_grid`` with the same structure as ``self.grid``
        but each cell additionally contains:

        * ``'time'``   – np.ndarray of time values in seconds.
        * ``'signal'`` – np.ndarray of voltage values in millivolts.
        """
        signal_grid = []
        sample_rate = 500

        for row_idx, row in enumerate(self.grid):
            baseline_y = self.baseline[row_idx]
            row_signals = []

            for cell in row:
                x_coords, sig = self.binarize_signal(cell['signal'])

                x_coords = x_coords[5:-5]
                sig = sig[5:-5]

                x_coords, sig = self.fill_gaps(x_coords, sig, method='linear')
                sig = self.smooth_signal(sig)

                signal_volts = (baseline_y - sig) * self.volt_per_pixel
                x_seconds = x_coords * self.time_per_pixel

                duration = x_seconds[-1] - x_seconds[0]
                num_samples = round(duration * sample_rate) + 1
                resampled_time = np.linspace(x_seconds[0], x_seconds[-1], num_samples)

                interpolator = interp1d(x_seconds, signal_volts, kind='linear',
                                        fill_value="extrapolate")
                resampled_signal = interpolator(resampled_time)

                row_signals.append({
                    'lead': cell['lead'],
                    'time': resampled_time,
                    'signal': resampled_signal
                })

            signal_grid.append(row_signals)

        self.signal_grid = signal_grid

    # ------------------------------------------------------------------
    # Metrics & comparison
    # ------------------------------------------------------------------

    def sliding_metrics(self, signal_a, signal_b, return_aligned_signals=False):
        """Compute alignment-aware similarity metrics between two signals.

        Slides the shorter signal across the longer one and finds the offset
        that maximises the Pearson correlation.  Reports RMSE and SNR at
        the best-aligned position.

        Parameters
        ----------
        signal_a, signal_b : array-like
            The two signals to compare (need not be the same length).
        return_aligned_signals : bool, optional
            If ``True``, also return the aligned signal pair (default ``False``).

        Returns
        -------
        max_corr : float
            Maximum Pearson correlation found across all offsets.
        rmse : float
            Root mean squared error at the best offset.
        snr : float
            Signal-to-noise ratio in dB at the best offset
            (``inf`` if noise power is zero).
        short_signal : np.ndarray
            The shorter signal (only when *return_aligned_signals* is ``True``).
        best_window : np.ndarray
            The aligned window from the longer signal
            (only when *return_aligned_signals* is ``True``).
        """
        len_a, len_b = len(signal_a), len(signal_b)
        if len_a > len_b:
            long_signal, short_signal = signal_a, signal_b
        else:
            long_signal, short_signal = signal_b, signal_a

        len_long, len_short          = len(long_signal), len(short_signal)
        max_corr, best_corr_offset   = -1, 0

        for i in range(len_long - len_short + 1):
            window = long_signal[i:i + len_short]
            corr, _ = pearsonr(window, short_signal)
            if corr > max_corr:
                max_corr         = corr
                best_corr_offset = i

        best_window  = long_signal[best_corr_offset:best_corr_offset + len_short]
        rmse         = mean_squared_error(best_window, short_signal)
        signal_power = np.mean(np.square(short_signal))
        noise_power  = np.mean(np.square(best_window - short_signal))
        snr          = (10 * np.log10(signal_power / noise_power)
                        if noise_power > 0 else np.inf)

        if return_aligned_signals:
            return max_corr, rmse, snr, short_signal, best_window
        return max_corr, rmse, snr

    def calculate_metrics_ptb(self, plot_signals=True, per_lead_scores=None):
        """Evaluate digitization quality against a WFDB ground-truth record.

        For each lead in ``self.signal_grid``, aligns the extracted signal
        to the corresponding WFDB signal using :meth:`sliding_metrics` and
        records Pearson r, RMSE, SNR, and p-value.  Requires ``self.wfdb_path``
        to point to a valid WFDB record.

        Sets the following instance attributes after completion:

        * ``self.average_pearson`` – mean Pearson r across all leads.
        * ``self.average_rmse``    – mean RMSE (mV).
        * ``self.average_snr``     – mean SNR (dB).
        * ``self.average_pval``    – mean p-value.

        Parameters
        ----------
        plot_signals : bool, optional
            Unused placeholder (default ``True``).
        per_lead_scores : dict or None, optional
            If provided, per-lead metrics are appended here for leads where
            Pearson r > 0.60.  Expected structure:
            ``{lead_name: {'pearson': [], 'rmse': [], 'snr': [], 'pval': []}}``.
        """
        record = wfdb.rdrecord(self.wfdb_path)
        avg_pearson, avg_rmse, avg_snr, avg_pval = [], [], [], []

        for row in self.signal_grid:
            for cell in row:
                if 'lead' not in cell or 'signal' not in cell:
                    continue
                try:
                    lead_index = record.sig_name.index(cell['lead'])
                except ValueError:
                    continue

                wfdb_signal    = record.p_signal[:, lead_index]
                wfdb_signal    = wfdb_signal[~np.isnan(wfdb_signal)]
                voltage_signal = np.array(cell['signal'])

                pearson_val, rmse, snr, sig1, sig2 = self.sliding_metrics(
                    voltage_signal, wfdb_signal, return_aligned_signals=True
                )

                try:
                    pearson, pval = pearsonr(sig1, sig2)
                except Exception:
                    pearson, pval = np.nan, np.nan

                cell.update({'pearson': pearson, 'rmse': rmse,
                             'snr': snr,         'pval': pval})

                avg_pearson.append(pearson)
                avg_rmse.append(rmse)
                avg_snr.append(snr)
                avg_pval.append(pval)

                if per_lead_scores is not None and pearson > 0.60:
                    lead = cell['lead']
                    if lead not in per_lead_scores:
                        per_lead_scores[lead] = {
                            'pearson': [], 'rmse': [], 'snr': [], 'pval': []
                        }
                    per_lead_scores[lead]['pearson'].append(pearson)
                    per_lead_scores[lead]['rmse'].append(rmse)
                    per_lead_scores[lead]['snr'].append(snr)
                    per_lead_scores[lead]['pval'].append(pval)

        self.average_pearson = np.mean(avg_pearson)
        self.average_rmse    = np.mean(avg_rmse)
        self.average_snr     = np.mean(avg_snr)
        self.average_pval    = np.mean(avg_pval)

    def plot_signals(self, title='', plot_wfdb=False):
        """Plot extracted signals, optionally overlaid with the WFDB ground truth.

        Generates one figure per lead.  When *plot_wfdb* is ``True``, the
        extracted signal and the best-aligned WFDB window are plotted
        together for visual comparison.

        Parameters
        ----------
        title : str, optional
            Figure title applied to every subplot (default ``''``).
        plot_wfdb : bool, optional
            If ``True``, overlay the ground-truth WFDB signal.  Requires
            ``self.wfdb_path`` and a loaded ``record`` in scope
            (default ``False``).
        """
        for row in self.signal_grid:
            for cell in row:
                voltage_signal = cell['signal']
                plt.figure(figsize=(10, 4))
                if not plot_wfdb:
                    plt.plot(voltage_signal, linewidth=1.5)
                else:
                    lead_index  = record.sig_name.index(cell['lead'])
                    wfdb_signal = record.p_signal[:, lead_index]
                    wfdb_signal = [x for x in wfdb_signal if not np.isnan(x)]
                    _, _, _, sig1, sig2 = self.sliding_metrics(
                        voltage_signal, wfdb_signal, return_aligned_signals=True
                    )
                    plt.plot(sig1, label='Extracted Signal', linewidth=1.5)
                    plt.plot(sig2, label='Ground Truth',     linewidth=1.5)
                plt.title(title)
                plt.legend()
                plt.xlabel("Time (ms)")
                plt.ylabel("Voltage (mV)")
                plt.tight_layout()
                plt.show()

    # ------------------------------------------------------------------
    # Export
    # ------------------------------------------------------------------

    def save_signals_as_wfdb(self, record_name, directory='.'):
        """Export extracted signals as a WFDB record (.hea + .dat files).

        Pads all lead signals to the same length with NaN (converted to 0),
        then writes a 500 Hz, 16-bit WFDB record in millivolts.

        Parameters
        ----------
        record_name : str
            Base name for the output files (no extension).
        directory : str, optional
            Destination directory (default ``'.'``).
        """
        signals, lead_names, max_length = [], [], 0
        for row in self.signal_grid:
            for cell in row:
                if 'signal' in cell and 'lead' in cell:
                    sig = np.array(cell['signal'])
                    signals.append(sig)
                    lead_names.append(cell['lead'])
                    max_length = max(max_length, len(sig))

        padded = [np.pad(s, (0, max_length - len(s)),
                         mode='constant', constant_values=np.nan)
                  for s in signals]
        signal_array = np.nan_to_num(np.array(padded).T, nan=0.0)

        wfdb.wrsamp(
            record_name=record_name,
            fs=500,
            units=['mV'] * len(lead_names),
            sig_name=lead_names,
            p_signal=signal_array,
            fmt=['16'] * len(lead_names),
            write_dir=directory,
        )

    def save_signals_as_csv(self, record_name, directory='.'):
        """
        Save ECG signals to CSV. Each column = one lead, each row = one time sample.
        """
        signals, lead_names, max_length = [], [], 0
        for row in self.signal_grid:
            for cell in row:
                if 'signal' in cell and 'lead' in cell:
                    sig = np.asarray(cell['signal'], dtype=float)
                    signals.append(sig)
                    lead_names.append(cell['lead'])
                    max_length = max(max_length, len(sig))

        if not signals:
            raise ValueError("No valid signals found in signal_grid.")

        padded = [np.pad(s, (0, max_length - len(s)),
                         mode='constant', constant_values=np.nan)
                  for s in signals]
        df = pd.DataFrame(np.vstack(padded).T, columns=lead_names)

        os.makedirs(directory, exist_ok=True)
        csv_path = os.path.join(directory, f"{record_name}.csv")
        df.to_csv(csv_path, index=False, float_format="%.6f")

    # ------------------------------------------------------------------
    # Full pipeline
    # ------------------------------------------------------------------

    def run_full_pipeline(self):
        """Run all digitization stages end-to-end for a single ECG image.

        Executes the following stages in order, with per-stage wall-clock
        timing printed to stdout via :meth:`_timed`:

        1. **Image loading & preprocessing** — ``load_image`` + ``preprocess_image``.
        2. **Segmentation** — ``segment_leads`` → ``make_segmentation_mask``
           → ``find_row_centers`` → ``get_roi``.
        3. **Sequential YOLO detections** — ``extract_lead_boxes``,
           ``extract_lead_name_boxes``, ``extract_reference_pulses``.
           (Sequential, not parallel, to avoid a PyTorch thread-pool deadlock.)
        4. **Retry loop** — if any of the three detections is empty, the
           image is reloaded at the next fallback target size from
           ``[2000, 2100, 1900, 2200, 1800, 1700]`` px and stages 1–3 are
           re-run.  Raises ``RuntimeError`` if all sizes fail.
        5. **Calibration** — ``get_reference_scale``.
        6. **Grid & signal extraction** — ``make_grid`` → ``extract_signals``.

        Raises
        ------
        RuntimeError
            If all fallback target sizes yield at least one empty detection.
        """
        FALLBACK_SIZES = [2000, 2100, 1900, 2200, 1800, 1700]

        for target_size in FALLBACK_SIZES:

            # -- Image loading & preprocessing --
            self._timed("load_image",       self.load_image, self.image_path,
                        target_size=target_size)
            self._timed("preprocess",       self.preprocess_image)

            # -- Segmentation (sequential: each step depends on the previous) --
            self._timed("segment_leads",    self.segment_leads)
            self._timed("make_mask",        self.make_segmentation_mask)
            self._timed("find_row_centers", self.find_row_centers)
            self._timed("get_roi",          self.get_roi)

            # -- Three YOLO calls sequentially --
            self._timed("extract_lead_boxes",       self.extract_lead_boxes)
            self._timed("extract_lead_name_boxes",  self.extract_lead_name_boxes)
            self._timed("extract_reference_pulses", self.extract_reference_pulses)

            # -- Check all three detections; retry at next size if any is empty --
            missing = []
            if len(self.reference_pulses) == 0:
                missing.append("pulses")
            if len(self.lead_bboxes) == 0:
                missing.append("lead boxes")
            if len(self.lead_name_bboxes) == 0:
                missing.append("lead names")

            if not missing:
                if target_size != FALLBACK_SIZES[0]:
                    print(f"    ↳ All detections found at target_size={target_size}",
                          flush=True)
                break
            else:
                print(f"    ↳ Missing {', '.join(missing)} at target_size={target_size},"
                      f" retrying...", flush=True)

        else:
            raise RuntimeError(
                f"Detections incomplete at all target sizes {FALLBACK_SIZES}. "
                f"Last missing: {', '.join(missing)}."
            )

        # -- Remainder of pipeline (sequential) --
        self._timed("get_reference_scale", self.get_reference_scale)
        self._timed("make_grid",           self.make_grid)
        self._timed("extract_signals",     self.extract_signals)