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