code/scripts/find_zero_baseline.py

#!/usr/bin/env python3
"""
Generate a zero-mV ECG image using ECG-image-kit style rendering,
then pass through stage01_processor.py to find exact baseline positions.

This helps verify/correct the ZERO_MV constants used in training.
"""

import os
import sys
import numpy as np
import cv2
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from pathlib import Path
from scipy import signal as scipy_signal

# Add scripts to path
sys.path.insert(0, str(Path(__file__).parent))

# =============================================================================
# Configuration
# =============================================================================
RAW_WIDTH = 2200
RAW_HEIGHT = 1700

TARGET_WIDTH = 4352
TARGET_HEIGHT = 1696

T0, T1 = 235, 4161
OUTPUT_GT_WIDTH = T1 - T0  # 3926

# Current assumed baselines
ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])
MV_TO_PIXEL = 78.5

# Lead layout
LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]

# 0001 style colors
GRID_COLOR_MAJOR = '#FFB6B6'
GRID_COLOR_MINOR = '#FFD0D0'
TRACE_COLOR = '#000000'
BG_COLOR = '#FFFFFF'


def generate_zero_mv_image():
    """Generate an ECG image with all signals at exactly 0 mV."""
    
    # Create zero signals for all leads (10 seconds at 500Hz)
    fs = 500
    duration = 10  # seconds
    num_samples = int(fs * duration)
    
    # All leads are exactly 0 mV
    lead_signals = {
        'I': np.zeros(num_samples),
        'II': np.zeros(num_samples),
        'III': np.zeros(num_samples),
        'aVR': np.zeros(num_samples),
        'aVL': np.zeros(num_samples),
        'aVF': np.zeros(num_samples),
        'V1': np.zeros(num_samples),
        'V2': np.zeros(num_samples),
        'V3': np.zeros(num_samples),
        'V4': np.zeros(num_samples),
        'V5': np.zeros(num_samples),
        'V6': np.zeros(num_samples),
    }
    
    # Create figure
    dpi = 100
    fig_width = RAW_WIDTH / dpi
    fig_height = RAW_HEIGHT / dpi
    
    fig, ax = plt.subplots(1, 1, figsize=(fig_width, fig_height), dpi=dpi)
    fig.patch.set_facecolor(BG_COLOR)
    ax.set_facecolor(BG_COLOR)
    
    ax.set_xlim(0, RAW_WIDTH)
    ax.set_ylim(RAW_HEIGHT, 0)
    ax.set_aspect('equal')
    ax.axis('off')
    
    # Draw grid
    major_spacing = 40
    minor_spacing = 8
    
    for x in range(0, RAW_WIDTH + 1, minor_spacing):
        lw = 0.5 if x % major_spacing == 0 else 0.2
        color = GRID_COLOR_MAJOR if x % major_spacing == 0 else GRID_COLOR_MINOR
        ax.axvline(x, color=color, linewidth=lw)
    
    for y in range(0, RAW_HEIGHT + 1, minor_spacing):
        lw = 0.5 if y % major_spacing == 0 else 0.2
        color = GRID_COLOR_MAJOR if y % major_spacing == 0 else GRID_COLOR_MINOR
        ax.axhline(y, color=color, linewidth=lw)
    
    # Scale factors
    raw_scale_y = RAW_HEIGHT / TARGET_HEIGHT
    raw_baselines = ZERO_MV * raw_scale_y
    raw_mv_to_pixel = MV_TO_PIXEL * raw_scale_y
    
    samples_per_column = int(2.5 * fs)
    column_width = RAW_WIDTH / 4
    
    # Draw horizontal lines at 0 mV baseline for each row
    for col_idx in range(3):
        for row_idx in range(4):
            lead_name = LEAD_LAYOUT[col_idx][row_idx]
            
            # 0 mV signal - straight horizontal line at baseline
            x_start = col_idx * column_width + 50
            x_end = (col_idx + 1) * column_width - 10
            y = raw_baselines[row_idx]  # 0 mV = baseline
            
            # Draw thin line (the signal)
            ax.plot([x_start, x_end], [y, y], color=TRACE_COLOR, linewidth=0.8)
    
    # Draw Lead II full
    x_start = 3 * column_width + 20
    x_end = RAW_WIDTH - 20
    y = raw_baselines[1]  # Row 1 for lead II
    ax.plot([x_start, x_end], [y, y], color=TRACE_COLOR, linewidth=0.8)
    
    plt.tight_layout(pad=0)
    
    # Save to buffer
    fig.canvas.draw()
    img = np.array(fig.canvas.buffer_rgba())[:, :, :3]  # RGB from RGBA
    plt.close(fig)
    
    # Resize to exact dimensions
    img = cv2.resize(img, (RAW_WIDTH, RAW_HEIGHT))
    
    return img, lead_signals


