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