| | |
| |
|
| | from dataclasses import dataclass |
| | from typing import Union |
| | import torch |
| |
|
| |
|
| | @dataclass |
| | class DensePoseChartPredictorOutput: |
| | """ |
| | Predictor output that contains segmentation and inner coordinates predictions for predefined |
| | body parts: |
| | * coarse segmentation, a tensor of shape [N, K, Hout, Wout] |
| | * fine segmentation, a tensor of shape [N, C, Hout, Wout] |
| | * U coordinates, a tensor of shape [N, C, Hout, Wout] |
| | * V coordinates, a tensor of shape [N, C, Hout, Wout] |
| | where |
| | - N is the number of instances |
| | - K is the number of coarse segmentation channels ( |
| | 2 = foreground / background, |
| | 15 = one of 14 body parts / background) |
| | - C is the number of fine segmentation channels ( |
| | 24 fine body parts / background) |
| | - Hout and Wout are height and width of predictions |
| | """ |
| |
|
| | coarse_segm: torch.Tensor |
| | fine_segm: torch.Tensor |
| | u: torch.Tensor |
| | v: torch.Tensor |
| |
|
| | def __len__(self): |
| | """ |
| | Number of instances (N) in the output |
| | """ |
| | return self.coarse_segm.size(0) |
| |
|
| | def __getitem__( |
| | self, item: Union[int, slice, torch.BoolTensor] |
| | ) -> "DensePoseChartPredictorOutput": |
| | """ |
| | Get outputs for the selected instance(s) |
| | |
| | Args: |
| | item (int or slice or tensor): selected items |
| | """ |
| | if isinstance(item, int): |
| | return DensePoseChartPredictorOutput( |
| | coarse_segm=self.coarse_segm[item].unsqueeze(0), |
| | fine_segm=self.fine_segm[item].unsqueeze(0), |
| | u=self.u[item].unsqueeze(0), |
| | v=self.v[item].unsqueeze(0), |
| | ) |
| | else: |
| | return DensePoseChartPredictorOutput( |
| | coarse_segm=self.coarse_segm[item], |
| | fine_segm=self.fine_segm[item], |
| | u=self.u[item], |
| | v=self.v[item], |
| | ) |
| |
|
| | def to(self, device: torch.device): |
| | """ |
| | Transfers all tensors to the given device |
| | """ |
| | coarse_segm = self.coarse_segm.to(device) |
| | fine_segm = self.fine_segm.to(device) |
| | u = self.u.to(device) |
| | v = self.v.to(device) |
| | return DensePoseChartPredictorOutput(coarse_segm=coarse_segm, fine_segm=fine_segm, u=u, v=v) |
| |
|
