File size: 1,740 Bytes
a2ccce3 | 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 | import numpy as np
import tensorflow as tf
# Load trained model
MODEL_PATH = "ptbxl_ecg_cnn_model.keras"
model = tf.keras.models.load_model(MODEL_PATH, compile=False)
# Diagnostic labels
TARGET_NAMES = ["NORM", "MI", "STTC", "CD", "HYP"]
# Default thresholds (can be replaced with tuned ones)
DEFAULT_THRESHOLDS = {
"NORM": 0.5,
"MI": 0.5,
"STTC": 0.5,
"CD": 0.5,
"HYP": 0.5
}
def predict_ecg(ecg_signal, thresholds=DEFAULT_THRESHOLDS):
"""
Predict cardiac abnormalities from a 12-lead ECG.
Parameters
----------
ecg_signal : np.ndarray
Shape (1000, 12), preprocessed ECG signal
thresholds : dict
Thresholds for each class
Returns
-------
predicted_labels : list
List of predicted diagnostic labels
probabilities : dict
Probability per diagnostic class
"""
if ecg_signal.shape != (1000, 12):
raise ValueError("ECG signal must have shape (1000, 12)")
# Add batch dimension
ecg_signal = np.expand_dims(ecg_signal, axis=0)
# Model prediction
probs = model.predict(ecg_signal, verbose=0)[0]
predicted_labels = []
probabilities = {}
for i, label in enumerate(TARGET_NAMES):
probabilities[label] = float(probs[i])
if probs[i] >= thresholds[label]:
predicted_labels.append(label)
return predicted_labels, probabilities
# Example usage (for testing only)
if __name__ == "__main__":
dummy_ecg = np.random.randn(1000, 12)
labels, probs = predict_ecg(dummy_ecg)
print("Predicted labels:", labels)
print("Probabilities:")
for k, v in probs.items():
print(f"{k}: {v:.3f}")
|