A Combined ResNet-DenseNet Architecture with ResU Blocks (ResU-Dense) for 12-lead ECG Abnormality Classification
import os
import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from utils import load_recording, normalize_ecg
from scipy import signal
import numpy as np
import random


def create_dataloader(opt):
    """Create training and validation data loaders."""
    print('Load the dataset...')

    # Read train and validation datasets from CSV
    train_header = pd.read_csv('./dataset/train_dataset.csv')
    valid_header = pd.read_csv('./dataset/test_dataset.csv')
    train = pd.DataFrame(train_header)
    valid = pd.DataFrame(valid_header)

    print(len(train))

    # Create training data loader
    train_loader = DataLoader(
        dataset=ECG_Dataset(opt, train, mode='train'),
        batch_size=opt.batch_size,
        shuffle=True,
        num_workers=0,
        pin_memory=True,
        drop_last=True,
        sampler=None
    )

    # Create validation data loader
    validatioin_loader = DataLoader(
        dataset=ECG_Dataset(opt, valid, mode='valid'),
        batch_size=opt.batch_size, shuffle=False, num_workers=0
    )

    return train_loader, validatioin_loader


def create_dataloader_for_test(opt):
    """Create test data loader."""

    print('Load the dataset...')
    if not os.path.exists('./dataset/test_dataset.csv'):
        exit()  # Exit if test dataset CSV doesn't exist
    else:
        test = pd.read_csv('./dataset/test_dataset.csv')
        test = pd.DataFrame(test)
        print(len(test))

    # Create test data loader
    test_loader = DataLoader(
        dataset=ECG_Dataset(opt, test, mode='test'),
        batch_size=1,
        shuffle=False,
        num_workers=0
    )

    return test_loader


class ECG_Dataset(Dataset):
    """ECG Dataset class that implements the required methods for PyTorch's Dataset class."""
    def __init__(self, opt, dataset, mode):
        # Initialize variables and load data
        self.fs = opt.fs
        self.samples = opt.samples
        self.mode = mode
        self.dataset = dataset

    def __len__(self):
        """Return the number of items in the dataset."""
        return len(self.dataset)

    def __getitem__(self, idx):
        """Return an item from the dataset at the specified index."""
        fs = self.dataset.iloc[idx]['fs']  # Fetch the sampling frequency
        targets = self.dataset.iloc[idx]['target']  # Fetch the targets

        # Load ECG recording
        inputs = load_recording(self.dataset.iloc[idx]['record'])

        # [Step 1] Normalization: Use z-score normalization
        inputs = normalize_ecg(inputs)

        # [Step 2] Resample data to 500Hz if needed
        if fs == float(1000):
            inputs = signal.resample_poly(inputs, up=1, down=2, axis=-1)  # to 500Hz
        elif fs == float(500):
            pass
        else:
            inputs = signal.resample(inputs, int(inputs.shape[1] * 500 / fs), axis=1)

        inputs = np.nan_to_num(inputs)  # Convert NaN values to zero
        inputs = torch.from_numpy(inputs)  # Convert to PyTorch tensor

        # [Step 3] Zero-Padding: Adjust input length to a fixed size (self.samples)
        if inputs.size(1) > self.samples:
            start_index = random.randint(0, inputs.size(1) - self.samples - 1)
            inputs = inputs[:, start_index:start_index + self.samples]
        else:
            pad_len = self.samples - inputs.size(1)
            inputs = torch.cat([inputs, torch.zeros([inputs.size(0), pad_len])], dim=1)

        # Process and convert targets to tensor
        targets = list(map(float, targets.replace('[', '').replace(']', '').split()))
        targets = torch.Tensor(targets)

        return inputs, targets
#
dataloader/dataloader.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
import torch
import torch.nn as nn
import torch.nn.functional as functional
import math


class Bottleneck(nn.Module):
    def __init__(self, nChannels, growthRate):
        super(Bottleneck, self).__init__()
        interChannels = 4 * growthRate
        self.bn1 = nn.BatchNorm1d(nChannels)
        self.conv1 = nn.Conv1d(nChannels, interChannels, kernel_size=1,
                               bias=False)
        self.bn2 = nn.BatchNorm1d(interChannels)
        self.conv2 = nn.Conv1d(interChannels, growthRate, kernel_size=3,
                               padding=1, bias=False)

    def forward(self, x):
        out = self.conv1(functional.relu(self.bn1(x)))
        out = self.conv2(functional.relu(self.bn2(out)))
        out = torch.cat((x, out), 1)
        return out


class SingleLayer(nn.Module):
    def __init__(self, nChannels, growthRate):
        super(SingleLayer, self).__init__()
        self.bn1 = nn.BatchNorm1d(nChannels)
        self.conv1 = nn.Conv1d(nChannels, growthRate, kernel_size=3,
                               padding=1, bias=False)

    def forward(self, x):
        out = self.conv1(functional.relu(self.bn1(x)))
        out = torch.cat((x, out), 1)
        return out


class Transition(nn.Module):
    def __init__(self, nChannels, nOutChannels, down=False):
        super(Transition, self).__init__()
        self.bn1 = nn.BatchNorm1d(nChannels)
        self.conv1 = nn.Conv1d(nChannels, nOutChannels, kernel_size=1,
                               bias=False)
        self.down = down

    def forward(self, x):
        out = self.conv1(functional.relu(self.bn1(x)))
        if self.down:
            out = functional.avg_pool1d(out, 2)
        return out


