File size: 7,203 Bytes
60b21d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | #!/usr/bin/env python3
# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
#
# SPDX-License-Identifier: MIT
"""
Test script for ECG-QA dataset implementation.
This script tests:
1. ECG-QA repository cloning
2. PTB-XL dataset downloading
3. Dataset loading and basic functionality
4. Sample data access
Usage:
python test_ecgqa.py
"""
import sys
def test_ecgqa_loader():
"""Test the ECG-QA loader functions."""
print("Testing ECG-QA loader...")
try:
from opentslm.time_series_datasets.ecg_qa.ecgqa_loader import (
does_ecg_qa_exist,
does_ptbxl_exist,
download_ecg_qa_if_not_exists,
download_ptbxl_if_not_exists
)
print(f"ECG-QA exists: {does_ecg_qa_exist()}")
print(f"PTB-XL exists: {does_ptbxl_exist()}")
# Download if needed (this might take a while)
print("Ensuring datasets are available...")
download_ecg_qa_if_not_exists()
download_ptbxl_if_not_exists()
print(f"After download - ECG-QA exists: {does_ecg_qa_exist()}")
print(f"After download - PTB-XL exists: {does_ptbxl_exist()}")
return True
except Exception as e:
print(f"Error in loader test: {e}")
return False
def test_ecgqa_dataset():
"""Test the ECG-QA dataset class."""
print("\nTesting ECGQADataset...")
try:
from opentslm.time_series_datasets.ecg_qa.ECGQADataset import ECGQADataset
# Try to create dataset instances with limited samples for faster testing
print("Creating dataset instances (limited to 5 samples each for testing)...")
dataset = ECGQADataset(split="train", EOS_TOKEN="", max_samples=5)
dataset_val = ECGQADataset(split="validation", EOS_TOKEN="", max_samples=5)
dataset_test = ECGQADataset(split="test", EOS_TOKEN="", max_samples=5)
print(f"Dataset sizes:")
print(f" Train: {len(dataset)} samples")
print(f" Validation: {len(dataset_val)} samples")
print(f" Test: {len(dataset_test)} samples")
if len(dataset) > 0:
print(f"\nExamining first training sample:")
sample = dataset[0]
print(f"Sample keys: {list(sample.keys())}")
print(f"Question: {sample.get('question', 'N/A')}")
print(f"Answer: {sample['answer']}")
print(f"Question type: {sample.get('question_type', 'N/A')}")
print(f"ECG IDs: {sample.get('ecg_id', 'N/A')}")
if 'time_series_text' in sample:
print(f"Number of time series: {len(sample['time_series_text'])}")
# Show first time series info and data
if len(sample['time_series_text']) > 0:
first_ts = sample['time_series_text'][0]
if hasattr(first_ts, 'text'):
print(f"First time series label: {first_ts.text}")
print(f"First time series length: {len(first_ts.time_series)}")
# Show the actual ECG data values
print(f"First 20 ECG data points: {first_ts.time_series[:20]}")
print(f"ECG data range: min={min(first_ts.time_series):.4f}, max={max(first_ts.time_series):.4f}")
# Show statistics
import numpy as np
ecg_data = np.array(first_ts.time_series)
print(f"ECG statistics: mean={np.mean(ecg_data):.4f}, std={np.std(ecg_data):.4f}")
# Show a few more leads if available
if len(sample['time_series_text']) > 1:
print(f"\nOther ECG leads available:")
for i, ts in enumerate(sample['time_series_text'][1:4], 1): # Show up to 3 more leads
if hasattr(ts, 'text'):
print(f" Lead {i+1}: {ts.text}")
print(f" First 10 values: {ts.time_series[:10]}")
ecg_data = np.array(ts.time_series)
print(f" Stats: mean={np.mean(ecg_data):.4f}, std={np.std(ecg_data):.4f}")
else:
print(f"Time series format issue - got {type(first_ts)}: {first_ts}")
# If it's a string, try to parse the actual time series data
if 'time_series' in sample:
print(f"\nRaw time_series field type: {type(sample['time_series'])}")
if isinstance(sample['time_series'], list) and len(sample['time_series']) > 0:
print(f"First time series data (first 20 points): {sample['time_series'][:20]}")
print(f"Time series length: {len(sample['time_series'])}")
import numpy as np
ts_data = np.array(sample['time_series'])
print(f"Time series stats: mean={np.mean(ts_data):.4f}, std={np.std(ts_data):.4f}")
print(f"\nPre-prompt: {sample['pre_prompt']}...")
print(f"Post-prompt: {sample['post_prompt']}...")
return True
except ImportError as e:
print(f"Import error (likely missing wfdb): {e}")
print("Please install wfdb: pip install wfdb")
return False
except Exception as e:
print(f"Error in dataset test: {e}")
import traceback
traceback.print_exc()
return False
def main():
"""Run all tests."""
print("ECG-QA Dataset Test Suite")
print("=" * 50)
# Test 1: Loader functionality
loader_success = test_ecgqa_loader()
# Test 2: Dataset functionality (only if loader works)
if loader_success:
dataset_success = test_ecgqa_dataset()
else:
dataset_success = False
print("Skipping dataset test due to loader failure")
# Summary
print("\n" + "=" * 50)
print("Test Results:")
print(f" Loader test: {'✓ PASS' if loader_success else '✗ FAIL'}")
print(f" Dataset test: {'✓ PASS' if dataset_success else '✗ FAIL'}")
if loader_success and dataset_success:
print("\n🎉 All tests passed! ECG-QA dataset is ready to use.")
print("\nTo use the full dataset (without sample limits), create ECGQADataset without max_samples parameter:")
print(" dataset = ECGQADataset(split='train', EOS_TOKEN='')")
else:
print("\n❌ Some tests failed. Check the error messages above.")
return loader_success and dataset_success
if __name__ == "__main__":
success = main()
sys.exit(0 if success else 1) |