A pose estimation model. A pose estimation model is composed of a backbone, optionally a neck, and an arbitrary number of heads. Outputs are computed as follows:
| 31 | |
| 32 | |
| 33 | class PoseModel(nn.Module): |
| 34 | """A pose estimation model. |
| 35 | |
| 36 | A pose estimation model is composed of a backbone, optionally a neck, and an |
| 37 | arbitrary number of heads. Outputs are computed as follows: |
| 38 | """ |
| 39 | |
| 40 | def __init__( |
| 41 | self, |
| 42 | cfg: dict, |
| 43 | backbone: BaseBackbone, |
| 44 | heads: dict[str, BaseHead], |
| 45 | neck: BaseNeck | None = None, |
| 46 | ) -> None: |
| 47 | """ |
| 48 | Args: |
| 49 | cfg: configuration dictionary for the model. |
| 50 | backbone: backbone network architecture. |
| 51 | heads: the heads for the model |
| 52 | neck: neck network architecture (default is None). Defaults to None. |
| 53 | """ |
| 54 | super().__init__() |
| 55 | self.cfg = cfg |
| 56 | self.backbone = backbone |
| 57 | self.heads = nn.ModuleDict(heads) |
| 58 | self.neck = neck |
| 59 | self.output_features = False |
| 60 | |
| 61 | self._strides = {name: _model_stride(self.backbone.stride, head.stride) for name, head in heads.items()} |
| 62 | |
| 63 | def forward(self, x: torch.Tensor, **backbone_kwargs) -> dict[str, dict[str, torch.Tensor]]: |
| 64 | """Forward pass of the PoseModel. |
| 65 | |
| 66 | Args: |
| 67 | x: input images |
| 68 | |
| 69 | Returns: |
| 70 | Outputs of head groups |
| 71 | """ |
| 72 | if x.dim() == 3: |
| 73 | x = x[None, :] |
| 74 | features = self.backbone(x, **backbone_kwargs) |
| 75 | if self.neck: |
| 76 | features = self.neck(features) |
| 77 | |
| 78 | outputs = {} |
| 79 | if self.output_features: |
| 80 | outputs["backbone"] = dict(features=features) |
| 81 | |
| 82 | for head_name, head in self.heads.items(): |
| 83 | outputs[head_name] = head(features) |
| 84 | return outputs |
| 85 | |
| 86 | def get_loss( |
| 87 | self, |
| 88 | outputs: dict[str, dict[str, torch.Tensor]], |
| 89 | targets: dict[str, dict[str, torch.Tensor]], |
| 90 | ) -> dict[str, torch.Tensor]: |
no outgoing calls