Encoder Decoder depther. EncoderDecoder typically consists of backbone and decode_head.
| 32 | |
| 33 | |
| 34 | class DepthEncoderDecoder(nn.Module): |
| 35 | """Encoder Decoder depther. |
| 36 | |
| 37 | EncoderDecoder typically consists of backbone and decode_head. |
| 38 | """ |
| 39 | |
| 40 | def __init__(self, backbone, decode_head): |
| 41 | super(DepthEncoderDecoder, self).__init__() |
| 42 | |
| 43 | self.backbone = backbone |
| 44 | self.decode_head = decode_head |
| 45 | self.align_corners = self.decode_head.align_corners |
| 46 | |
| 47 | def extract_feat(self, img): |
| 48 | """Extract features from images.""" |
| 49 | return self.backbone(img) |
| 50 | |
| 51 | def encode_decode(self, img, img_metas, rescale=True, size=None): |
| 52 | """Encode images with backbone and decode into a depth estimation |
| 53 | map of the same size as input.""" |
| 54 | x = self.extract_feat(img) |
| 55 | out = self._decode_head_forward_test(x, img_metas) |
| 56 | # crop the pred depth to the certain range. |
| 57 | out = torch.clamp(out, min=self.decode_head.min_depth, max=self.decode_head.max_depth) |
| 58 | if rescale: |
| 59 | if size is None: |
| 60 | if img_metas is not None: |
| 61 | size = img_metas[0]["ori_shape"][:2] |
| 62 | else: |
| 63 | size = img.shape[2:] |
| 64 | out = resize(input=out, size=size, mode="bilinear", align_corners=self.align_corners) |
| 65 | return out |
| 66 | |
| 67 | def _decode_head_forward_train(self, img, x, img_metas, depth_gt, **kwargs): |
| 68 | """Run forward function and calculate loss for decode head in |
| 69 | training.""" |
| 70 | losses = dict() |
| 71 | loss_decode = self.decode_head.forward_train(img, x, img_metas, depth_gt, **kwargs) |
| 72 | losses.update(add_prefix(loss_decode, "decode")) |
| 73 | return losses |
| 74 | |
| 75 | def _decode_head_forward_test(self, x, img_metas): |
| 76 | """Run forward function and calculate loss for decode head in |
| 77 | inference.""" |
| 78 | depth_pred = self.decode_head.forward_test(x, img_metas) |
| 79 | return depth_pred |
| 80 | |
| 81 | def forward_dummy(self, img): |
| 82 | """Dummy forward function.""" |
| 83 | depth = self.encode_decode(img, None) |
| 84 | |
| 85 | return depth |
| 86 | |
| 87 | def forward_train(self, img, img_metas, depth_gt, **kwargs): |
| 88 | """Forward function for training. |
| 89 | |
| 90 | Args: |
| 91 | img (Tensor): Input images. |
no outgoing calls
no test coverage detected