cameras.py 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. import torch
  2. from torch import nn
  3. import numpy as np
  4. from utils.graphics_utils import getWorld2View2, getProjectionMatrix
  5. class Camera(nn.Module):
  6. def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask,
  7. image_name, uid,
  8. trans=np.array([0.0, 0.0, 0.0]), scale=1.0
  9. ):
  10. super(Camera, self).__init__()
  11. self.uid = uid
  12. self.colmap_id = colmap_id
  13. self.R = R
  14. self.T = T
  15. self.FoVx = FoVx
  16. self.FoVy = FoVy
  17. self.image_name = image_name
  18. self.original_image = image.clamp(0.0, 1.0).cuda()
  19. self.image_width = self.original_image.shape[2]
  20. self.image_height = self.original_image.shape[1]
  21. if gt_alpha_mask is not None:
  22. self.original_image *= gt_alpha_mask.cuda()
  23. else:
  24. self.original_image *= torch.ones((1, self.image_height, self.image_width), device="cuda")
  25. self.zfar = 100.0
  26. self.znear = 0.01
  27. self.trans = trans
  28. self.scale = scale
  29. self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda()
  30. self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda()
  31. self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0)
  32. self.camera_center = self.world_view_transform.inverse()[3, :3]
  33. class MiniCam:
  34. def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform):
  35. self.image_width = width
  36. self.image_height = height
  37. self.FoVy = fovy
  38. self.FoVx = fovx
  39. self.znear = znear
  40. self.zfar = zfar
  41. self.world_view_transform = world_view_transform
  42. self.full_proj_transform = full_proj_transform
  43. view_inv = torch.inverse(self.world_view_transform)
  44. self.camera_center = view_inv[3][:3]