Encode images with backbone and decode into a semantic segmentation map of the same size as input.
(self, rgb, modal_x)
| 223 | ) |
| 224 | |
| 225 | def encode_decode(self, rgb, modal_x): |
| 226 | """Encode images with backbone and decode into a semantic segmentation |
| 227 | map of the same size as input.""" |
| 228 | orisize = rgb.shape |
| 229 | # print('builder',rgb.shape,modal_x.shape) |
| 230 | x = self.backbone(rgb, modal_x) |
| 231 | if len(x) == 2: # if output is (rgb,depth) only use rgb |
| 232 | x = x[0] |
| 233 | out = self.decode_head.forward(x) |
| 234 | out = F.interpolate(out, size=orisize[-2:], mode="bilinear", align_corners=False) |
| 235 | if self.aux_head: |
| 236 | aux_fm = self.aux_head(x[0][self.aux_index]) |
| 237 | aux_fm = F.interpolate(aux_fm, size=orisize[2:], mode="bilinear", align_corners=False) |
| 238 | return out, aux_fm |
| 239 | return out |
| 240 | |
| 241 | def forward(self, rgb, modal_x=None, label=None): |
| 242 | # print('builder',rgb.shape,modal_x.shape) |