class ResidualUBlock(nn.Module):
    def __init__(self, out_ch, mid_ch, layers, downsampling=True):
        super(ResidualUBlock, self).__init__()
        self.downsample = downsampling  # Flag to decide if down-sampling is needed
        K = 9  # Kernel size
        P = (K - 1) // 2  # Padding calculation

        # Initial convolutional layer
        self.conv1 = nn.Conv1d(in_channels=out_ch,
                               out_channels=out_ch,
                               kernel_size=K,
                               padding=P,
                               bias=False)
        self.bn1 = nn.BatchNorm1d(out_ch)

        self.encoders = nn.ModuleList()  # Encoder layers
        self.decoders = nn.ModuleList()  # Decoder layers

        # Creating encoder-decoder blocks
        for idx in range(layers):
            # Encoder block definition
            if idx == 0:
                # First encoder has different input channels
                self.encoders.append(nn.Sequential(
                    nn.Conv1d(
                        in_channels=out_ch,
                        out_channels=mid_ch,
                        kernel_size=K,
                        stride=2,
                        padding=P,
                        bias=False
                    ),
                    nn.BatchNorm1d(mid_ch),
                    nn.LeakyReLU()
                ))
            else:
                # Subsequent encoders use mid_ch as input
                self.encoders.append(nn.Sequential(
                    nn.Conv1d(
                        in_channels=mid_ch,
                        out_channels=mid_ch,
                        kernel_size=K,
                        stride=2,
                        padding=P,
                        bias=False
                    ),
                    nn.BatchNorm1d(mid_ch),
                    nn.LeakyReLU()
                ))

            # Decoder block definition
            if idx == layers - 1:
                # Last decoder has different output channels
                self.decoders.append(nn.Sequential(
                    nn.ConvTranspose1d(
                        in_channels=mid_ch * 2,
                        out_channels=out_ch,
                        kernel_size=K,
                        stride=2,
                        padding=P,
                        output_padding=1,
                        bias=False
                    ),
                    nn.BatchNorm1d(out_ch),
                    nn.LeakyReLU()
                ))
            else:
                # Subsequent decoders output to mid_ch
                self.decoders.append(nn.Sequential(
                    nn.ConvTranspose1d(
                        in_channels=mid_ch * 2,
                        out_channels=mid_ch,
                        kernel_size=K,
                        stride=2,
                        padding=P,
                        output_padding=1,
                        bias=False
                    ),
                    nn.BatchNorm1d(mid_ch),
                    nn.LeakyReLU()
                ))

        # Bottleneck layer (center of U-Net)
        self.bottleneck = nn.Sequential(
            nn.Conv1d(
                in_channels=mid_ch,
                out_channels=mid_ch,
                kernel_size=K,
                padding=P,
                bias=False
            ),
            nn.BatchNorm1d(mid_ch),
            nn.LeakyReLU()
        )

        # Down-sampling layers (if required)
        if self.downsample:
            self.idfunc_0 = nn.AvgPool1d(kernel_size=2, stride=2)
            self.idfunc_1 = nn.Conv1d(in_channels=out_ch,
                                      out_channels=out_ch,
                                      kernel_size=1,
                                      bias=False)

    def forward(self, x):
        x_in = functional.leaky_relu(self.bn1(self.conv1(x)))

        out = x_in
        encoder_out = []
        for idx, layer in enumerate(self.encoders):
            # If output size is not divisible by 4, padding is added
            if out.size(-1) % 4 != 0:
                out = functional.pad(out, [1, 0, 0, 0])
            out = layer(out)
            encoder_out.append(out)

        out = self.bottleneck(out)

        for idx, layer in enumerate(self.decoders):
            out = layer(torch.cat([out, encoder_out[-1 - idx]], dim=1))

        # Trim the output to match the size of x_in (input)
        out = out[..., :x_in.size(-1)]

        out += x_in

        # If down-sampling is required, apply down-sampling layers
        if self.downsample:
            out = self.idfunc_0(out)
            out = self.idfunc_1(out)

        return out


def _make_dense(nChannels, growthRate, nDenseBlocks, bottleneck):
    layers = []
    for i in range(int(nDenseBlocks)):
        if bottleneck:
            layers.append(Bottleneck(nChannels, growthRate))
        else:
            layers.append(SingleLayer(nChannels, growthRate))
        nChannels += growthRate
    return nn.Sequential(*layers)


class ResU_Dense(nn.Module):
    def __init__(self, nOUT, in_ch=12, out_ch=256, mid_ch=64):
        super(ResU_Dense, self).__init__()
        # Initial convolutional layer
        self.conv = nn.Conv1d(in_channels=in_ch,
                              out_channels=out_ch,
                              kernel_size=15,
                              padding=7,
                              stride=2,
                              bias=False)
        self.bn = nn.BatchNorm1d(out_ch)

        # Define Residual U-blocks
        self.rub_0 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=6)
        self.rub_1 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=5)
        self.rub_2 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=4)
        self.rub_3 = ResidualUBlock(out_ch=out_ch, mid_ch=mid_ch, layers=3)

        # Parameters for dense blocks
        growthRate = 12
        reduction = 0.5
        nChannels = out_ch
        nDenseBlocks = 16

        # Define dense blocks and transitions
        self.dense1 = _make_dense(nChannels, growthRate=12, nDenseBlocks=nDenseBlocks, bottleneck=True)
        nChannels += nDenseBlocks * growthRate
        nOutChannels = int(math.floor(nChannels * reduction))
        self.trans1 = Transition(nChannels, nOutChannels)

        nChannels = nOutChannels
        self.dense2 = _make_dense(nChannels, growthRate=12, nDenseBlocks=nDenseBlocks, bottleneck=True)
        nChannels += nDenseBlocks * growthRate
        self.trans2 = Transition(nChannels, out_ch)

        # Multihead attention layer
        self.mha = nn.MultiheadAttention(out_ch, 8)
        # Max pooling layer to reduce dimensions
        self.pool = nn.AdaptiveMaxPool1d(output_size=1)

        # Fully connected layer
        self.fc_1 = nn.Linear(out_ch, nOUT)

    def forward(self, x):
        x = functional.leaky_relu(self.bn(self.conv(x)))

        # Pass through the residual U-blocks
        x = self.rub_0(x)
        x = self.rub_1(x)
        x = self.rub_2(x)
        x = self.rub_3(x)

        # Pass through the dense blocks and transitions
        x = self.trans1(self.dense1(x))
        x = self.trans2(self.dense2(x))

        # Apply dropout for regularization
        x = functional.dropout(x, p=0.5, training=self.training)

        x = x.permute(2, 0, 1)
        x, _ = self.mha(x, x, x)
        x = x.permute(1, 2, 0)

        # Reduce dimensions with pooling
        x = self.pool(x).squeeze(2)

        # Fully connected layer for final output
        x = self.fc_1(x)
        return x

