| import torch |
| import torch.nn as nn |
| from models.bts import encoder_image, bts_gated_fuse |
| from models.radar import encoder_radar_sparse_conv, encoder_radar_sub, decoder_radar |
|
|
| class CaFNet(nn.Module): |
| def __init__(self, params, threshold=0.4): |
| super(CaFNet, self).__init__() |
| self.threshold = threshold |
| self.encoder = encoder_image(params) |
| self.encoder_radar1 = encoder_radar_sparse_conv(params) |
| self.decoder_radar = decoder_radar(params, self.encoder.feat_out_channels, self.encoder_radar1.feat_out_channels) |
| self.encoder_radar2 = encoder_radar_sub(params) |
| self.decoder = bts_gated_fuse(params, self.encoder.feat_out_channels, self.encoder_radar2.feat_out_channels, params.bts_size) |
|
|
|
|
| def forward(self, x, radar, focal): |
|
|
| skip_feat = self.encoder(x) |
| skip_feat_radar = self.encoder_radar1(radar) |
| rad_confidence, rad_depth = self.decoder_radar(skip_feat, skip_feat_radar) |
| mask = (rad_confidence > self.threshold).float() |
| radar_new_input = torch.cat([mask*rad_depth, radar], axis=1) |
| skip_feat_radar_new = self.encoder_radar2(radar_new_input) |
| |
| depth_8x8_scaled, depth_4x4_scaled, depth_2x2_scaled, reduc1x1, final_depth = self.decoder(skip_feat, skip_feat_radar_new, focal, rad_confidence) |
|
|
| return depth_8x8_scaled, depth_4x4_scaled, depth_2x2_scaled, reduc1x1, final_depth, rad_confidence, rad_depth |
|
|
