Документация Engee
Notebook

Мониторинг ЭКГ с использованием радара непрерывного излучения и глубокого обучения (Часть 2)

Для мониторирования сердца используют ЭКГ: на кожу ставят электроды и фиксируют активность сердца. Но из-за кратковременности записи, дискомфорта от наклеек и риска артефактов при движении такой метод не всегда удобен.

Радарные системы предлагают бесконтактный мониторинг: они улавливают микродвижения грудной клетки от сердцебиения и дыхания, не ограничивая пациента в движениях и не требуя контакта с кожей. Используя нейросети можно восстанавливать электрокардиограмму по механическим колебаниям радарного сигнала, превращая движение грудной клетки в привычные P-QRS-T-волны.

В рамках демо-примера мы сфокусируемся на биорадарах для наблюдения за сердцем и применении ИИ-алгоритмов для трансформации радарного сигнала в эквивалент ЭКГ.

Для ознакомления с 1 частью проекта, посвященной анализу и сравнивению признаков ЭКГ сигнала и отраженного эхо-радара, перейдите по ссылке.

Подготовка к работе

Импортируем необходимые библиотеки

Импортируем необходимые для работы библиотеки из Julia и Python ( skimage, scipy, neurokit2), а также импортируем в скрипт файл utils.jl, который содержит все необходимые функции

In [ ]: 
include("$(@__DIR__)/install_packages.jl");

Для работы скрипта необходимо установить несколько питоновских пакетов, функции которыъ вызываются через PyCall. Для установки необходимо перейти в командную строку и нажать символ ";". Далее откроется shell, куда нужно прописать:

pip install -r requirements.txt

важно находится в рабочей директории.

In [ ]: 
using Glob
using DSP
using CUDA
using Flux 
using Statistics
using cuDNN
using BSON: @save, @load
In [ ]: 
using Measures

Поскольку используемые функции имеют достаточно большое количество строк кода, они были вынесены в отдельные файлы .jl, которые импортируются ниже

Файл model.jl реализует структуру используемой модели. Файл dataset.jl используется для реализации логики модуля набора данных, prepare_data.jl содержит функции, которые выполняют подготовку данных для обучения, ну а файл utils.jl содержит все остальные необходимые для работы проекта функции

In [ ]: 
Dirpath = "$(@__DIR__)"
include(joinpath(Dirpath, "model.jl"))
include(joinpath(Dirpath, "dataset.jl"))
include(joinpath(Dirpath, "prepare_data.jl"))
include(joinpath(Dirpath, "utils.jl"));

Подготовка данных

В работе используется открытый набор данных, представленный Schellenberger et al. (Scientific Data, 2020)​. Этот датасет содержит ~24 часа синхронных записей радара и эталонных датчиков (ЭКГ, импедансная кардиография, непрерывное АД) от 30 здоровых добровольцев​. Из датасета был извлечен одно изменерие, над которым проводилось исследование

Создадим папку, куда будут сохраняться результаты обработки сигналов.

In [ ]: 
prepareRadarData = joinpath(Dirpath, "prepareRadarData")
isdir(prepareRadarData) || (mkdir(prepareRadarData); println("Папка prepareRadarData создана"));

Далее выполним все те преобразования над сырыми ЭКГ и радарными данными, что были проделаны в демо-примере анализу этих данных, а именно:

  1. Реконструкция эллипса
  2. Извлечение жизненноважных сиглавнов - дыхание, пульс
  3. Вычисление ускорения для грудной клетки и последующая фильтрация
  4. Фильтрация ЭКГ

Реализация этой функции лежит в файле prepare_data.jl

In [ ]: 
root = joinpath(Dirpath, "data")

pat = r"^GDN\d{4}_.+\.mat$"
order_filter = 6
files = [p for p in glob("**/*.mat", root) if occursin(pat, basename(p))]
save_path = joinpath(Dirpath, "prepareRadarData")
isdir(save_path) || mkdir(save_path)
for file in files
    prepareData(file, save_path, order_filter)
end

C помощью блока выше мы обработали данные и сохранили их в отдельную папку

Создание загрузчика данных

Далее мы создадим даталоадер, который будет батчами загружать данные в модель. В блоке ниже мы инициализуем датасет, поделим данные на тренировочное, валидационное и тестовое множество и создадим загрузчики данных

Логика датасета прописана в файле dataset.jl

Датасет нарезает ЭКГ и радарный сигнал на окна с наложением. Нарезка позволяет снизить нагрузку на ИНС изза использования длинных последовательностей, а использование наложения позволяет учитывать временные зависимости

In [ ]: 
ds = RadarECG("prepareRadarData");
N = length(ds)
train_len = floor(Int, 0.7 * N)
val_len   = floor(Int, 0.15 * N)
test_len  = N - train_len - val_len


train_idx = 1:train_len
val_idx   = (train_len+1):(train_len+val_len)
test_idx  = (train_len+val_len+1):N