models/ResU_Dense.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
from .ResU_Dense import ResU_Dense

models/__init__.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
"""
Docstring for Options:
This class serves to provide a centralized argument parser. It defines the hyperparameters and settings
used throughout the ECG classification training process, including data processing, model architecture,
and training configurations.
"""


class Options:
    def __init__(self):
        pass

    def init(self, parser):
        # Global settings
        parser.add_argument('--batch_size', type=int, default=24,
                            help='Size of each batch during training.')
        parser.add_argument('--nepoch', type=int, default=30,
                            help='Total number of training epochs.')
        parser.add_argument('--optimizer', type=str, default='adam',
                            help='Type of optimizer to use during training (e.g., "adam", "sgd").')
        parser.add_argument('--lr_initial', type=float, default=1e-3,
                            help='Initial learning rate.')
        parser.add_argument("--decay_epoch", type=int, default=20,
                            help="Epoch number from which to start learning rate decay.")
        parser.add_argument('--weight_decay', type=float, default=0.1,
                            help='Weight decay factor for regularization.')
        parser.add_argument('--valid_ratio', type=float, default=0.2,
                            help='Proportion of data to be used for validation.')

        # Training settings
        parser.add_argument('--arch', type=str, default='ResU_Dense',
                            help='Model architecture to be used for training (e.g., "ResU_Dense").')
        parser.add_argument('--loss_type', type=str, default='base',
                            help='Type of loss function to be used (e.g., "base", "mse").')
        parser.add_argument('--loss_oper', type=str, default='base',
                            help='Operation type of the loss function (e.g., "add", "multiply").')
        parser.add_argument('--device', type=str, default='cuda',
                            help='Device to be used for training ("cuda" for GPU, "cpu" for CPU).')

        # Network settings
        parser.add_argument('--class_num', type=int, default=26,
                            help='Number of classes for ECG classification.')
        parser.add_argument('--leads', type=int, default=12,
                            help='Number of ECG leads to be considered.')

        # Pretrained model settings
        parser.add_argument('--env', type=str, default='0901',
                            help='Name used for logging purposes.')
        parser.add_argument('--pretrained', type=bool, default=False,
                            help='Option to load weights from a pretrained model.')
        parser.add_argument('--pretrain_model_path', type=str,
                            default='./log/ResU_Dense/models/chkpt_opt.pt',
                            help='Path to the pretrained model weights.')

        # Dataset settings
        parser.add_argument('--fs', type=int, default=500,
                            help='Sampling frequency of the ECG data.')
        parser.add_argument('--samples', type=int, default=4096,
                            help='Number of samples to consider from the ECG data.')
        parser.add_argument('--dirs_for_train', type=str,
                            default='../Dataset/',
                            help='Directory path where ECG dataset for training and validation is located.')

        return parser

options.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
"""
Test interface for speech enhancement!
You can just run this file.
"""
import os
import argparse
import torch
import options
import utils
import datetime
import random
import numpy as np
from dataloader import create_dataloader_for_test
from sklearn.metrics import accuracy_score

######################################################################################################################
#                                                  Parser init                                                       #
######################################################################################################################
# Parse command-line arguments and configurations for the experiment
opt = options.Options().init(argparse.ArgumentParser(description='speech enhancement')).parse_args()
print(opt)

######################################################################################################################
#                                    Set a model (check point) and a log folder                                      #
######################################################################################################################
# Get the absolute path of the current script and initialize directories for logging and models
dir_name = os.path.dirname(os.path.abspath(__file__))
print(dir_name)

log_dir = os.path.join(dir_name, 'log', opt.arch + '_' + opt.env)
utils.mkdir(log_dir)
print("Now time is : ", datetime.datetime.now().isoformat())
tboard_dir = './log/{}_{}/logs'.format(opt.arch, opt.env)  # os.path.join(log_dir, 'logs')
model_dir = './log/{}_{}/models'.format(opt.arch, opt.env)  # os.path.join(log_dir, 'models')
utils.mkdir(model_dir)  # make a dir if there is no dir (given path)
utils.mkdir(tboard_dir)

######################################################################################################################
#                                                   Model init                                                       #
######################################################################################################################
# Set the device (either CPU or GPU) and initialize the random seeds for reproducibility
DEVICE = torch.device(opt.device)

random.seed(1234)
np.random.seed(1234)
torch.manual_seed(1234)
torch.cuda.manual_seed_all(1234)

# Initialize the model and print its size
model = utils.get_arch(opt)

