run.py

import torch
import os
import gc
import random

from transformer import ECGDETR
from dataset import ECGDataset
from qrs import QRS
from utils import prepare_data, create_dataset, load_ecg


if __name__ == '__main__':
    ecg, annotation = load_ecg('mitbit/223')
    d = ECGDataset(ecg[:12000])
    m = ECGDETR(QRS.METHODS[1:])
    pred = m.predict(ecg[:12000])
    m.dataset.plot_ecg(30, 32)


    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    path = 'mitbit'

    numbers = ['200', '208', '233', '100', '101', '103', '113', '115', '109', '111', '214', '118', '124', '212', '231',
               '209', '220', '222', '223', '232', '106', '119']  # 207

    #
    import pickle
    # dataset = create_dataset(numbers, path)
    # with open('dataset_new.pkl', 'wb') as f:
    #     pickle.dump(dataset, f)

    with open('dataset_new.pkl', 'rb') as f:
        dataset = pickle.load(f)

    batch_labels = 2000
    max_items = 2200
    methods = QRS.METHODS[1:]  # ['dwt', 'stft', 'spectrogram', 'cwt', 'fft', 'welch']

    epochs = 1000

    sp = []
    random.shuffle(dataset.split_ecg)
    for label in QRS.LABELS:
        sp.extend(dataset[label][:max_items])
    random.shuffle(sp)
    dataset.split_ecg = sp
    for method in methods:
        if not os.path.exists(f'models/{method}'):
            os.makedirs(f'models/{method}')
        model = ECGDETR([method])
        for i in range(0, max_items, batch_labels):
            if i < max_items * .9:
                train_images, train_lbls = prepare_data(i, [method], dataset, batch_labels=batch_labels)
                model.train_model(train_images.reshape(-1, 1, *QRS.IMG_SIZE),
                                  train_lbls.reshape(-1, len(QRS.LABELS)),
                                  epochs=epochs, early_stopping_patience=epochs // 10)
                print(f"Batch: {i} done")
            else:
                train_images, train_lbls = prepare_data(i, [method], dataset, batch_labels=max_items % batch_labels)
                model.evaluate(train_images.reshape(-1, 1, *QRS.IMG_SIZE),
                               train_lbls.reshape(-1, len(QRS.LABELS)))
            del train_images, train_lbls
            gc.collect()

        # model.save(f'models')
        model.plot_training()