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()