total_params = utils.cal_total_params(model)
print('total params   : %d (%.2f M, %.2f MBytes)\n' %
      (total_params,
       total_params / 1000000.0,
       total_params * 4.0 / 1000000.0))

# Load a pre-trained model
print('Load the pretrained model...')
chkpt = torch.load(opt.pretrain_model_path)
model.load_state_dict(chkpt['model'])

model = model.to(DEVICE)

######################################################################################################################
######################################################################################################################
#                                             Main program - test                                                    #
######################################################################################################################
######################################################################################################################
# Start testing
test_loader = create_dataloader_for_test(opt)
test_log_fp = open(model_dir + '/test_log.txt', 'a')

print('Test start...')
data_num = 0

t_all = []
o_all = []
b_all = []
p_all = []

# test
model.eval()
with torch.no_grad():
    for inputs, targets in utils.Bar(test_loader):
        # to cuda
        inputs = inputs.float().to(DEVICE)
        targets = targets.float().to(DEVICE)

        outputs = model(inputs)
        q = torch.sigmoid(outputs)

        t_all.append(targets.data.cpu().numpy())
        o_all.append(q.data.cpu().numpy())

t_all = np.concatenate(t_all, axis=0)
o_all = np.concatenate(o_all, axis=0)

# Find the threshold to classify outputs
threshold = utils.find_threshold(t_all, o_all, num_classes=opt.class_num)
print('Threshold: ', threshold)

# Convert the outputs to binary labels based on the threshold
for idx in range(len(o_all)):
    labels = utils.get_binary_outputs_np(o_all[idx], threshold)
    b_all.append(labels)

b_all = np.array(b_all)


######################################################################################################################
#                                                   Get scores                                                       #
######################################################################################################################
# Compute evaluation metrics (AUROC, AUPRC, F-Measure, and Kappa) for the predictions
num_classes = opt.class_num

s_all = np.zeros((len(b_all), num_classes), dtype=np.float64)
for i in range(len(b_all)):
    s_all[i] = [True if v == 1 else False for v in b_all[i]]

auroc, auprc, _, _ = utils.compute_auc(t_all, s_all)
f_measure, _ = utils.compute_f_measure(t_all, b_all)
kappa = utils.compute_kappa(t_all, b_all, num_classes=num_classes)

print("####################################################################")
print(opt.pretrain_model_path)
print("AUROC : {:.4}".format(auroc))
print("AUPRC : {:.4}".format(auprc))
print("F-Measure : {:.4}".format(f_measure))
print("Kappa : {:.4}".format(kappa))
print("####################################################################")

print('System has been finished.')

test_interface.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
"""
Train interface for speech enhancement!
You can just run this file.
"""
import os
import argparse
import torch
import options
import utils
import datetime
import random
import numpy as np
import time
from dataloader import create_dataloader

######################################################################################################################
#                                                  Parser init                                                       #
######################################################################################################################
# Parse command-line arguments and configurations for the training experiment
opt = options.Options().init(argparse.ArgumentParser(description='speech enhancement')).parse_args()
print(opt)

######################################################################################################################
#                                    Set a model (check point) and a log folder                                      #
######################################################################################################################
# Get the absolute path of the current script and initialize directories for logging and model storage
dir_name = os.path.dirname(os.path.abspath(__file__))  # absolute path
print(dir_name)

log_dir = os.path.join(dir_name, 'log', opt.arch + '_' + opt.env)
utils.mkdir(log_dir)
print("Now time is : ", datetime.datetime.now().isoformat())
tboard_dir = os.path.join(log_dir, 'logs')
model_dir = os.path.join(log_dir, 'models')
utils.mkdir(model_dir)  # make a dir if there is no dir (given path)
utils.mkdir(tboard_dir)

######################################################################################################################
#                                                   Model init                                                       #
######################################################################################################################
# Set up the computation device (either CPU or GPU) and initialize random seeds for reproducibility
DEVICE = torch.device(opt.device)

random.seed(1234)
np.random.seed(1234)
torch.manual_seed(1234)
torch.cuda.manual_seed(1234)
torch.cuda.manual_seed_all(1234)

# Define and initialize the model, then print the number of parameters
model = utils.get_arch(opt)

total_params = utils.cal_total_params(model)
print('total params (gen)  : %d (%.2f M, %.2f MBytes)\n' %
      (total_params,
       total_params / 1000000.0,
       total_params * 4.0 / 1000000.0))

# Define training settings including loss function, optimization method, and learning rate scheduler
trainer, validator = utils.get_train_mode(opt)
loss_calculator = utils.get_loss(opt)

optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr_initial)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=opt.decay_epoch, gamma=0.1)

# If a pre-trained model is provided, load its parameters
epoch_start_idx = 1
if opt.pretrained:
    print('Load the pretrained model...')
    chkpt = torch.load(opt.pretrain_model_path)
    model.load_state_dict(chkpt['model'])
    optimizer.load_state_dict(chkpt['optimizer'])
    epoch_start_idx = chkpt['epoch'] + 1
    print('Resuming Start Epoch: ', epoch_start_idx)

    utils.optimizer_to(optimizer, DEVICE)

model = model.to(DEVICE)

######################################################################################################################
#                                               Create Dataloader                                                    #
######################################################################################################################
# Create dataloaders for training and validation
train_loader, valid_loader = create_dataloader(opt)
print("Sizeof training set: ", train_loader.__len__(),
      ", sizeof validation set: ", valid_loader.__len__())

