| 219 | return self.rho ** 2 * dist |
| 220 | |
| 221 | class PerspectiveCamera(nn.Module): |
| 222 | |
| 223 | FOCAL_LENGTH = 50*128 |
| 224 | |
| 225 | def __init__(self, rotation=None, translation=None, |
| 226 | focal_length_x=None, focal_length_y=None, |
| 227 | batch_size=1, |
| 228 | center=None, dtype=torch.float32): |
| 229 | super(PerspectiveCamera, self).__init__() |
| 230 | self.batch_size = batch_size |
| 231 | self.dtype = dtype |
| 232 | # Make a buffer so that PyTorch does not complain when creating |
| 233 | # the camera matrix |
| 234 | self.register_buffer('zero', |
| 235 | torch.zeros([batch_size], dtype=dtype)) |
| 236 | |
| 237 | if focal_length_x is None or type(focal_length_x) == float: |
| 238 | focal_length_x = torch.full( |
| 239 | [batch_size], |
| 240 | self.FOCAL_LENGTH if focal_length_x is None else |
| 241 | focal_length_x, |
| 242 | dtype=dtype) |
| 243 | |
| 244 | if focal_length_y is None or type(focal_length_y) == float: |
| 245 | focal_length_y = torch.full( |
| 246 | [batch_size], |
| 247 | self.FOCAL_LENGTH if focal_length_y is None else |
| 248 | focal_length_y, |
| 249 | dtype=dtype) |
| 250 | |
| 251 | self.register_buffer('focal_length_x', focal_length_x) |
| 252 | self.register_buffer('focal_length_y', focal_length_y) |
| 253 | |
| 254 | if center is None: |
| 255 | center = torch.zeros([batch_size, 2], dtype=dtype) |
| 256 | self.register_buffer('center', center) |
| 257 | |
| 258 | if rotation is None: |
| 259 | rotation = torch.eye( |
| 260 | 3, dtype=dtype).unsqueeze(dim=0).repeat(batch_size, 1, 1) |
| 261 | |
| 262 | rotation = nn.Parameter(rotation, requires_grad=False) |
| 263 | self.register_parameter('rotation', rotation) |
| 264 | |
| 265 | if translation is None: |
| 266 | translation = torch.zeros([batch_size, 3], dtype=dtype) |
| 267 | |
| 268 | translation = nn.Parameter(translation, |
| 269 | requires_grad=True) |
| 270 | self.register_parameter('translation', translation) |
| 271 | |
| 272 | def forward(self, points): |
| 273 | device = points.device |
| 274 | with torch.no_grad(): |
| 275 | camera_mat = torch.zeros([self.batch_size, 2, 2], |
| 276 | dtype=self.dtype, device=points.device) |
| 277 | camera_mat[:, 0, 0] = self.focal_length_x |
| 278 | camera_mat[:, 1, 1] = self.focal_length_y |