def process_through_stage01(img):
    """Process image through stage01_processor.py"""
    from subprocess import run, PIPE
    import tempfile
    
    # Save image temporarily
    with tempfile.TemporaryDirectory() as tmpdir:
        in_path = Path(tmpdir) / 'zero_mv_input.png'
        out_path = Path(tmpdir) / 'zero_mv_output.png'
        
        cv2.imwrite(str(in_path), cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
        
        # Run stage01_processor.py
        stage01_dir = Path('/home/azureuser/ecg-digitization/src/stage01')
        processor_script = stage01_dir / 'stage01_processor.py'
        
        if processor_script.exists():
            cmd = [
                'python3', str(processor_script),
                str(in_path), str(out_path)
            ]
            result = run(cmd, cwd=str(stage01_dir), capture_output=True, text=True)
            print("Stage01 stdout:", result.stdout)
            print("Stage01 stderr:", result.stderr)
            
            if out_path.exists():
                processed = cv2.imread(str(out_path))
                processed = cv2.cvtColor(processed, cv2.COLOR_BGR2RGB)
                return processed
        
    return None


def find_baselines_from_image(img):
    """
    Find the baseline y-positions from a zero-mV processed image.
    The black trace lines should be exactly at the baselines.
    """
    # Convert to grayscale
    if len(img.shape) == 3:
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    else:
        gray = img
    
    # Find dark pixels (the trace)
    dark_mask = gray < 50  # Black pixels
    
    height, width = gray.shape
    print(f"\nProcessed image shape: {height} x {width}")
    
    # For each row region (4 rows), find the horizontal line
    # Row heights in the processed image (TARGET_HEIGHT = 1696)
    row_height = height // 4
    
    baselines = []
    
    for row_idx in range(4):
        y_start = row_idx * row_height
        y_end = (row_idx + 1) * row_height
        
        # Find the y-coordinate with most dark pixels in this region
        row_region = dark_mask[y_start:y_end, :]
        y_profile = row_region.sum(axis=1)  # Sum across width
        
        if y_profile.max() > 0:
            # Find the peak (where the line is)
            local_y = np.argmax(y_profile)
            global_y = y_start + local_y
            baselines.append(global_y)
            print(f"Row {row_idx}: baseline at y={global_y:.1f} (local: {local_y})")
        else:
            print(f"Row {row_idx}: no trace found, using expected {ZERO_MV[row_idx]:.1f}")
            baselines.append(ZERO_MV[row_idx])
    
    return np.array(baselines)


def main():
    print("=" * 70)
    print("Finding Zero-mV Baseline Positions for ECG Image Kit Synthetic Data")
    print("=" * 70)
    
    # Generate zero-mV image
    print("\n1. Generating zero-mV ECG image...")
    img, signals = generate_zero_mv_image()
    print(f"   Raw image shape: {img.shape}")
    
    # Save raw image
    out_dir = Path('/data/ecg-digitization/baseline_test')
    out_dir.mkdir(parents=True, exist_ok=True)
    
    raw_path = out_dir / 'zero_mv_raw.png'
    cv2.imwrite(str(raw_path), cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
    print(f"   Saved: {raw_path}")
    
    # Check if stage01_processor exists
    stage01_processor = Path('/home/azureuser/ecg-digitization/src/stage01/stage01_processor.py')
    
    if stage01_processor.exists():
        print("\n2. Processing through stage01...")
        processed = process_through_stage01(img)
        
        if processed is not None:
            proc_path = out_dir / 'zero_mv_processed.png'
            cv2.imwrite(str(proc_path), cv2.cvtColor(processed, cv2.COLOR_RGB2BGR))
            print(f"   Saved: {proc_path}")
            
            print("\n3. Finding baselines from processed image...")
            measured_baselines = find_baselines_from_image(processed)
        else:
            print("   Stage01 processing failed, using raw image")
            measured_baselines = find_baselines_from_image(img)
    else:
        print("\n2. Stage01 processor not found, analyzing raw image directly")
        # Scale raw image to target dimensions
        img_scaled = cv2.resize(img, (TARGET_WIDTH, TARGET_HEIGHT))
        measured_baselines = find_baselines_from_image(img_scaled)
    
    # Compare with expected
    print("\n" + "=" * 70)
    print("BASELINE COMPARISON")
    print("=" * 70)
    print(f"{'Row':<8} {'Expected':<15} {'Measured':<15} {'Diff (px)':<15} {'Diff (mV)':<15}")
    print("-" * 70)
    
    for i in range(4):
        expected = ZERO_MV[i]
        measured = measured_baselines[i]
        diff_px = measured - expected
        diff_mv = diff_px / MV_TO_PIXEL
        print(f"{i:<8} {expected:<15.1f} {measured:<15.1f} {diff_px:<+15.2f} {diff_mv:<+15.4f}")
    
    print("\n" + "=" * 70)
    print("RECOMMENDED ZERO_MV VALUES:")
    print("=" * 70)
    print(f"ZERO_MV = np.array([{measured_baselines[0]:.1f}, {measured_baselines[1]:.1f}, {measured_baselines[2]:.1f}, {measured_baselines[3]:.1f}])")
    
    return measured_baselines


if __name__ == '__main__':
    main()