Биорадар: Часть 2. Обучение ИНС реконструкции ЭКГ

using Flux
using Flux: @functor 

using Flux
using Flux: unsqueeze
using Flux: Conv, ConvTranspose, BatchNorm, MaxPool, leakyrelu, cat


  
struct SEBlock
    fc1::Dense
    fc2::Dense
  end
  Flux.@functor SEBlock
  
function SEBlock(channels::Int; reduction::Int=8)

fc1 = Dense(channels, div(channels, reduction), relu)
fc2 = Dense(div(channels, reduction), channels, σ)
SEBlock(fc1, fc2)
end
  
function (m::SEBlock)(x)
# x: (L, C, N)
s = mean(x, dims=1)       
s = dropdims(s, dims=1) 
s = m.fc2(m.fc1(s))       
s = Flux.unsqueeze(s, 1)  
return x .* s            
end
  

struct ConvBlock
    block
    residual
end
Flux.@layer ConvBlock
  
function ConvBlock(cin::Int, cout::Int; k=5, d=2, reduction=8)
    pad = ((k - 1) ÷ 2) * d
    block = Chain(
        Conv((7,), cin => cout; pad=(pad,), dilation=(d,)),
        x -> leakyrelu.(x, 0.1),
        Dropout(0.25),
        Conv((5,), cout => cout; pad=(pad,), dilation=(d,)),
        x -> leakyrelu.(x, 0.1),
        Dropout(0.15),
        Conv((3,), cout => cout; pad=(pad,), dilation=(d,)),
        x -> leakyrelu.(x, 0.1),
        Dropout(0.1),
        SEBlock(cout; reduction=reduction),
    )
    res = cin == cout ? identity : Conv((1,), cin => cout)
    ConvBlock(block, res)
end



(m::ConvBlock)(x) = leakyrelu.(m.block(x) .+ m.residual(x), 0.1)



struct UpBlock
    up::ConvTranspose
    conv::ConvBlock
end


Flux.@layer UpBlock

function UpBlock(cin::Integer, cout::Integer)
    up   = ConvTranspose((4,), cin => cout; stride=2, pad=1)
    conv = ConvBlock(cin, cout)
    UpBlock(up, conv)
end

function (u::UpBlock)(x, skip)
    x = u.up(x)                       


    if size(x,1) > size(skip,1)             
        x = x[1:size(skip,1), :, :]         
    end


    x_cat = cat(x, skip; dims = 2)         


    return u.conv(x_cat)
end


struct Radar2ECG_UNet_SE
    c1; p1; c2; p2; c3; p3; c4; p4; c5; p5; c6
    u5; u4; u3; u2; u1
    out
end
Flux.@layer Radar2ECG_UNet_SE

function Radar2ECG_UNet_SE(base=32)
    c1 = ConvBlock(1,    base)
    p1 = MaxPool((2,); pad=(0,), stride=(2,))
    c2 = ConvBlock(base, base*2)
    p2 = MaxPool((2,); pad=(0,), stride=(2,))
    c3 = ConvBlock(base*2, base*4; d=2)
    p3 = MaxPool((2,); pad=(0,), stride=(2,))
    c4 = ConvBlock(base*4, base*8; d=4)
    p4 = MaxPool((2,); pad=(0,), stride=(2,))
    c5 = ConvBlock(base*8, base*16; d=8)
    p5 = MaxPool((2,); pad=(0,), stride=(2,))
    c6 = ConvBlock(base*16, base*32; d=16)

    u5 = UpBlock(base*32, base*16)
    u4 = UpBlock(base*16, base*8)
    u3 = UpBlock(base*8,  base*4)
    u2 = UpBlock(base*4,  base*2)
    u1 = UpBlock(base*2,  base)
    out = Conv((1,), base => 1)
    Radar2ECG_UNet_SE(c1,p1,c2,p2,c3,p3,c4,p4,c5,p5,c6,u5,u4,u3,u2,u1,out)
end

function (m::Radar2ECG_UNet_SE)(x)
    s1 = m.c1(x)
    s2 = m.c2(m.p1(s1))
    s3 = m.c3(m.p2(s2))
    s4 = m.c4(m.p3(s3))
    s5 = m.c5(m.p4(s4))
    b  = m.c6(m.p5(s5))

    d5 = m.u5(b,  s5)
    d4 = m.u4(d5, s4)
    d3 = m.u3(d4, s3)
    d2 = m.u2(d3, s2)
    d1 = m.u1(d2, s1)

    return m.out(d1)
end