|
|
|
|
|
from dataclasses import dataclass |
|
from typing import Union |
|
import torch |
|
|
|
|
|
@dataclass |
|
class DensePoseEmbeddingPredictorOutput: |
|
""" |
|
Predictor output that contains embedding and coarse segmentation data: |
|
* embedding: float tensor of size [N, D, H, W], contains estimated embeddings |
|
* coarse_segm: float tensor of size [N, K, H, W] |
|
Here D = MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE |
|
K = MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS |
|
""" |
|
|
|
embedding: torch.Tensor |
|
coarse_segm: 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] |
|
) -> "DensePoseEmbeddingPredictorOutput": |
|
""" |
|
Get outputs for the selected instance(s) |
|
|
|
Args: |
|
item (int or slice or tensor): selected items |
|
""" |
|
if isinstance(item, int): |
|
return DensePoseEmbeddingPredictorOutput( |
|
coarse_segm=self.coarse_segm[item].unsqueeze(0), |
|
embedding=self.embedding[item].unsqueeze(0), |
|
) |
|
else: |
|
return DensePoseEmbeddingPredictorOutput( |
|
coarse_segm=self.coarse_segm[item], embedding=self.embedding[item] |
|
) |
|
|
|
def to(self, device: torch.device): |
|
""" |
|
Transfers all tensors to the given device |
|
""" |
|
coarse_segm = self.coarse_segm.to(device) |
|
embedding = self.embedding.to(device) |
|
return DensePoseEmbeddingPredictorOutput(coarse_segm=coarse_segm, embedding=embedding) |
|
|
