import torch import torch.nn as nn import torch.nn.functional as F class DoubleConvBlock(nn.Module): """ (2D conv => BN => LeakyReLU) * 2 """ def __init__(self, in_ch, out_ch, k_size, pad, dil): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), nn.BatchNorm2d(out_ch), nn.LeakyReLU(inplace=True), nn.Conv2d(out_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), nn.BatchNorm2d(out_ch), nn.LeakyReLU(inplace=True) ) def forward(self, x): x = self.block(x) return x class Double3DConvBlock(nn.Module): """ (3D conv => BN => LeakyReLU) * 2 """ def __init__(self, in_ch, out_ch, k_size, pad, dil): super().__init__() self.block = nn.Sequential( nn.Conv3d(in_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), nn.BatchNorm3d(out_ch), nn.LeakyReLU(inplace=True), nn.Conv3d(out_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), nn.BatchNorm3d(out_ch), nn.LeakyReLU(inplace=True) ) def forward(self, x): x = self.block(x) return x class ConvBlock(nn.Module): """ (2D conv => BN => LeakyReLU) """ def __init__(self, in_ch, out_ch, k_size, pad, dil): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size=k_size, padding=pad, dilation=dil), nn.BatchNorm2d(out_ch), nn.LeakyReLU(inplace=True) ) def forward(self, x): x = self.block(x) return x # reda peak_conv zlw @20220624 # Peak Conv with fixed guard banwidth (gb) and dense reference receptive field (RRF) class PeakConv2D(nn.Module): def __init__(self, in_ch, out_ch, bias, refer_band, guard_band_dop, guard_band, signal_type): """ Args: refer_band: a single integer value or integer vector of two values. guard_band_dop: guard bandwidth for doppler input guard_band: guard banwidth for non-doppler input signal_type in ('range_doppler', 'angle_doppler', 'range_angle'). k_size, padding and stride are computed using the input fixed guard bandwidth input: x in the shape of BxCxHxW """ super(PeakConv2D, self).__init__() self.signal_type = signal_type if self.signal_type in ('range_doppler', 'angle_doppler'): # guard_band should be bi-lateral self.guard_band = (guard_band_dop, guard_band_dop) # padding order: l-r-t-b self.padding = (self.guard_band[1] + refer_band, self.guard_band[1] + refer_band, self.guard_band[0] + refer_band, self.guard_band[0] + refer_band) else: self.guard_band = (guard_band, guard_band) # padding order: l-r-t-b self.padding = (self.guard_band[1] + refer_band, self.guard_band[1] + refer_band, self.guard_band[0] + refer_band, self.guard_band[0] + refer_band) self.refer_band = refer_band self.kernel_size = (self.guard_band[0] + self.guard_band[1] + self.refer_band * 2), 4 self.stride = self.kernel_size # padding order: l-r-t-b self.zero_padding = nn.ZeroPad2d(self.padding) self.bias = bias self.pc_strd = 1 # conv block for peak receptive field (PRF) in x # take PRF as input and output the center point output of PRF; # kernel channel = in_ch*2: depth_concat(p_c, p_r) # since each p_c is expanded as a tensor of kernel scale RF, the stride = kernel_size self.peak_conv = nn.Conv2d(in_ch, out_ch, kernel_size=self.kernel_size, padding=0, stride=self.stride, bias=self.bias) #self.peak_conv.register_backward_hook(self._set_lr) @staticmethod # set learning rate for peak_conv def _set_lr(module, grad_input, grad_output): grad_input = (grad_input[i] * 0.1 for i in range(len(grad_input))) grad_output = (grad_output[i] * 0.1 for i in range(len(grad_output))) def forward(self, x): h, w = x.shape[2], x.shape[3] rb = self.refer_band gb = self.guard_band k_pk = self.kernel_size N = k_pk[0] * k_pk[1] dtype = x.data.type() b = x.size(0) x = self.zero_padding(x) # let p_r denote the reference point positions # p_r shape: Bx2NxHxW p_r = self._get_p_r(b, h, w, N, dtype) # sample x_r (b, c, h, w, N) using p_r x_r = self._sample_x(x, p_r, N) # sample x_c (b, c, h, w, N) using p_c p_c = self._get_p_c(b, h, w, N, dtype) x_c = self._sample_x(x, p_c, N) # getting x_prf for peak_conv # (b, c, h, w, N) ==> (b, 2c, h, w, N) # x_prf = torch.cat((x_c, x_r), 1) x_prf = x_c - x_r # (b, 2c, h*k_pk[0], w*k_pk[1]) x_prf = self._reshape_x_prf(k_pk[0], k_pk[1], x_prf) # out: (b, c, h, w) out = self.peak_conv(x_prf) return out # getting p_c (the center point coords) from the padded grid of input x def _get_p_c(self, b, h, w, N, dtype): # generating pc_grid # for UserWarning in meshgrid, indexing='ij' @20240627 p_c_x, p_c_y = torch.meshgrid( torch.arange(self.padding[2], h * self.pc_strd + self.padding[2], self.pc_strd), torch.arange(self.padding[0], w * self.pc_strd + self.padding[0], self.pc_strd), indexing='ij') p_c_x = torch.flatten(p_c_x).view(1, 1, h, w).repeat(1, N, 1, 1) p_c_y = torch.flatten(p_c_y).view(1, 1, h, w).repeat(1, N, 1, 1) # p_c: 1x2NxHxW p_c = torch.cat([p_c_x, p_c_y], 1).type(dtype) # (B, 2N, h, w) p_c = p_c.repeat(b, 1, 1, 1) return p_c # generating peak receptive field grid def _gen_prf_grid(self, rb, gb, N, dtype): h_t = -(rb + gb[0]) h_d = rb + gb[0] w_l = -(rb + gb[1]) w_r = rb + gb[1] w_prf = (rb + gb[1]) * 2 + 1 h_prf = (rb + gb[0]) * 2 + 1 prf_x_idx, prf_y_idx = torch.meshgrid( torch.arange(h_t, h_d + 1), torch.arange(w_l, w_r + 1), indexing='ij') # grid_size*2, 1 grid_size = (h_d*2+1) * (w_r*2+1) # taking positions clockwise prf_xt = prf_x_idx[0:rb] prf_xr = prf_x_idx[rb:(h_prf - 2 * rb + 1), (w_prf - rb):w_prf] prf_xd = prf_x_idx[(h_prf - rb):h_prf] prf_xl = prf_x_idx[rb:(h_prf - 2 * rb + 1), 0:rb] prf_x = torch.cat([torch.flatten(prf_xt), torch.flatten(prf_xr), torch.flatten(prf_xd), torch.flatten(prf_xl)], 0) prf_yt = prf_y_idx[0:rb] prf_yr = prf_y_idx[rb:(h_prf - 2 * rb + 1), (w_prf - rb):w_prf] prf_yd = prf_y_idx[(h_prf - rb):h_prf] prf_yl = prf_y_idx[rb:(h_prf - 2 * rb + 1), 0:rb] prf_y = torch.cat([torch.flatten(prf_yt), torch.flatten(prf_yr), torch.flatten(prf_yd), torch.flatten(prf_yl)], 0) prf = torch.cat([prf_x, prf_y], 0) prf = prf.view(1, 2 * N, 1, 1).type(dtype) return prf # getting p_r positions from each p_c def _get_p_r(self, b, h, w, N, dtype): rb = self.refer_band gb = self.guard_band # (1, 2N, 1, 1) prf = self._gen_prf_grid(rb, gb, N, dtype) # (B, 2N, h, w) p_c = self._get_p_c(b, h, w, N, dtype) # (B, 2N, h, w) p_r = p_c + prf return p_r # sampling x using p_r or p_c def _sample_x(self, x_pad, p, N): p = p.contiguous().permute(0, 2, 3, 1).floor() b, h, w, _ = p.size() h_pad = x_pad.size(2) w_pad = x_pad.size(3) c = x_pad.size(1) # (b, c, h_pad*w_pad) # strech each spatial channel of x_pad as 1-D vector x_pad = x_pad.contiguous().view(b, c, -1) # (b, h_pad, w_pad, N) # transform spatial coord of p into the 1-D index index = p[..., :N] * w_pad + p[..., N:] # index_x = torch.clamp(p[..., :N], 0, h_pad-1) # index_y = torch.clamp(p[..., N:], 0, w_pad-1) # index = index_x * w_pad + index_y # (b, c, h*w*N) index = index.contiguous().unsqueeze(dim=1).expand(-1, c, -1, -1, -1).contiguous().view(b, c, -1) index = index.long().to(x_pad.device) x_r = x_pad.gather(dim=-1, index=index).contiguous().view(b, c, h, w, N) return x_r @staticmethod # reshape the x_prf def _reshape_x_prf(k_h, k_w, x_prf): b, c, h, w, N = x_prf.size() x_prf = torch.cat([x_prf[..., s:s + k_w].contiguous().view(b, c, h, w * k_w) for s in range(0, N, k_w)], dim=-1) x_prf = x_prf.contiguous().view(b, c, h * k_h, w * k_w) return x_prf # double PeakConv2D block definition zlw @20220621 class DoublePeakConv2D(nn.Module): """ (2D PeakConv => BN => LeakyReLU) * 2 """ def __init__(self, in_ch, out_ch, bias, rb, gb_dop, gb, sig): super(DoublePeakConv2D, self).__init__() self.bias = bias self.refer_band = rb self.guard_band_dop = gb_dop self.guard_band = gb self.signal_type = sig self.pk_conv1 = PeakConv2D(in_ch, out_ch, bias=self.bias, refer_band=self.refer_band, guard_band_dop=self.guard_band_dop, guard_band=self.guard_band, signal_type=self.signal_type) self.bn1 = nn.BatchNorm2d(out_ch) self.act1 = nn.LeakyReLU(inplace=True) self.pk_conv2 = PeakConv2D(out_ch, out_ch, bias=self.bias, refer_band=self.refer_band, guard_band_dop=self.guard_band_dop, guard_band=self.guard_band, signal_type=self.signal_type) self.bn2 = nn.BatchNorm2d(out_ch) self.act2 = nn.LeakyReLU(inplace=True) def forward(self, x): x = self.pk_conv1(x) # print("x:", type(x)) x = self.bn1(x) x = self.act1(x) x = self.pk_conv2(x) x = self.bn2(x) x = self.act2(x) return x class ASPPBlock(nn.Module): """Atrous Spatial Pyramid Pooling Parallel conv blocks with different dilation rate """ def __init__(self, in_ch, out_ch, shape): super().__init__() self.shape = shape self.global_avg_pool = nn.AvgPool2d(self.shape) self.conv1_1x1 = nn.Conv2d(in_ch, out_ch, kernel_size=1, padding=0, dilation=1) self.single_conv_block1_1x1 = ConvBlock(in_ch, out_ch, k_size=1, pad=0, dil=1) self.single_conv_block1_3x3 = ConvBlock(in_ch, out_ch, k_size=3, pad=6, dil=6) self.single_conv_block2_3x3 = ConvBlock(in_ch, out_ch, k_size=3, pad=12, dil=12) self.single_conv_block3_3x3 = ConvBlock(in_ch, out_ch, k_size=3, pad=18, dil=18) def forward(self, x): x1 = F.interpolate(self.global_avg_pool(x), size=self.shape, align_corners=False, mode='bilinear') x1 = self.conv1_1x1(x1) x2 = self.single_conv_block1_1x1(x) x3 = self.single_conv_block1_3x3(x) x4 = self.single_conv_block2_3x3(x) x5 = self.single_conv_block3_3x3(x) x_cat = torch.cat((x2, x3, x4, x5, x1), 1) return x_cat class EncodingBranch(nn.Module): """ Encoding branch for a single radar view PARAMETERS ---------- signal_type: str Type of radar view. Supported: 'range_doppler', 'range_angle' and 'angle_doppler' """ def __init__(self, signal_type, n_frames, in_ch=1): super().__init__() self.signal_type = signal_type self.n_frames = n_frames self.in_ch = in_ch self.double_3dconv_block1 = Double3DConvBlock(in_ch=self.in_ch, out_ch=32, k_size=(1+n_frames//2, 3, 3), pad=(0, 1, 1), dil=1) self.doppler_max_pool = nn.MaxPool2d(2, stride=(1, 2)) self.max_pool = nn.MaxPool2d(2, stride=2) # self.double_conv_block2 = DoubleConvBlock(in_ch=128, out_ch=128, k_size=3, # pad=1, dil=1) self.double_conv_block2 = DoubleConvBlock(in_ch=32, out_ch=64, k_size=3, pad=1, dil=1) # self.double_pkc_block = DoublePeakConv2D(in_ch=32, out_ch=64, # bias=True, # rb=1, # gb_dop=1, # gb=1, # sig=self.signal_type) self.single_conv_block1_1x1 = ConvBlock(in_ch=64, out_ch=128, k_size=1, pad=0, dil=1) def forward(self, x): x1 = self.double_3dconv_block1(x) x1 = torch.squeeze(x1, 2) x1_pad = F.pad(x1, (0, 0, 0, 0), "constant", 0) x1_down = self.max_pool(x1_pad) x2 = self.double_conv_block2(x1_down) x2_pad = F.pad(x2, (0, 0, 0, 1), "constant", 0) x2_down = self.doppler_max_pool(x2_pad) # x3 = self.single_conv_block1_1x1(x2_down) # return input of ASPP block + latent features return x2_down, x3 class TMVANet_SV(nn.Module): """ Temporal Singel-View with ASPP Network (TMVA-Net) PARAMETERS ---------- n_classes: int Number of classes used for the semantic segmentation task n_frames: int Total numer of frames used as a sequence """ def __init__(self, n_classes, n_frames, in_ch, w, h): super().__init__() self.n_classes = n_classes self.n_frames = n_frames self.in_ch = in_ch self.rd_shape = (int(w//2), int(h//4)) # Backbone (encoding) self.rd_encoding_branch = EncodingBranch('range_doppler', self.n_frames, self.in_ch) # ASPP Blocks self.rd_aspp_block = ASPPBlock(in_ch=64, out_ch=128, shape=self.rd_shape) self.rd_single_conv_block1_1x1 = ConvBlock(in_ch=640, out_ch=128, k_size=1, pad=0, dil=1) # Decoding # self.rd_single_conv_block2_1x1 = ConvBlock(in_ch=128, out_ch=128, k_size=1, pad=0, dil=1) # Pallel range-Doppler (RD) and range-angle (RA) decoding branches self.rd_upconv1 = nn.ConvTranspose2d(256, 128, (2, 2), stride=(2, 2)) self.rd_double_conv_block1 = DoubleConvBlock(in_ch=128, out_ch=128, k_size=3, pad=1, dil=1) self.rd_upconv2 = nn.ConvTranspose2d(128, 128, (1, 2), stride=(1, 2)) self.rd_double_conv_block2 = DoubleConvBlock(in_ch=128, out_ch=128, k_size=3, pad=1, dil=1) # Final 1D convs self.rd_final = nn.Conv2d(in_channels=128, out_channels=n_classes, kernel_size=1) def forward(self, x_rd): # Backbone rd_features, rd_latent = self.rd_encoding_branch(x_rd) # ASPP blocks x2_rd = self.rd_aspp_block(rd_features) x2_rd = self.rd_single_conv_block1_1x1(x2_rd) # Latent Space + ASPP features x4_rd = torch.cat((x2_rd, rd_latent), 1) # Parallel decoding branches with upconvs x4_rd = self.rd_upconv1(x4_rd) x4_rd = self.rd_double_conv_block1(x4_rd) x4_rd = self.rd_upconv2(x4_rd) x4_rd = self.rd_double_conv_block2(x4_rd) # Final 1D convolutions x4_rd = self.rd_final(x4_rd) return x4_rd if __name__ == '__main__': data = torch.randn(8,1,1,128,1024) model = TMVANet_SV(2,1,1,128,1024) pre = model(data) print(pre)