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