######################################################################################################################
######################################################################################################################
#                                             Main program - train                                                   #
######################################################################################################################
######################################################################################################################
# Initialize tensorboard writer and training log file
writer = utils.Writer(tboard_dir)
train_log_fp = open(model_dir + '/train_log.txt', 'a')
max_epoch = 0
max_auprc = 0.0

print('Train start...')
# Main training loop
for epoch in range(epoch_start_idx, opt.nepoch + 1):
    st_time = time.time()

    # Train
    train_loss = trainer(model, train_loader, loss_calculator, optimizer,
                         writer, epoch, DEVICE, opt)

    # Save the model's state after each epoch
    save_path = str(model_dir + '/chkpt_%d.pt' % epoch)
    torch.save({
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'epoch': epoch
    }, save_path)

    # Update the learning rate based on the scheduler
    scheduler.step()

    # Validate
    valid_loss, valid_auprc, valid_targets, valid_outputs = validator(model, valid_loader, loss_calculator,
                                                                      writer, epoch, DEVICE, opt)

    # Print the training and validation loss and AUPRC for the current epoch
    print('EPOCH[{}] T {:.6f} |  V {:.6f}  takes {:.3f} seconds'
          .format(epoch, train_loss, valid_loss, time.time() - st_time))
    print('AUPRC {:.6f}'.format(valid_auprc))

    # Write results to the training log
    train_log_fp.write('EPOCH[{}] T {:.6f} |  V {:.6f}  takes {:.3f} seconds'
                       .format(epoch, train_loss, valid_loss, time.time() - st_time))
    train_log_fp.write('AUPRC {:.6f}'.format(valid_auprc))


print('Training has been finished.')
train_log_fp.close()

train_interface.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
from .progress import *
from .scores import *
from .tensorboard import *
from .model_init import *
from .trainer import *
from .scores import *

utils/__init__.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
# get architecture
def get_arch(opt):
    """Retrieve the specified neural network architecture based on user's choice."""

    # Extracting relevant configuration options
    arch = opt.arch
    class_num = opt.class_num
    leads = opt.leads

    print('You choose ' + arch + '...')

    # Check if the chosen architecture is 'ResU_Dense'
    if arch == 'ResU_Dense':
        # Import the specified architecture from models
        from models import ResU_Dense
        model = ResU_Dense(nOUT=class_num)
    else:
        # If the architecture is not found, raise an exception
        raise Exception("Arch error!")

    return model


# get trainer and validator (train method)
def get_train_mode(opt):
    """Retrieve the specified training and validation methods based on the user's choice of loss type."""

    loss_type = opt.loss_type

    print('You choose ' + loss_type + 'trainer ...')

    # Check if the chosen loss type is 'base'
    if loss_type == 'base':
        # Importing the training and validation methods for the 'base' loss type
        from .trainer import base_train
        from .trainer import base_valid
        trainer = base_train
        validator = base_valid
    else:
        # If the loss type is not found, raise an exception
        raise Exception("Loss type error!")

    return trainer, validator


# get loss function
def get_loss(opt):
    """Retrieve the specified loss function based on user's choice."""

    loss_oper = opt.loss_oper
    DEVICE = opt.device
    class_num = opt.class_num

    print('You choose ' + loss_oper + ' loss function ...')

    # Check if the chosen loss operation is 'base'
    if loss_oper == 'base':
        # Use Binary Cross-Entropy with Logits as the loss function
        import torch.nn as nn
        bce_loss = nn.BCEWithLogitsLoss().to(DEVICE)
        return bce_loss
    else:
        # If the loss operation type is not found, raise an exception
        raise Exception("Loss type error!")

utils/model_init.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
import os
import time
import torch
import torch.nn as nn
from sklearn.metrics import f1_score
import numpy as np


######################################################################################################################
#                                                   for dataset                                                      #
######################################################################################################################

# Load recording file as an array.
def load_recording(recording_file, header=None, leads=None, key='val'):
    from scipy.io import loadmat
    recording = loadmat(recording_file)[key]
    if header and leads:
        recording = choose_leads(recording, header, leads)
    return recording


# Select specific leads from the recording.
def choose_leads(recording, header, leads):
    num_leads = len(leads)
    num_samples = np.shape(recording)[1]
    chosen_recording = np.zeros((num_leads, num_samples), recording.dtype)
    available_leads = get_leads(header)
    for i, lead in enumerate(leads):
        if lead in available_leads:
            j = available_leads.index(lead)
            chosen_recording[i, :] = recording[j, :]
    return chosen_recording


# Extract lead names from the header.
def get_leads(header):
    leads = list()
    for i, l in enumerate(header.split('\n')):
        entries = l.split(' ')
        if i == 0:
            num_leads = int(entries[1])
        elif i <= num_leads:
            leads.append(entries[-1])
        else:
            break
    return tuple(leads)


# Create a new directory.
def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)


# Normalize the ECG data using Z-score normalization.
def normalize_ecg(ecg_data):
    mean = np.mean(ecg_data, axis=1, keepdims=True)
    std = np.std(ecg_data, axis=1, keepdims=True)
    return (ecg_data - mean) / (std + 1e-8)  # Prevent division by zero


######################################################################################################################
#                                                  Parser init                                                       #
######################################################################################################################
# Check if a variable is a number or represents a number.
def is_number(x):
    try:
        float(x)
        return True
    except (ValueError, TypeError):
        return False


# Check if a variable is a a finite number or represents a finite number.
def is_finite_number(x):
    if is_number(x):
        return np.isfinite(float(x))
    else:
        return False


