| 12 | |
| 13 | |
| 14 | class ImageEncoder(nn.Module): |
| 15 | def __init__( |
| 16 | self, |
| 17 | trunk: nn.Module, |
| 18 | neck: nn.Module, |
| 19 | scalp: int = 0, |
| 20 | ): |
| 21 | super().__init__() |
| 22 | self.trunk = trunk |
| 23 | self.neck = neck |
| 24 | self.scalp = scalp |
| 25 | assert ( |
| 26 | self.trunk.channel_list == self.neck.backbone_channel_list |
| 27 | ), f"Channel dims of trunk and neck do not match. Trunk: {self.trunk.channel_list}, neck: {self.neck.backbone_channel_list}" |
| 28 | |
| 29 | def forward(self, sample: torch.Tensor): |
| 30 | # Forward through backbone |
| 31 | features, pos = self.neck(self.trunk(sample)) |
| 32 | if self.scalp > 0: |
| 33 | # Discard the lowest resolution features |
| 34 | features, pos = features[: -self.scalp], pos[: -self.scalp] |
| 35 | |
| 36 | src = features[-1] |
| 37 | output = { |
| 38 | "vision_features": src, |
| 39 | "vision_pos_enc": pos, |
| 40 | "backbone_fpn": features, |
| 41 | } |
| 42 | return output |
| 43 | |
| 44 | |
| 45 | class FpnNeck(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected