app.py

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