|
|
|
|
|
|
|
import torch
|
|
|
|
from densepose.structures.data_relative import DensePoseDataRelative
|
|
|
|
|
|
class DensePoseList:
|
|
|
|
_TORCH_DEVICE_CPU = torch.device("cpu")
|
|
|
|
def __init__(self, densepose_datas, boxes_xyxy_abs, image_size_hw, device=_TORCH_DEVICE_CPU):
|
|
assert len(densepose_datas) == len(
|
|
boxes_xyxy_abs
|
|
), "Attempt to initialize DensePoseList with {} DensePose datas " "and {} boxes".format(
|
|
len(densepose_datas), len(boxes_xyxy_abs)
|
|
)
|
|
self.densepose_datas = []
|
|
for densepose_data in densepose_datas:
|
|
assert isinstance(densepose_data, DensePoseDataRelative) or densepose_data is None, (
|
|
"Attempt to initialize DensePoseList with DensePose datas "
|
|
"of type {}, expected DensePoseDataRelative".format(type(densepose_data))
|
|
)
|
|
densepose_data_ondevice = (
|
|
densepose_data.to(device) if densepose_data is not None else None
|
|
)
|
|
self.densepose_datas.append(densepose_data_ondevice)
|
|
self.boxes_xyxy_abs = boxes_xyxy_abs.to(device)
|
|
self.image_size_hw = image_size_hw
|
|
self.device = device
|
|
|
|
def to(self, device):
|
|
if self.device == device:
|
|
return self
|
|
return DensePoseList(self.densepose_datas, self.boxes_xyxy_abs, self.image_size_hw, device)
|
|
|
|
def __iter__(self):
|
|
return iter(self.densepose_datas)
|
|
|
|
def __len__(self):
|
|
return len(self.densepose_datas)
|
|
|
|
def __repr__(self):
|
|
s = self.__class__.__name__ + "("
|
|
s += "num_instances={}, ".format(len(self.densepose_datas))
|
|
s += "image_width={}, ".format(self.image_size_hw[1])
|
|
s += "image_height={})".format(self.image_size_hw[0])
|
|
return s
|
|
|
|
def __getitem__(self, item):
|
|
if isinstance(item, int):
|
|
densepose_data_rel = self.densepose_datas[item]
|
|
return densepose_data_rel
|
|
elif isinstance(item, slice):
|
|
densepose_datas_rel = self.densepose_datas[item]
|
|
boxes_xyxy_abs = self.boxes_xyxy_abs[item]
|
|
return DensePoseList(
|
|
densepose_datas_rel, boxes_xyxy_abs, self.image_size_hw, self.device
|
|
)
|
|
elif isinstance(item, torch.Tensor) and (item.dtype == torch.bool):
|
|
densepose_datas_rel = [self.densepose_datas[i] for i, x in enumerate(item) if x > 0]
|
|
boxes_xyxy_abs = self.boxes_xyxy_abs[item]
|
|
return DensePoseList(
|
|
densepose_datas_rel, boxes_xyxy_abs, self.image_size_hw, self.device
|
|
)
|
|
else:
|
|
densepose_datas_rel = [self.densepose_datas[i] for i in item]
|
|
boxes_xyxy_abs = self.boxes_xyxy_abs[item]
|
|
return DensePoseList(
|
|
densepose_datas_rel, boxes_xyxy_abs, self.image_size_hw, self.device
|
|
)
|
|
|
