Мониторинг ЭКГ с использованием радара непрерывного излучения и глубокого обучения (Часть 2)
Для мониторирования сердца используют ЭКГ: на кожу ставят электроды и фиксируют активность сердца. Но из-за кратковременности записи, дискомфорта от наклеек и риска артефактов при движении такой метод не всегда удобен.
Радарные системы предлагают бесконтактный мониторинг: они улавливают микродвижения грудной клетки от сердцебиения и дыхания, не ограничивая пациента в движениях и не требуя контакта с кожей. Используя нейросети можно восстанавливать электрокардиограмму по механическим колебаниям радарного сигнала, превращая движение грудной клетки в привычные P-QRS-T-волны.
В рамках демо-примера мы сфокусируемся на биорадарах для наблюдения за сердцем и применении ИИ-алгоритмов для трансформации радарного сигнала в эквивалент ЭКГ.
Для ознакомления с 1 частью проекта, посвященной анализу и сравнивению признаков ЭКГ сигнала и отраженного эхо-радара, перейдите по ссылке.
Подготовка к работе
Импортируем необходимые библиотеки
Импортируем необходимые для работы библиотеки из Julia и Python ( skimage, scipy, neurokit2), а также импортируем в скрипт файл utils.jl, который содержит все необходимые функции
include("$(@__DIR__)/install_packages.jl");
Для работы скрипта необходимо установить несколько питоновских пакетов, функции которыъ вызываются через PyCall. Для установки необходимо перейти в командную строку и нажать символ ";". Далее откроется shell, куда нужно прописать:
pip install -r requirements.txt
важно находится в рабочей директории.
using Glob
using DSP
using CUDA
using Flux
using Statistics
using cuDNN
using BSON: @save, @load
using Measures
Поскольку используемые функции имеют достаточно большое количество строк кода, они были вынесены в отдельные файлы .jl, которые импортируются ниже
Файл model.jl реализует структуру используемой модели. Файл dataset.jl используется для реализации логики модуля набора данных, prepare_data.jl содержит функции, которые выполняют подготовку данных для обучения, ну а файл utils.jl содержит все остальные необходимые для работы проекта функции
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 здоровых добровольцев. Из датасета был извлечен одно изменерие, над которым проводилось исследование
Создадим папку, куда будут сохраняться результаты обработки сигналов.
prepareRadarData = joinpath(Dirpath, "prepareRadarData")
isdir(prepareRadarData) || (mkdir(prepareRadarData); println("Папка prepareRadarData создана"));
Далее выполним все те преобразования над сырыми ЭКГ и радарными данными, что были проделаны в демо-примере анализу этих данных, а именно:
- Реконструкция эллипса
- Извлечение жизненноважных сиглавнов - дыхание, пульс
- Вычисление ускорения для грудной клетки и последующая фильтрация
- Фильтрация ЭКГ
Реализация этой функции лежит в файле prepare_data.jl
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
Датасет нарезает ЭКГ и радарный сигнал на окна с наложением. Нарезка позволяет снизить нагрузку на ИНС изза использования длинных последовательностей, а использование наложения позволяет учитывать временные зависимости
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);
Посмотрим, как выглядят ЭКГ и радарный сигналы, нарезаный на окна с наложением
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)
После опредилим устройство, на котором будет обучаться модель. Если есть доступ к ГПУ, оубчение будет ускорено на несколько порядков.
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. Инициализируем модель и переведем ее на используемое устройство
Инициализация модели, оптимизатора, функции потерь и других параметров
model = Radar2ECG_UNet_SE(32);
device == gpu ? model = gpu(model) : nothing;
Определим некоторые параметры: скорость обучения, количество эпох, имя модели - оно нужно для автоматического создания директории в рабочем пространстве, куда будут сохраняться в процессе обучения веса модели и графики различных метрик
Также инициализируем оптимизатор, функцию потерь, которая реализована в файле utils.jl, ну и воспользуемся вспомогательной функцией get_training_dir, которая создаст необходимые директории для сохранения результатов обучения
fs_ecg = 2000;
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)
Определим списки, где будут храниться данные вычисленных метрик. В данном примере вычисляем MAE, RMSE, R2, Accuracy по уровню tau
Также инициализируем переменную params, которая хранит некоторые параметры цикла обучения
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. После окончания обучения результаты сохраняются в созданную папку
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
Инференс модели на тестовых данных
Проверим, как работает модель на тестовых данных
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)
Оценка обученной модели
Отобразим графики ошибок на трейне и на валидации
plot(train_losses, label="train loss", lw=3)
plot!(valid_losses, label="valid loss", lw = 3)
xlabel!("Эпоха")
ylabel!("Ошибка")
Визуализируем метрики, накопленные во время обучения - R^2, RMSE, MAE, Accuracy
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 подобная нейросеть на предобработанных радарных данных для восстановления ЭКГ. Модель неплохо справляется со своей задачей показывая хорошие цифры на вычисленных метриках