#!/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()