Spaces:
Runtime error
Runtime error
| import tensorflow as tf | |
| from tensorflow import keras | |
| from tensorflow.keras import layers, regularizers | |
| import numpy as np | |
| import json | |
| import gradio as gr | |
| from scipy.signal import butter, sosfiltfilt, iirnotch, tf2sos | |
| L2 = regularizers.l2(2e-4) | |
| # ============ temporal branch =============== | |
| def se_block_1d(x, reduction=8, name=None): | |
| channels = x.shape[-1] | |
| se = layers.GlobalAveragePooling1D()(x) | |
| se = layers.Dense(max(channels // reduction, 4), activation="relu")(se) | |
| se = layers.Dense(channels, activation="sigmoid")(se) | |
| se = layers.Reshape((1, channels))(se) | |
| return layers.Multiply(name=name)([x, se]) | |
| def tcn_residual_block(x, filters, kernel_size, dilation_rate, dropout_rate=0.3): | |
| previous = x | |
| x = layers.Conv1D(filters, kernel_size, padding="causal", dilation_rate=dilation_rate, | |
| kernel_regularizer=L2)(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.Dropout(dropout_rate)(x) | |
| x = layers.Conv1D(filters, kernel_size, padding="causal", dilation_rate=dilation_rate, | |
| kernel_regularizer=L2)(x) | |
| x = layers.BatchNormalization()(x) | |
| if previous.shape[-1] != filters: | |
| previous = layers.Conv1D(filters, 1, padding="same")(previous) | |
| x = layers.Add()([previous, x]) | |
| x = layers.Activation("relu")(x) | |
| return x | |
| def build_tcn(x, filters=64, kernel_size=3, num_blocks=4, dropout_rate=0.4): | |
| for i in range(num_blocks): | |
| x = tcn_residual_block(x, filters, kernel_size, dilation_rate=2 ** i, dropout_rate=dropout_rate) | |
| return x | |
| def temporal_branch(signal_input): | |
| x = layers.Conv1D(48, kernel_size=7, padding="same", kernel_regularizer=L2)(signal_input) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.Conv1D(96, kernel_size=5, padding="same", kernel_regularizer=L2)(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.MaxPooling1D(2)(x) | |
| x = layers.Conv1D(128, kernel_size=5, padding="same", kernel_regularizer=L2)(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.MaxPooling1D(2)(x) | |
| x = build_tcn(x, filters=64, num_blocks=4, dropout_rate=0.3) | |
| x = se_block_1d(x, reduction=8, name="se_attention") | |
| return layers.GlobalAveragePooling1D()(x) | |
| # ===================== Spectral branch ============================== | |
| class STFTLayer(layers.Layer): | |
| def __init__(self, frame_length=256, frame_step= 64, fft_length=256, **kwargs): | |
| super().__init__(**kwargs) | |
| self.frame_length = frame_length | |
| self.frame_step = frame_step | |
| self.fft_length = fft_length | |
| def call(self, x): | |
| x = tf.transpose(x, [0, 2, 1]) | |
| spec = tf.signal.stft(x, frame_length=self.frame_length, | |
| frame_step=self.frame_step, fft_length=self.fft_length) | |
| mag = tf.math.log1p(tf.abs(spec)) | |
| return tf.transpose(mag, [0, 3, 2, 1]) | |
| def compute_output_shape(self, input_shape): | |
| batch, time, leads = input_shape | |
| frames = 1 + (time - self.frame_length) // self.frame_step | |
| freq_bins = self.fft_length // 2 + 1 | |
| return (batch, freq_bins, frames, leads) | |
| def get_config(self): | |
| config = super().get_config() | |
| config.update(frame_length=self.frame_length, | |
| frame_step=self.frame_step, | |
| fft_length=self.fft_length) | |
| return config | |
| def transformer_encoder_block(x, num_heads=4, key_dim=32, ff_dim=128, dropout_rate=0.3): | |
| attn_out = layers.MultiHeadAttention(num_heads=num_heads, key_dim=key_dim, dropout=dropout_rate)(x, x) | |
| x = layers.Add()([x, attn_out]) | |
| x = layers.LayerNormalization()(x) | |
| ff = layers.Dense(ff_dim, activation="gelu", kernel_regularizer=L2)(x) | |
| ff = layers.Dense(x.shape[-1])(ff) | |
| ff = layers.Dropout(dropout_rate)(ff) | |
| x = layers.Add()([x, ff]) | |
| x = layers.LayerNormalization()(x) | |
| return x | |
| class AttentionPooling1D(layers.Layer): | |
| def __init__(self, units=64, **kwargs): | |
| super().__init__(**kwargs) | |
| self.units = units | |
| def build(self, input_shape): | |
| input_shape = tuple(input_shape) | |
| self.score_dense1 = layers.Dense(self.units, activation="tanh") | |
| self.score_dense2 = layers.Dense(1) | |
| self.score_dense1.build(input_shape) # new | |
| self.score_dense2.build(input_shape[:-1] + (self.units,)) # new | |
| super().build(input_shape) | |
| def call(self, x): | |
| scores = self.score_dense2(self.score_dense1(x)) | |
| weights = tf.nn.softmax(scores, axis=1) | |
| return tf.reduce_sum(x * weights, axis=1) | |
| def compute_output_shape(self, input_shape): | |
| return (input_shape[0], input_shape[-1]) | |
| def get_config(self): # added | |
| config = super().get_config() | |
| config.update(units=self.units) | |
| return config | |
| def spectral_branch(signal_input): | |
| x = STFTLayer()(signal_input) | |
| x = layers.Conv2D(24, (3, 3), padding="same", kernel_regularizer=L2)(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.MaxPooling2D((2, 1))(x) | |
| x = layers.Conv2D(48, (3, 3), padding="same", kernel_regularizer=L2)(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.MaxPooling2D((2, 1))(x) | |
| shape = x.shape | |
| x = layers.Reshape((shape[1] * shape[2], shape[3]))(x) | |
| x = transformer_encoder_block(x, num_heads=4, key_dim=32, ff_dim=128) | |
| return AttentionPooling1D()(x) | |
| # ================== Metadata branch ============================== | |
| def patient_branch(patient_input): | |
| p = layers.Dense(32, activation="relu")(patient_input) | |
| p = layers.BatchNormalization()(p) | |
| p = layers.Dropout(0.2)(p) | |
| p = layers.Dense(16, activation="relu")(p) | |
| p = layers.BatchNormalization()(p) | |
| return p | |
| # =================== Fusion + classifier ============================ | |
| def gated_multimodal_fusion(temporal_vec, spectral_vec, dim=96): | |
| t = layers.Dense(dim, activation="tanh")(temporal_vec) | |
| s = layers.Dense(dim, activation="tanh")(spectral_vec) | |
| gate_input = layers.Concatenate()([temporal_vec, spectral_vec]) | |
| # gate_hidden = layers.Dense(dim, activation="relu")(gate_input) | |
| z = layers.Dense(dim, activation="sigmoid", name="fusion_gate")(gate_input) | |
| fused = layers.Add(name="fused_features")([ | |
| layers.Multiply()([z, t]), | |
| layers.Multiply()([layers.Lambda(lambda x: 1 - x)(z), s]) | |
| ]) | |
| return fused | |
| def classifier(fusion, num_classes): | |
| x = layers.Dense(256, kernel_regularizer=L2)(fusion) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.Dropout(0.4)(x) | |
| x = layers.Dense(128, kernel_regularizer=L2)(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.Dropout(0.3)(x) | |
| return layers.Dense(num_classes, activation="sigmoid", name="diagnosis_output")(x) | |
| # =================== Model Build ==================== | |
| def build_model(signal_len, num_leads, num_classes): | |
| input_ecg = layers.Input(shape=(signal_len, num_leads), name="ecg_signal") | |
| patient_input = layers.Input(shape=(2,), name="patient_input") | |
| temporal = temporal_branch(input_ecg) | |
| spectral = spectral_branch(input_ecg) | |
| patient_features = patient_branch(patient_input) | |
| fusion = gated_multimodal_fusion(temporal, spectral) | |
| fusion = layers.Concatenate(name="fusion_with_patient")([fusion, patient_features]) | |
| output = classifier(fusion, num_classes) | |
| return keras.Model(inputs=[input_ecg, patient_input], outputs=output, name="ECG_DualBranch_v2") | |
| # laod pre trained model | |
| model = build_model( signal_len=1000, num_leads=12, num_classes=5 ) | |
| model.load_weights( "ecg_model.weights.h5" ) | |
| p = json.load(open("preprocess.json")) | |
| def prepare(ecg): # ecg: raw signal, shape (1000, 12), 100 Hz | |
| sos = butter(4, [0.5/50, 40/50], btype="band", output="sos") | |
| x = sosfiltfilt(sos, ecg.astype(np.float32), axis=0) | |
| x = sosfiltfilt(tf2sos(*iirnotch(50, 30, fs=100)), x, axis=0) | |
| return ((x - np.array(p["lead_mean"])) / np.array(p["lead_std"])).astype(np.float32) | |
| def predict(ecg_file, age, sex): | |
| # read uploaded ECG | |
| ecg=np.fromfile( ecg_file, dtype=np.int16 ) | |
| ecg=( ecg.reshape(-1,12)/1000 ) | |
| ecg=ecg.astype(np.float32) | |
| x = prepare(ecg)[None] # shape (1, 1000, 12) | |
| patient = np.array([[(age - p["age_mean"]) / p["age_std"], sex]], dtype=np.float32) | |
| probs = model.predict([x, patient], verbose=0)[0] | |
| return dict(zip(p["classes"], probs.round(3))) | |
| # ============================ Gradio UI =========================== | |
| import gradio as gr | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.File(label="Upload ECG .dat file"), | |
| gr.Number(label="Age", value=50), | |
| gr.Radio([0, 1], label="Sex (0 Female, 1 Male)") | |
| ], | |
| outputs=gr.Label(label="Prediction", num_top_classes=5), | |
| title="12 Lead ECG Multi-label Classification" | |
| ) | |
| demo.launch() |