X_train = [ ds[i][1] for i in train_idx ]  
Y_train = [ ds[i][2] for i in train_idx ]

X_val   = [ ds[i][1] for i in val_idx   ]
Y_val   = [ ds[i][2] for i in val_idx   ]

X_test  = [ ds[i][1] for i in test_idx  ]
Y_test  = [ ds[i][2] for i in test_idx  ]


batch_size = 64
train_loader = Flux.DataLoader((X_train, Y_train);
    batchsize = batch_size,
    shuffle   = false,
    partial      = true,
    parallel = true,
    collate   = mycollate
)
val_loader = Flux.DataLoader((X_val, Y_val);   batchsize=batch_size, shuffle=true, parallel = true, collate   = mycollate)
test_loader= Flux.DataLoader((X_test, Y_test);  batchsize=batch_size, shuffle=true, parallel = true, collate   = mycollate);

Посмотрим, как выглядят ЭКГ и радарный сигналы, нарезаный на окна с наложением

In [ ]: 
radar_batch, ecg_batch = first(train_loader)  
radar_batch = cpu(radar_batch)
ecg_batch = cpu(ecg_batch)
win, B = size(radar_batch)

p = plot(size=(500,500), layout = (2,1), margin=15mm)

for (arr, name) in zip((radar_batch, ecg_batch), ("radar", "ecg"))
    target = name == "radar" ? 1 : 2
    for j in 0:3
        x = (1:win) .+ j*ds.step
        y = arr[:, j+1]
        if target == "radar"
            plot!(p, x, y; subplot=target, lw=2, showlegend=false, alpha = 0.7, xlabel="Отсчет", ylabel="Ускорение, м/с^2")
        else
            plot!(p, x, y; subplot=target, lw=2, showlegend=false, alpha = 0.7, xlabel="Отсчет", ylabel="Амплитуда, мВ")
        end
    end

end
title!(p, "Радар", subplot=1)
title!(p, "ЭКГ",   subplot=2)



display(p)

После опредилим устройство, на котором будет обучаться модель. Если есть доступ к ГПУ, оубчение будет ускорено на несколько порядков.

In [ ]: 
device = CUDA.functional() ? gpu : cpu;
device == gpu ? train_loader = gpu.(train_loader) : nothing
device == gpu ? val_loader = gpu.(val_loader) : nothing
device == gpu ? test_loader = gpu.(test_loader) : nothing;

Архитектура модели представлена в файле model.jl. Инициализируем модель и переведем ее на используемое устройство

Инициализация модели, оптимизатора, функции потерь и других параметров

In [ ]: 
model = Radar2ECG_UNet_SE(32);
device == gpu ? model = gpu(model) : nothing;

Определим некоторые параметры: скорость обучения, количество эпох, имя модели - оно нужно для автоматического создания директории в рабочем пространстве, куда будут сохраняться в процессе обучения веса модели и графики различных метрик

Также инициализируем оптимизатор, функцию потерь, которая реализована в файле utils.jl, ну и воспользуемся вспомогательной функцией get_training_dir, которая создаст необходимые директории для сохранения результатов обучения

In [ ]: 
fs_ecg = 2000;
In [ ]: 
learning_rate = 1e-4
num_epochs = 100
model_name = "trainedModel"

lossFunction = loss;
# opt = Flux.Adam(learning_rate)
opt = Flux.Optimiser(WeightDecay(1e-6), Adam(learning_rate))

epoch_number = 0
train_dir, name_new_dir = get_training_dir(Dirpath, model_name)
print(train_dir, " ", name_new_dir)

/user/Диплом/AI/model/trainedModel/train1 train1

Определим списки, где будут храниться данные вычисленных метрик. В данном примере вычисляем MAE, RMSE, R2, Accuracy по уровню tau

Также инициализируем переменную params, которая хранит некоторые параметры цикла обучения

In [ ]: 
tau = 0.05
best_val_loss = Inf
train_losses = Float64[]   
valid_losses = Float64[]  
MAEList      = Float64[]
RMSEList     = Float64[]
R2List       = Float64[]
AccList      = Float64[]
params = Dict("epochs"=> num_epochs, "batch_size"=>batch_size, "learning_rate"=> learning_rate, "tau"=>tau, "model"=>"Radar2ECG_UNet", "info"=>"Train model without scheduler." );

Обучения модели

Далее запустим цикл обучения модели. Функции train, validate находятся в файле utils.jl. После окончания обучения результаты сохраняются в созданную папку

In [ ]: 
no_improve_epochs = 0
best_model    = nothing

