| 3 | |
| 4 | |
| 5 | class Pinhole2D: |
| 6 | def __init__(self, K=None, fx=None, fy=None, cx=None, cy=None, h=0, w=0): |
| 7 | if K is not None: |
| 8 | assert fx is None and fy is None and cx is None and cy is None |
| 9 | self.fx = K[0, 0] |
| 10 | self.fy = K[1, 1] |
| 11 | self.cx = K[0, 2] |
| 12 | self.cy = K[1, 2] |
| 13 | else: |
| 14 | assert \ |
| 15 | fx is not None and fy is not None and \ |
| 16 | cx is not None and cy is not None |
| 17 | self.fx = fx |
| 18 | self.fy = fy |
| 19 | self.cx = cx |
| 20 | self.cy = cy |
| 21 | self.h = h |
| 22 | self.w = w |
| 23 | |
| 24 | def get_K(self): |
| 25 | return np.array([[self.fx, 0, self.cx], [0, self.fy, self.cy], |
| 26 | [0, 0, 1]]) |
| 27 | |
| 28 | def project_ndc(self, vertices, eps=1e-9): |
| 29 | """ |
| 30 | vertices: torch.Tensor of shape (N, 3), 3 stands for xyz |
| 31 | """ |
| 32 | assert isinstance(vertices, torch.Tensor) |
| 33 | assert len(vertices.shape) == 2 |
| 34 | assert vertices.shape[1] == 3 |
| 35 | K = torch.tensor(self.get_K(), |
| 36 | device=vertices.device, |
| 37 | dtype=vertices.dtype) |
| 38 | |
| 39 | # apply intrinsics |
| 40 | vertices_ndc = vertices @ K.transpose(1, 0) |
| 41 | |
| 42 | # divide xy by z, leave z unchanged |
| 43 | vertices_ndc[:, [0, 1]] /= vertices_ndc[:, [2]] + eps |
| 44 | |
| 45 | # convert x from [0, w) to [-1, 1] range |
| 46 | # convert y from [0, h) to [-1, 1] range |
| 47 | wh = torch.tensor([self.w, self.h], |
| 48 | device=vertices.device, |
| 49 | dtype=vertices.dtype).unsqueeze(0) |
| 50 | vertices_ndc[:, [0, 1]] = 2 * vertices_ndc[:, [0, 1]] / wh - 1 |
| 51 | |
| 52 | return vertices_ndc |