# Load a table with row and column names.
def load_table(table_file):
    # The table should have the following form:
    # ,    a,   b,   c
    # a, 1.2, 2.3, 3.4
    # b, 4.5, 5.6, 6.7
    # c, 7.8, 8.9, 9.0
    table = list()
    with open(table_file, 'r') as f:
        for i, l in enumerate(f):
            arrs = [arr.strip() for arr in l.split(',')]
            table.append(arrs)

    # Define the numbers of rows and columns and check for errors.
    num_rows = len(table) - 1
    if num_rows < 1:
        raise Exception('The table {} is empty.'.format(table_file))
    row_lengths = set(len(table[i]) - 1 for i in range(num_rows))
    if len(row_lengths) != 1:
        raise Exception('The table {} has rows with different lengths.'.format(table_file))
    num_cols = min(row_lengths)
    if num_cols < 1:
        raise Exception('The table {} is empty.'.format(table_file))

    # Find the row and column labels.
    rows = [table[0][j + 1] for j in range(num_rows)]
    cols = [table[i + 1][0] for i in range(num_cols)]

    # Find the entries of the table.
    values = np.zeros((num_rows, num_cols), dtype=np.float64)
    for i in range(num_rows):
        for j in range(num_cols):
            value = table[i + 1][j + 1]
            if is_finite_number(value):
                values[i, j] = float(value)
            else:
                values[i, j] = float('nan')

    return rows, cols, values


######################################################################################################################
#                                                   for training                                                     #
######################################################################################################################
def optimizer_to(optim, device):
    for param in optim.state.values():
        # Not sure there are any global tensors in the state dict
        if isinstance(param, torch.Tensor):
            param.data = param.data.to(device)
            if param._grad is not None:
                param._grad.data = param._grad.data.to(device)
        elif isinstance(param, dict):
            for subparam in param.values():
                if isinstance(subparam, torch.Tensor):
                    subparam.data = subparam.data.to(device)
                    if subparam._grad is not None:
                        subparam._grad.data = subparam._grad.data.to(device)


# Calculate total number of parameters in a model.
def cal_total_params(our_model):
    total_parameters = 0
    for variable in our_model.parameters():
        shape = variable.size()
        variable_parameters = 1
        for dim in shape:
            variable_parameters *= dim
        total_parameters += variable_parameters

    return total_parameters


# Display a progress bar during training/validation.
class Bar(object):
    def __init__(self, dataloader):
        if not hasattr(dataloader, 'dataset'):
            raise ValueError('Attribute `dataset` not exists in dataloder.')
        if not hasattr(dataloader, 'batch_size'):
            raise ValueError('Attribute `batch_size` not exists in dataloder.')

        self.dataloader = dataloader
        self.iterator = iter(dataloader)
        self.dataset = dataloader.dataset
        self.batch_size = dataloader.batch_size
        self._idx = 0
        self._batch_idx = 0
        self._time = []
        self._DISPLAY_LENGTH = 50

    def __len__(self):
        return len(self.dataloader)

    def __iter__(self):
        return self

    def __next__(self):
        if len(self._time) < 2:
            self._time.append(time.time())

        self._batch_idx += self.batch_size
        if self._batch_idx > len(self.dataset):
            self._batch_idx = len(self.dataset)

        try:
            batch = next(self.iterator)
            self._display()
        except StopIteration:
            raise StopIteration()

        self._idx += 1
        if self._idx >= len(self.dataloader):
            self._reset()

        return batch

    def _display(self):
        if len(self._time) > 1:
            t = (self._time[-1] - self._time[-2])
            eta = t * (len(self.dataloader) - self._idx)
        else:
            eta = 0

        rate = self._idx / len(self.dataloader)
        len_bar = int(rate * self._DISPLAY_LENGTH)
        bar = ('=' * len_bar + '>').ljust(self._DISPLAY_LENGTH, '.')
        idx = str(self._batch_idx).rjust(len(str(len(self.dataset))), ' ')

        tmpl = '\r{}/{}: [{}] - ETA {:.1f}s'.format(
            idx,
            len(self.dataset),
            bar,
            eta
        )
        print(tmpl, end='')
        if self._batch_idx == len(self.dataset):
            print()

    def _reset(self):
        self._idx = 0
        self._batch_idx = 0
        self._time = []


######################################################################################################################
#                                               for validation and test                                              #
######################################################################################################################
# Compute a modified confusion matrix for multi-class, multi-label tasks.
def compute_modified_confusion_matrix(labels, outputs):
    # Compute a binary multi-class, multi-label confusion matrix, where the rows
    # are the labels and the columns are the outputs.
    num_recordings, num_classes = np.shape(labels)
    A = np.zeros((num_classes, num_classes))

    # Iterate over all of the recordings.
    for i in range(num_recordings):
        # Calculate the number of positive labels and/or outputs.
        normalization = float(max(np.sum(np.any((labels[i, :], outputs[i, :]), axis=0)), 1))
        # Iterate over all of the classes.
        for j in range(num_classes):
            # Assign full and/or partial credit for each positive class.
            if labels[i, j]:
                for k in range(num_classes):
                    if outputs[i, k]:
                        A[j, k] += 1.0 / normalization

    return A


# Determine the threshold for each class based on F1 score.
def find_threshold(targets, scores, num_classes=26):
    classes = ['164889003', '164890007', '6374002', '426627000', '733534002',
               '713427006', '270492004', '713426002', '39732003', '445118002',
               '164947007', '251146004', '111975006', '698252002', '426783006',
               '284470004', '10370003', '365413008', '427172004', '164917005',
               '47665007', '427393009', '426177001', '427084000', '164934002',
               '59931005']

    thresholds = []
    for class_idx in range(num_classes):
        pre_targets = targets[:, class_idx]
        pre_scores = scores[:, class_idx]

        # get best threshold based on f1 score
        best_threshold, best_f1 = 0, 0
        for threshold in range(0, 101):
            binary_preds = torch.tensor(pre_scores) > threshold / 100
            f1 = f1_score(pre_targets, binary_preds)
            if f1 > best_f1:
                best_f1 = f1
                best_threshold = threshold / 100
        thresholds.append(best_threshold)
    return thresholds