try
    for epoch in 1:num_epochs
        println("-"^50 * "\n")
        println("EPOCH $(epoch):")

        opt, train_loss = train!(
            model, train_loader, opt,
            lossFunction, device, 1, num_epochs
        )

        val_loss, metrics = validate(model, val_loader, lossFunction, gpu)
        @show "R2" metrics.R2

        if val_loss < best_val_loss
            best_val_loss = val_loss
            no_improve_epochs = 0
        else
            no_improve_epochs += 1
            if no_improve_epochs >= 10
                old_lr = opt[2].eta
                opt[2].eta *= 0.9
                @info "LR reduced: $old_lr → $(opt[2].eta) at epoch $epoch"
                no_improve_epochs = 0
            end
        end

        println("Epoch $epoch/$num_epochs | train $(round(train_loss, digits=4)) | val $(round(val_loss, digits=4))")
        println("Current LR = $(opt[2].eta)")

        push!(train_losses, train_loss)
        push!(valid_losses, val_loss)
        push!(MAEList,  metrics.MAE)
        push!(RMSEList, metrics.RMSE)
        push!(R2List,   metrics.R2)
        push!(AccList,  metrics[Symbol("Acc@0.05")])


        if val_loss == best_val_loss
            best_model = deepcopy(model)
            @info "Валидационная потеря упала" val_loss=best_val_loss
        end

    end

catch e
    if e isa InterruptException
        println("Training interrupted by user.")
    else
        rethrow(e)
    end

finally
    best_model = cpu(best_model)
    output_dir_models = joinpath(train_dir, "models_bson")
    isdir(output_dir_models) || mkdir(output_dir_models)
    best_path = joinpath(output_dir_models, "model_best.bson")
    @info "Saving best model (val_loss=$(best_val_loss)) to $best_path"
    @save best_path best_model
    save_plots!(MAEList, RMSEList, R2List, AccList,
                train_losses, valid_losses;
                model_dir = "model",
                run_name  = name_new_dir,
                plot_name = model_name)

    params_path = joinpath(train_dir, "params.json")
    open(params_path, "w") do io
        JSON.print(io, params)
    end
    println("All done. Artifacts saved in $train_dir")
end

Инференс модели на тестовых данных

Проверим, как работает модель на тестовых данных

In [ ]: 
radar, ecg = test_loader[2]


# radar, ecg = batch


radar_ = reshape(radar, size(radar,1), 1, size(radar,2)) |> device
ecg_gpu = ecg |> device

pred_gpu = model(radar_)  


radar_cpu = Array(radar)                        
pred_cpu  = Array(dropdims(pred_gpu; dims=2))  
ecg_cpu   = Array(ecg_gpu)                     

B, L = size(radar_cpu)
rows, cols = 4, 2
nplots      = min(rows*cols, L)      # вдруг каналов < 8

plt = plot(layout       = (rows, cols),   # 4×2
           size         = (1500, 1500),
           margin       = 15mm,
           legend       = :outertopright,
           legendfontsize = 8)

for i in 1:nplots
    # подписи в легенде один-раз, чтобы не плодить дубликаты
    lbls = i == 1 ? ("Радар", "Восстановленная ECG", "Эталонная ECG") : ("", "", "")

    plot!(plt, radar_cpu[:, i + 2], lw = 2, alpha = 0.7,
          label = lbls[1], subplot = i)

    plot!(plt, pred_cpu[:,  i + 2], lw = 2,
          label = lbls[2], subplot = i)

    plot!(plt, ecg_cpu[:,   i + 2], lw = 2, linestyle = :dash,
          label = lbls[3], subplot = i)

    title!(plt, "Пример $i", subplot = i)
    xlabel!(plt, "Время [s]", subplot = i)
    ylabel!(plt, "Амплитуда, мВ", subplot = i)
end

display(plt)

Оценка обученной модели

Отобразим графики ошибок на трейне и на валидации

In [ ]: 
plot(train_losses, label="train loss", lw=3)
plot!(valid_losses, label="valid loss", lw = 3)
xlabel!("Эпоха")
ylabel!("Ошибка")
Out[0]:

Визуализируем метрики, накопленные во время обучения - R^2, RMSE, MAE, Accuracy

In [ ]: 
plt = plot(layout = (1,4), size = (900,300),
           legend = :outertopright,   
           legendfontsize = 8,
           margin = 5Plots.mm)

metricsList = [RMSEList, MAEList, R2List, AccList]
metric_names = ["RMSE", "MAE", "R²", "Accuracy"]

plt = plot(
    layout       = (1, 4),
    size         = (1000, 300),
    margin       = 5Plots.mm,
    legend       = false,         
)

for (i, name) in enumerate(metric_names)
    plot!(
        plt, metricsList[i], subplot = i,lw = 2,alpha   = 0.8,label   = "")
    title!(plt, name,      subplot = i)
    xlabel!(plt, "Epoch",    subplot = i)
    ylabel!(plt, name,       subplot = i)
end

display(plt)

Заключение

В работе была обучена кастомная U-net подобная нейросеть на предобработанных радарных данных для восстановления ЭКГ. Модель неплохо справляется со своей задачей показывая хорошие цифры на вычисленных метриках