mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 19:33:46 +08:00
Ethan Weber
146332cfb6
* in progress. camera_indices broken * faster dataloader. starting on pixel sampler logic * imagesampler and pixelsampler * lint and black * continuing merge * pytest working * fixed docs * docs * minor changes
274 lines
8.0 KiB
Python
274 lines
8.0 KiB
Python
"""
|
|
Camera Models
|
|
"""
|
|
from abc import abstractmethod
|
|
from dataclasses import dataclass
|
|
from typing import Optional, Type
|
|
|
|
import torch
|
|
from torch.nn.functional import normalize
|
|
from torchtyping import TensorType
|
|
|
|
from pyrad.structures.rays import RayBundle
|
|
|
|
|
|
@dataclass
|
|
class Rays:
|
|
"""Camera rays. Can be arbitrary dimension"""
|
|
|
|
origin: TensorType[..., 3]
|
|
direction: TensorType[..., 3]
|
|
|
|
|
|
class Camera:
|
|
"""Base Camera. Intended to be subclassed"""
|
|
|
|
def __init__(self, camera_to_world: Optional[TensorType[3, 4]] = torch.eye(4)[:3]) -> None:
|
|
self.camera_to_world = camera_to_world
|
|
|
|
@abstractmethod
|
|
def get_num_intrinsics_params(self) -> int:
|
|
"""
|
|
Returns:
|
|
int: number of optimizable intrinsic parameters
|
|
"""
|
|
return
|
|
|
|
@abstractmethod
|
|
def get_intrinsics(self) -> torch.Tensor:
|
|
"""
|
|
Returns:
|
|
torch.Tensor: Intrinsics matrix
|
|
"""
|
|
return
|
|
|
|
def get_camera_to_world(self) -> TensorType[3, 4]:
|
|
"""
|
|
Returns:
|
|
int: Camera to world transformation
|
|
"""
|
|
return self.camera_to_world
|
|
|
|
@abstractmethod
|
|
def get_image_height(self) -> int:
|
|
"""
|
|
Returns:
|
|
int: Image height
|
|
"""
|
|
return
|
|
|
|
@abstractmethod
|
|
def get_image_width(self) -> int:
|
|
"""
|
|
Returns:
|
|
int: Image width
|
|
"""
|
|
return
|
|
|
|
def get_image_coords(self, pixel_offset: float = 0.5) -> TensorType["image_height", "image_width", 2]:
|
|
"""
|
|
Args:
|
|
pixel_offset (float): Offset for each pixel. Defaults to center of pixel (0.5)
|
|
|
|
Returns:
|
|
TensorType["image_height", "image_width", 2]: Grid of image coordinates.
|
|
"""
|
|
image_height = self.get_image_height()
|
|
image_width = self.get_image_width()
|
|
image_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing="ij")
|
|
image_coords = torch.stack(image_coords, dim=-1) + pixel_offset # stored as (y, x) coordinates
|
|
return image_coords
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def generate_rays(
|
|
cls,
|
|
intrinsics: TensorType["num_rays", "num_intrinsics_params"],
|
|
camera_to_world: TensorType["num_rays", 3, 4],
|
|
coords: TensorType["image_height", "image_width", 2],
|
|
) -> RayBundle:
|
|
"""
|
|
Args:
|
|
intrinsics (TensorType["num_rays", "num_intrinsics_params"]): Camera intrinsics
|
|
camera_to_world (TensorType["num_rays", 3, 4]): Camera to world transformation matrix
|
|
coords (TensorType["image_height", "image_width", 2]): Image grid coordinates
|
|
|
|
Returns:
|
|
RayBundle: A bundle of rays for each grid coordinate.
|
|
"""
|
|
return
|
|
|
|
def generate_camera_rays(self, device=None) -> RayBundle:
|
|
"""Generate rays for the camera.
|
|
|
|
Returns:
|
|
Rays: Camera rays of shape [height, width]
|
|
"""
|
|
if device is None:
|
|
device = self.camera_to_world.device
|
|
height = self.get_image_height()
|
|
width = self.get_image_width()
|
|
intrinsics = (
|
|
self.get_intrinsics().unsqueeze(0).repeat(height, width, 1).to(device)
|
|
) # (num_rays, num_intrinsics_params)
|
|
camera_to_world = self.camera_to_world.unsqueeze(0).repeat(height, width, 1, 1).to(device) # (num_rays, 3, 4)
|
|
coords = self.get_image_coords().to(device)
|
|
ray_bundle = self.generate_rays(intrinsics, camera_to_world, coords)
|
|
return ray_bundle
|
|
|
|
|
|
class PinholeCamera(Camera):
|
|
"""Pinhole camera model."""
|
|
|
|
def __init__(
|
|
self, cx: float, cy: float, fx: float, fy: float, camera_to_world: Optional[TensorType[3, 4]] = torch.eye(4)[:3]
|
|
):
|
|
super().__init__(camera_to_world)
|
|
self.cx = cx
|
|
self.cy = cy
|
|
self.fx = fx
|
|
self.fy = fy
|
|
|
|
def get_num_intrinsics_params(self) -> int:
|
|
return 4
|
|
|
|
def get_image_width(self) -> int:
|
|
return int(self.cx * 2.0)
|
|
|
|
def get_image_height(self) -> int:
|
|
return int(self.cy * 2.0)
|
|
|
|
def get_intrinsics_matrix(self) -> TensorType[3, 3]:
|
|
"""
|
|
Returns:
|
|
TensorType[3, 3]: Pinhole camera intrinsics matrix
|
|
"""
|
|
K = torch.tensor(
|
|
[[self.fx, 0.0, self.cx], [0.0, self.fy, self.cy], [0, 0, 1.0]],
|
|
dtype=torch.float32,
|
|
)
|
|
return K
|
|
|
|
def get_intrinsics(self) -> torch.Tensor:
|
|
return torch.tensor([self.cx, self.cy, self.fx, self.fy])
|
|
|
|
@classmethod
|
|
def fx_index(cls):
|
|
"""TODO(ethan): redo this in a better way.
|
|
Ideally we can dynamically grab the focal length parameters depending on
|
|
Simple vs. not Simple Pinhole Model.
|
|
|
|
Returns:
|
|
_type_: _description_
|
|
"""
|
|
return 2
|
|
|
|
@classmethod
|
|
def fy_index(cls):
|
|
"""_summary_
|
|
|
|
Returns:
|
|
_type_: _description_
|
|
"""
|
|
return 3
|
|
|
|
@classmethod
|
|
def generate_rays(
|
|
cls,
|
|
intrinsics: TensorType[..., "num_intrinsics_params"],
|
|
camera_to_world: TensorType[..., 3, 4],
|
|
coords: TensorType[..., 2],
|
|
) -> RayBundle:
|
|
|
|
cx = intrinsics[..., 0:1]
|
|
cy = intrinsics[..., 1:2]
|
|
fx = intrinsics[..., cls.fx_index() : cls.fx_index() + 1]
|
|
fy = intrinsics[..., cls.fy_index() : cls.fy_index() + 1]
|
|
y = coords[..., 0:1] # (..., 1)
|
|
x = coords[..., 1:2] # (..., 1)
|
|
original_directions = torch.cat([(x - cx) / fx, -(y - cy) / fy, -torch.ones_like(x)], -1) # (..., 3)
|
|
rotation = camera_to_world[..., :3, :3] # (..., 3, 3)
|
|
directions = torch.sum(
|
|
original_directions[..., None, :] * rotation, dim=-1
|
|
) # (..., 1, 3) * (..., 3, 3) -> (..., 3)
|
|
directions = normalize(directions, dim=-1)
|
|
origins = camera_to_world[..., :3, 3] # (..., 3)
|
|
return RayBundle(origins=origins, directions=directions)
|
|
|
|
|
|
class SimplePinholeCamera(PinholeCamera):
|
|
"""Simple Pinhole Camera model."""
|
|
|
|
def __init__(self, cx: float, cy: float, f: float, camera_to_world: Optional[TensorType[3, 4]] = torch.eye(4)[:3]):
|
|
super().__init__(cx, cy, f, f, camera_to_world)
|
|
|
|
@classmethod
|
|
def fx_index(cls):
|
|
return 2
|
|
|
|
@classmethod
|
|
def fy_index(cls):
|
|
return 2
|
|
|
|
|
|
class EquirectangularCamera(Camera):
|
|
"""Equirectangular (360 degree) camera model."""
|
|
|
|
def __init__(self, height: int, width: int, camera_to_world: Optional[TensorType[3, 4]] = torch.eye(4)[:3]):
|
|
super().__init__(camera_to_world)
|
|
self.height = height
|
|
self.width = width
|
|
|
|
def get_num_intrinsics_params(self):
|
|
return 2
|
|
|
|
def get_image_width(self):
|
|
return self.width
|
|
|
|
def get_image_height(self):
|
|
return self.height
|
|
|
|
def get_intrinsics(self) -> torch.Tensor:
|
|
return torch.tensor([self.height, self.width])
|
|
|
|
@classmethod
|
|
def generate_rays(
|
|
cls,
|
|
intrinsics: TensorType["num_rays", 2],
|
|
camera_to_world: TensorType["num_rays", 3, 4],
|
|
coords: TensorType["num_rays", 2],
|
|
) -> RayBundle:
|
|
|
|
phi = coords[:, 0:1] # (num_rays, 1)
|
|
theta = coords[:, 1:2] # (num_rays, 1)
|
|
|
|
height = intrinsics[:, 0:1]
|
|
width = intrinsics[:, 1:2]
|
|
|
|
phi = phi / height * torch.pi
|
|
theta = -theta / width * 2 * torch.pi
|
|
|
|
directions = torch.stack(
|
|
[torch.cos(theta) * torch.sin(phi), torch.sin(theta) * torch.sin(phi), torch.cos(phi)], dim=-1
|
|
)
|
|
origins = camera_to_world[:, :3, 3] # (num_rays, 3)
|
|
return RayBundle(origins=origins, directions=directions)
|
|
|
|
|
|
def get_camera_model(num_intrinsics_params: int) -> Type[Camera]:
|
|
"""Returns the camera model given the specified number of intrinsics parameters.
|
|
|
|
Args:
|
|
num_intrinsics_params (int): Number of intrinsic parametes.
|
|
|
|
Returns:
|
|
Camera: Camera model
|
|
"""
|
|
if num_intrinsics_params == 2:
|
|
return EquirectangularCamera
|
|
if num_intrinsics_params == 3:
|
|
return SimplePinholeCamera
|
|
if num_intrinsics_params == 4:
|
|
return PinholeCamera
|
|
raise NotImplementedError
|