| 1597 | |
| 1598 | |
| 1599 | class Camera: |
| 1600 | def __init__( |
| 1601 | self, |
| 1602 | H_c2w: torch.Tensor, # (b, q, 4, 4) camera pose in the world coord |
| 1603 | intrinsic: torch.Tensor, # (b, q, 3, 3) camera intrinsics |
| 1604 | width_px: int, |
| 1605 | height_px: int, |
| 1606 | ): |
| 1607 | self.H_c2w = H_c2w |
| 1608 | self.intrinsic = intrinsic |
| 1609 | self.width_px = width_px |
| 1610 | self.height_px = height_px |
| 1611 | |
| 1612 | self.attr_names = ['H_c2w', 'intrinsic', 'width_px', 'height_px'] |
| 1613 | |
| 1614 | def index_select(self, dim: int, index: torch.Tensor) -> 'Camera': |
| 1615 | camera = self.clone() |
| 1616 | for attr_name in ['H_c2w', 'intrinsic']: |
| 1617 | arr = getattr(camera, attr_name, None) |
| 1618 | if arr is not None: |
| 1619 | setattr(camera, attr_name, torch.index_select(arr, dim=dim, index=index)) |
| 1620 | return camera |
| 1621 | |
| 1622 | def chunk(self, chunks: int, dim: int = 0) -> T.List['Camera']: |
| 1623 | out_dict = dict() |
| 1624 | total = None |
| 1625 | for attr_name in ['H_c2w', 'intrinsic']: |
| 1626 | arr = getattr(self, attr_name, None) |
| 1627 | if arr is not None: |
| 1628 | chunked_arr = arr.chunk(chunks=chunks, dim=dim) |
| 1629 | out_dict[attr_name] = chunked_arr |
| 1630 | if total is None: |
| 1631 | total = len(chunked_arr) |
| 1632 | else: |
| 1633 | assert len(chunked_arr) == total |
| 1634 | cameras = [] |
| 1635 | for i in range(total): |
| 1636 | d = dict() |
| 1637 | for attr_name in out_dict: |
| 1638 | d[attr_name] = out_dict[attr_name][i] |
| 1639 | camera = Camera(**d, width_px=self.width_px, height_px=self.height_px) |
| 1640 | cameras.append(camera) |
| 1641 | return cameras |
| 1642 | |
| 1643 | def to(self, device: torch.device) -> 'Camera': |
| 1644 | self.H_c2w = self.H_c2w.to(device=device) |
| 1645 | self.intrinsic = self.intrinsic.to(device=device) |
| 1646 | return self |
| 1647 | |
| 1648 | def detach(self) -> 'Camera': |
| 1649 | self.H_c2w = self.H_c2w.detach() |
| 1650 | self.intrinsic = self.intrinsic.detach() |
| 1651 | return self |
| 1652 | |
| 1653 | def clone(self) -> 'Camera': |
| 1654 | return Camera( |
| 1655 | H_c2w=self.H_c2w.clone() if self.H_c2w is not None else None, |
| 1656 | intrinsic=self.intrinsic.clone() if self.intrinsic is not None else None, |
no outgoing calls
no test coverage detected