# Get binary outputs based on a threshold.
def get_binary_outputs(q, threshold):
    labels = torch.zeros_like(q)
    for i in range(len(threshold)):
        pre_q = q[i]
        pre_threshold = threshold[i]

        labels[i] = pre_q >= pre_threshold
    return labels


# Get binary outputs (as numpy arrays) based on a threshold.
def get_binary_outputs_np(q, threshold):
    labels = np.zeros_like(q)
    for i in range(len(threshold)):
        pre_q = q[:, i]
        pre_threshold = threshold[i]

        labels[:, i] = pre_q >= pre_threshold
    return labels

utils/progress.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
import numpy as np
from sklearn.metrics import cohen_kappa_score


def compute_auc(labels, outputs):
    num_recordings, num_classes = np.shape(labels)

    # Compute and summarize the confusion matrices for each class across at distinct output values.
    auroc = np.zeros(num_classes)
    auprc = np.zeros(num_classes)

    for k in range(num_classes):
        # We only need to compute TPs, FPs, FNs, and TNs at distinct output values.
        thresholds = np.unique(outputs[:, k])
        thresholds = np.append(thresholds, thresholds[-1] + 1)
        thresholds = thresholds[::-1]
        num_thresholds = len(thresholds)

        # Initialize the TPs, FPs, FNs, and TNs.
        tp = np.zeros(num_thresholds)
        fp = np.zeros(num_thresholds)
        fn = np.zeros(num_thresholds)
        tn = np.zeros(num_thresholds)
        fn[0] = np.sum(labels[:, k] == 1)
        tn[0] = np.sum(labels[:, k] == 0)

        # Find the indices that result in sorted output values.
        idx = np.argsort(outputs[:, k])[::-1]

        # Compute the TPs, FPs, FNs, and TNs for class k across thresholds.
        i = 0
        for j in range(1, num_thresholds):
            # Initialize TPs, FPs, FNs, and TNs using values at previous threshold.
            tp[j] = tp[j - 1]
            fp[j] = fp[j - 1]
            fn[j] = fn[j - 1]
            tn[j] = tn[j - 1]

            # Update the TPs, FPs, FNs, and TNs at i-th output value.
            while i < num_recordings and outputs[idx[i], k] >= thresholds[j]:
                if labels[idx[i], k]:
                    tp[j] += 1
                    fn[j] -= 1
                else:
                    fp[j] += 1
                    tn[j] -= 1
                i += 1

        # Summarize the TPs, FPs, FNs, and TNs for class k.
        tpr = np.zeros(num_thresholds)
        tnr = np.zeros(num_thresholds)
        ppv = np.zeros(num_thresholds)
        for j in range(num_thresholds):
            if tp[j] + fn[j]:
                tpr[j] = float(tp[j]) / float(tp[j] + fn[j])
            else:
                tpr[j] = float('nan')
            if fp[j] + tn[j]:
                tnr[j] = float(tn[j]) / float(fp[j] + tn[j])
            else:
                tnr[j] = float('nan')
            if tp[j] + fp[j]:
                ppv[j] = float(tp[j]) / float(tp[j] + fp[j])
            else:
                ppv[j] = float('nan')

        # Compute AUROC as the area under a piecewise linear function with TPR/
        # sensitivity (x-axis) and TNR/specificity (y-axis) and AUPRC as the area
        # under a piecewise constant with TPR/recall (x-axis) and PPV/precision
        # (y-axis) for class k.
        for j in range(num_thresholds - 1):
            auroc[k] += 0.5 * (tpr[j + 1] - tpr[j]) * (tnr[j + 1] + tnr[j])
            auprc[k] += (tpr[j + 1] - tpr[j]) * ppv[j + 1]

    # Compute macro AUROC and macro AUPRC across classes.
    if np.any(np.isfinite(auroc)):
        macro_auroc = np.nanmean(auroc)
    else:
        macro_auroc = float('nan')
    if np.any(np.isfinite(auprc)):
        macro_auprc = np.nanmean(auprc)
    else:
        macro_auprc = float('nan')

    return macro_auroc, macro_auprc, auroc, auprc


def compute_kappa(labels, outputs, num_classes=26):
    avg_kappa = np.zeros(num_classes)
    for class_num in range(num_classes):
        current_t_all = labels[:, class_num]
        current_binary_outputs = outputs[:, class_num]
        kappa = cohen_kappa_score(current_t_all, current_binary_outputs)
        if kappa:
            avg_kappa[class_num] = kappa
        else:
            avg_kappa[class_num] = float('nan')
    if np.any(np.isfinite(avg_kappa)):
        avg_kappa = np.nanmean(avg_kappa)
    else:
        avg_kappa = float('nan')
    return avg_kappa


# Check if a variable is a number or represents a number.
def is_number(x):
    try:
        float(x)
        return True
    except (ValueError, TypeError):
        return False


