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