File size: 4,602 Bytes
4f6b78d 60fd7ba 4f6b78d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact [email protected]
#
import torch
from torch import nn
import numpy as np
from utils_das3r.graphics_utils import getWorld2View2, getProjectionMatrix
class Camera(nn.Module):
def __init__(self, colmap_id, intr, R, T, original_pose, FoVx, FoVy, image, gt_alpha_mask, dynamic_mask, enlarged_dynamic_mask,
dyna_avg_map, dyna_max_map, gt_dynamic_mask,
conf_map, depth_map,
image_name, uid,
trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device = "cuda"
):
super(Camera, self).__init__()
self.uid = uid
self.colmap_id = colmap_id
self.R = R
self.T = T
self.FoVx = FoVx
self.FoVy = FoVy
self.image_name = image_name
try:
self.data_device = torch.device(data_device)
except Exception as e:
print(e)
print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" )
self.data_device = torch.device("cuda")
self.original_image = image.clamp(0.0, 1.0).to(self.data_device)
self.image_width = self.original_image.shape[2]
self.image_height = self.original_image.shape[1]
if original_pose is not None:
self.original_pose = torch.tensor(original_pose, dtype=torch.float32).to(self.data_device)
if intr is not None:
self.intr = intr
if conf_map is not None:
self.conf_map = conf_map.to(self.data_device)
if depth_map is not None:
self.depth_map = depth_map.to(self.data_device)
if dynamic_mask is not None:
self.dynamic_mask = dynamic_mask.to(self.data_device)
if gt_dynamic_mask is not None:
gt_dynamic_mask = gt_dynamic_mask.to(self.data_device)
gt_dynamic_mask = gt_dynamic_mask.unsqueeze(0).repeat(3, 1, 1).unsqueeze(0).float()
self.gt_dynamic_mask = torch.nn.functional.interpolate(
gt_dynamic_mask,
size=(self.image_height, self.image_width),
mode="nearest",
).squeeze(0)
if enlarged_dynamic_mask is not None:
self.enlarged_dynamic_mask = enlarged_dynamic_mask.to(self.data_device)
if dyna_avg_map is not None:
self.dyna_avg_map = dyna_avg_map.to(self.data_device)
if dyna_max_map is not None:
self.dyna_max_map = dyna_max_map.to(self.data_device)
if gt_alpha_mask is not None:
self.original_image *= gt_alpha_mask.to(self.data_device)
else:
self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device)
self.zfar = 100.0
self.znear = 0.01
self.trans = trans
self.scale = scale
self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda()
self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda()
self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0)
self.camera_center = self.world_view_transform.inverse()[3, :3]
def get_full_proj_transform(self, FoVx, FoVy):
self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=FoVx, fovY=FoVy).transpose(0,1).cuda()
return (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0)
def get_projection_matrix(self, FoVx, FoVy):
return getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=FoVx, fovY=FoVy).transpose(0,1).cuda()
class MiniCam:
def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform):
self.image_width = width
self.image_height = height
self.FoVy = fovy
self.FoVx = fovx
self.znear = znear
self.zfar = zfar
self.world_view_transform = world_view_transform
self.full_proj_transform = full_proj_transform
view_inv = torch.inverse(self.world_view_transform)
self.camera_center = view_inv[3][:3]
|