def compute_confusion_matrices(labels, outputs, normalize=False):
    # Compute a binary confusion matrix for each class k:
    #
    #     [TN_k FN_k]
    #     [FP_k TP_k]
    #
    # If the normalize variable is set to true, then normalize the contributions
    # to the confusion matrix by the number of labels per recording.
    num_recordings, num_classes = np.shape(labels)

    if not normalize:
        A = np.zeros((num_classes, 2, 2))
        for i in range(num_recordings):
            for j in range(num_classes):
                if labels[i, j] == 1 and outputs[i, j] == 1:  # TP
                    A[j, 1, 1] += 1
                elif labels[i, j] == 0 and outputs[i, j] == 1:  # FP
                    A[j, 1, 0] += 1
                elif labels[i, j] == 1 and outputs[i, j] == 0:  # FN
                    A[j, 0, 1] += 1
                elif labels[i, j] == 0 and outputs[i, j] == 0:  # TN
                    A[j, 0, 0] += 1
                else:  # This condition should not happen.
                    raise ValueError('Error in computing the confusion matrix.')
    else:
        A = np.zeros((num_classes, 2, 2))
        for i in range(num_recordings):
            normalization = float(max(np.sum(labels[i, :]), 1))
            for j in range(num_classes):
                if labels[i, j] == 1 and outputs[i, j] == 1:  # TP
                    A[j, 1, 1] += 1.0 / normalization
                elif labels[i, j] == 0 and outputs[i, j] == 1:  # FP
                    A[j, 1, 0] += 1.0 / normalization
                elif labels[i, j] == 1 and outputs[i, j] == 0:  # FN
                    A[j, 0, 1] += 1.0 / normalization
                elif labels[i, j] == 0 and outputs[i, j] == 0:  # TN
                    A[j, 0, 0] += 1.0 / normalization
                else:  # This condition should not happen.
                    raise ValueError('Error in computing the confusion matrix.')

    return A


def compute_f_measure(labels, outputs):
    num_recordings, num_classes = np.shape(labels)

    A = compute_confusion_matrices(labels, outputs)

    f_measure = np.zeros(num_classes)
    for k in range(num_classes):
        tp, fp, fn, tn = A[k, 1, 1], A[k, 1, 0], A[k, 0, 1], A[k, 0, 0]
        if 2 * tp + fp + fn:
            f_measure[k] = float(2 * tp) / float(2 * tp + fp + fn)
        else:
            f_measure[k] = float('nan')

    if np.any(np.isfinite(f_measure)):
        macro_f_measure = np.nanmean(f_measure)
    else:
        macro_f_measure = float('nan')

    return macro_f_measure, f_measure


utils/scores.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
from tensorboardX import SummaryWriter


# Define a custom writer class that extends SummaryWriter to log training/validation metrics.
class Writer(SummaryWriter):
    def __init__(self, logdir):
        super(Writer, self).__init__(logdir)

    # Method to log training loss.
    def log_train_loss(self, loss_type, train_loss, step):
        self.add_scalar('train_{}_loss'.format(loss_type), train_loss, step)

    # Method to log validation loss.
    def log_valid_loss(self, loss_type, valid_loss, step):
        self.add_scalar('valid_{}_loss'.format(loss_type), valid_loss, step)

    # Method to log other performance metrics (e.g., accuracy, F1-score).
    def log_score(self, metrics_name, metrics, step):
        # Add a scalar value to the writer with the given metric name.
        self.add_scalar(metrics_name, metrics, step)




utils/tensorboard.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification
import torch
from .progress import Bar, np
from sklearn.metrics import average_precision_score


######################################################################################################################
#                                               train loss function                                                  #
######################################################################################################################
# Define the training function for the model.
def base_train(model, train_loader, loss_calculator, optimizer, writer,
               EPOCH, DEVICE, opt):
    # Initialization of variables to track training loss and number of batches.
    train_loss = 0
    batch_num = 0

    # Train
    model.train()

    # Iterate through batches of data from the train_loader.
    for inputs, targets in Bar(train_loader):
        batch_num += 1

        # to cuda
        inputs = inputs.float().to(DEVICE)
        targets = targets.float().to(DEVICE)

        outputs = model(inputs)

        loss = loss_calculator(input=outputs, target=targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()  # Update model parameters based on gradients.

        train_loss += loss.item()
    train_loss /= batch_num

    # Log the average training loss to tensorboard.
    writer.log_train_loss('total', train_loss, EPOCH)

    return train_loss


######################################################################################################################
#                                               valid loss function                                                  #
######################################################################################################################
# Define the validation function for the model.
def base_valid(model, valid_loader, loss_calculator, writer, EPOCH, DEVICE, opt):
    # Initialization of variables to track validation loss and number of batches.
    valid_loss = 0
    batch_num = 0

    t_all = []
    o_all = []

    # Validation
    model.eval()

    with torch.no_grad():
        for inputs, targets in Bar(valid_loader):
            batch_num += 1

            # to cuda
            inputs = inputs.float().to(DEVICE)
            targets = targets.float().to(DEVICE)

            outputs = model(inputs)
            p = torch.sigmoid(outputs)
            loss = loss_calculator(input=outputs, target=targets)

            valid_loss += loss

            t_all.append(targets.data.cpu().numpy())
            o_all.append(p.data.cpu().numpy())
        valid_loss /= batch_num

    t_all = np.concatenate(t_all, axis=0)
    o_all = np.concatenate(o_all, axis=0)

    # Compute the Area Under the Precision-Recall Curve (AUPRC).
    valid_auprc = average_precision_score(y_true=t_all, y_score=o_all)
    # Log the average validation loss to tensorboard.
    writer.log_valid_loss('total', valid_loss, EPOCH)
    writer.log_score('AUPRC', valid_auprc, EPOCH)

    return valid_loss, valid_auprc, t_all, o_all


utils/trainer.py
seorim0/Multi-label-12-lead-ECG-abnormality-classification