|
|
|
|
|
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) |
|
|
