def get_renderer(resolution, n_pts_per_ray):
# Changing the rendering resolution is a bit involved.
raysampler = NDCGridRaysampler(
image_width=resolution,
image_height=resolution,
n_pts_per_ray=n_pts_per_ray,
min_depth=args.camera_radius -
args.volume_extent_world * np.sqrt(3) / 2,
max_depth=args.camera_radius +
args.volume_extent_world * np.sqrt(3) / 2,
)
raymarcher = EmissionAbsorptionRaymarcher()
renderer = VolumeRenderer(raysampler=raysampler, raymarcher=raymarcher)
return renderer
def renderer(
volume_size=(25, 25, 25),
batch_size=10,
shape="sphere",
raymarcher_type=EmissionAbsorptionRaymarcher,
n_rays_per_image=10,
n_pts_per_ray=10,
):
# get the volumes
volumes = init_boundary_volume(volume_size=volume_size,
batch_size=batch_size,
shape=shape)[0]
# init the mc raysampler
raysampler = MonteCarloRaysampler(
min_x=-1.0,
max_x=1.0,
min_y=-1.0,
max_y=1.0,
n_rays_per_image=n_rays_per_image,
n_pts_per_ray=n_pts_per_ray,
min_depth=0.1,
max_depth=2.0,
).to(volumes.device)
# get the raymarcher
raymarcher = raymarcher_type()
renderer = VolumeRenderer(raysampler=raysampler,
raymarcher=raymarcher,
sample_mode="bilinear")
# generate NDC camera extrinsics and intrinsics
cameras = init_cameras(batch_size, image_size=None, ndc=True)
def run_renderer():
renderer(cameras=cameras, volumes=volumes)
return run_renderer
def test_rotating_cube_volume_render(self):
"""
Generates 4 renders of 4 sides of a volume representing a 3D cube.
Since each side of the cube is homogeneously colored with
a different color, this should result in 4 images of homogeneous color
with the depth of each pixel equal to a constant.
"""
# batch_size = 4 sides of the cube
batch_size = 4
image_size = (50, 40)
for volume_size in ([25, 25, 25], ):
for sample_mode in ("bilinear", "nearest"):
volume_translation = torch.zeros(4, 3)
volume_translation.requires_grad = True
volumes, volume_voxel_size, _ = init_boundary_volume(
volume_size=volume_size,
batch_size=batch_size,
shape="cube",
volume_translation=volume_translation,
)
# generate camera extrinsics and intrinsics
cameras = init_cameras(batch_size, image_size=image_size)
# enable the gradient caching for the camera variables
# the list of differentiable camera vars
cam_vars = ("R", "T", "focal_length", "principal_point")
for cam_var in cam_vars:
getattr(cameras, cam_var).requires_grad = True
# enable the grad for volume vars as well
volumes.features().requires_grad = True
volumes.densities().requires_grad = True
raysampler = MultinomialRaysampler(
min_x=0.5,
max_x=image_size[1] - 0.5,
min_y=0.5,
max_y=image_size[0] - 0.5,
image_width=image_size[1],
image_height=image_size[0],
n_pts_per_ray=128,
min_depth=0.01,
max_depth=3.0,
)
raymarcher = EmissionAbsorptionRaymarcher()
renderer = VolumeRenderer(
raysampler=raysampler,
raymarcher=raymarcher,
sample_mode=sample_mode,
)
images_opacities = renderer(cameras=cameras,
volumes=volumes)[0]
images, opacities = images_opacities[
..., :3], images_opacities[..., 3]
# check that the renderer does not erase gradients
loss = images_opacities.sum()
loss.backward()
for check_var in (
*[getattr(cameras, cam_var) for cam_var in cam_vars],
volumes.features(),
volumes.densities(),
volume_translation,
):
self.assertIsNotNone(check_var.grad)
# ao opacities should be exactly the same as the ea ones
# we can further get the ea opacities from a feature-less
# version of our volumes
raymarcher_ao = AbsorptionOnlyRaymarcher()
renderer_ao = VolumeRenderer(
raysampler=raysampler,
raymarcher=raymarcher_ao,
sample_mode=sample_mode,
)
volumes_featureless = Volumes(
densities=volumes.densities(),
volume_translation=volume_translation,
voxel_size=volume_voxel_size,
)
opacities_ao = renderer_ao(cameras=cameras,
volumes=volumes_featureless)[0][...,
0]
self.assertClose(opacities, opacities_ao)
# colors of the sides of the cube
gt_clr_sides = torch.tensor(
[
[1.0, 0.0, 0.0],
[0.0, 1.0, 1.0],
[1.0, 1.0, 1.0],
[0.0, 1.0, 0.0],
],
dtype=torch.float32,
device=images.device,
)
if DEBUG:
outdir = tempfile.gettempdir() + "/test_volume_renderer"
os.makedirs(outdir, exist_ok=True)
for imidx, (image,
opacity) in enumerate(zip(images, opacities)):
for image_ in (image, opacity):
image_pil = Image.fromarray(
(image_.detach().cpu().numpy() * 255.0).astype(
np.uint8))
outfile = (outdir + f"/rgb_{sample_mode}" +
f"_{str(volume_size).replace(' ','')}" +
f"_{imidx:003d}")
if image_ is image:
outfile += "_rgb.png"
else:
outfile += "_opacity.png"
image_pil.save(outfile)
print(f"exported {outfile}")
border = 10
for image, opacity, gt_color in zip(images, opacities,
gt_clr_sides):
image_crop = image[border:-border, border:-border]
opacity_crop = opacity[border:-border, border:-border]
# check mean and std difference from gt
err = ((
image_crop -
gt_color[None, None].expand_as(image_crop)).abs().mean(
dim=-1))
zero = err.new_zeros(1)[0]
self.assertClose(err.mean(), zero, atol=1e-2)
self.assertClose(err.std(), zero, atol=1e-2)
err_opacity = (opacity_crop - 1.0).abs()
self.assertClose(err_opacity.mean(), zero, atol=1e-2)
self.assertClose(err_opacity.std(), zero, atol=1e-2)
def _rotating_gif(self,
image_size,
n_frames=50,
fps=15,
volume_size=(100, 100, 100)):
"""
Render a gif animation of a rotating cube/sphere (runs only if `DEBUG==True`).
"""
if not DEBUG:
# do not run this if debug is False
return
for shape in ("sphere", "cube"):
for sample_mode in ("bilinear", "nearest"):
volumes = init_boundary_volume(volume_size=volume_size,
batch_size=n_frames,
shape=shape)[0]
# generate camera extrinsics and intrinsics
cameras = init_cameras(n_frames, image_size=image_size)
# init the grid raysampler
raysampler = MultinomialRaysampler(
min_x=0.5,
max_x=image_size[1] - 0.5,
min_y=0.5,
max_y=image_size[0] - 0.5,
image_width=image_size[1],
image_height=image_size[0],
n_pts_per_ray=256,
min_depth=0.5,
max_depth=2.0,
)
# get the EA raymarcher
raymarcher = EmissionAbsorptionRaymarcher()
# initialize the renderer
renderer = VolumeRenderer(
raysampler=raysampler,
raymarcher=raymarcher,
sample_mode=sample_mode,
)
# run the renderer
images_opacities = renderer(cameras=cameras,
volumes=volumes)[0]
# split output to the alpha channel and rendered images
images, opacities = images_opacities[
..., :3], images_opacities[..., 3]
# export the gif
outdir = tempfile.gettempdir() + "/test_volume_renderer_gifs"
os.makedirs(outdir, exist_ok=True)
frames = []
for image, opacity in zip(images, opacities):
image_pil = Image.fromarray(
(torch.cat((image, opacity[..., None].repeat(1, 1, 3)),
dim=1).detach().cpu().numpy() *
255.0).astype(np.uint8))
frames.append(image_pil)
outfile = os.path.join(outdir, f"{shape}_{sample_mode}.gif")
frames[0].save(
outfile,
save_all=True,
append_images=frames[1:],
duration=n_frames // fps,
loop=0,
)
print(f"exported {outfile}")
def test_monte_carlo_rendering(self,
n_frames=20,
volume_size=(30, 30, 30),
image_size=(40, 50)):
"""
Tests that rendering with the MonteCarloRaysampler matches the
rendering with MultinomialRaysampler sampled at the corresponding
MonteCarlo locations.
"""
volumes = init_boundary_volume(volume_size=volume_size,
batch_size=n_frames,
shape="sphere")[0]
# generate camera extrinsics and intrinsics
cameras = init_cameras(n_frames, image_size=image_size)
# init the grid raysampler
raysampler_multinomial = MultinomialRaysampler(
min_x=0.5,
max_x=image_size[1] - 0.5,
min_y=0.5,
max_y=image_size[0] - 0.5,
image_width=image_size[1],
image_height=image_size[0],
n_pts_per_ray=256,
min_depth=0.5,
max_depth=2.0,
)
# init the mc raysampler
raysampler_mc = MonteCarloRaysampler(
min_x=0.5,
max_x=image_size[1] - 0.5,
min_y=0.5,
max_y=image_size[0] - 0.5,
n_rays_per_image=3000,
n_pts_per_ray=256,
min_depth=0.5,
max_depth=2.0,
)
# get the EA raymarcher
raymarcher = EmissionAbsorptionRaymarcher()
# get both mc and grid renders
(
(images_opacities_mc, ray_bundle_mc),
(images_opacities_grid, ray_bundle_grid),
) = [
VolumeRenderer(
raysampler=raysampler_multinomial,
raymarcher=raymarcher,
sample_mode="bilinear",
)(cameras=cameras, volumes=volumes)
for raysampler in (raysampler_mc, raysampler_multinomial)
]
# convert the mc sampling locations to [-1, 1]
sample_loc = ray_bundle_mc.xys.clone()
sample_loc[..., 0] = 2 * (sample_loc[..., 0] / image_size[1]) - 1
sample_loc[..., 1] = 2 * (sample_loc[..., 1] / image_size[0]) - 1
# sample the grid render at the mc locations
images_opacities_mc_ = torch.nn.functional.grid_sample(
images_opacities_grid.permute(0, 3, 1, 2),
sample_loc,
align_corners=False)
# check that the samples are the same
self.assertClose(images_opacities_mc.permute(0, 3, 1, 2),
images_opacities_mc_,
atol=1e-4)
def test_compare_with_pointclouds_renderer(self,
batch_size=11,
volume_size=(30, 30, 30),
image_size=(200, 250)):
"""
Generate a volume and its corresponding point cloud and check whether
PointsRenderer returns the same images as the corresponding VolumeRenderer.
"""
# generate NDC camera extrinsics and intrinsics
cameras = init_cameras(batch_size, image_size=image_size, ndc=True)
# init the boundary volume
for shape in ("sphere", "cube"):
if not DEBUG and shape == "cube":
# do not run numeric checks for the cube as the
# differences in rendering equations make the renders incomparable
continue
# get rand offset of the volume
volume_translation = torch.randn(batch_size, 3) * 0.1
# volume_translation[2] = 0.1
volumes = init_boundary_volume(
volume_size=volume_size,
batch_size=batch_size,
shape=shape,
volume_translation=volume_translation,
)[0]
# convert the volumes to a pointcloud
points = []
points_features = []
for densities_one, features_one, grid_one in zip(
volumes.densities(),
volumes.features(),
volumes.get_coord_grid(world_coordinates=True),
):
opaque = densities_one.view(-1) > 1e-4
points.append(grid_one.view(-1, 3)[opaque])
points_features.append(features_one.reshape(3, -1).t()[opaque])
pointclouds = Pointclouds(points, features=points_features)
# init the grid raysampler with the ndc grid
coord_range = 1.0
half_pix_size = coord_range / max(*image_size)
raysampler = NDCMultinomialRaysampler(
image_width=image_size[1],
image_height=image_size[0],
n_pts_per_ray=256,
min_depth=0.1,
max_depth=2.0,
)
# get the EA raymarcher
raymarcher = EmissionAbsorptionRaymarcher()
# jitter the camera intrinsics a bit for each render
cameras_randomized = cameras.clone()
cameras_randomized.principal_point = (
torch.randn_like(cameras.principal_point) * 0.3)
cameras_randomized.focal_length = (
cameras.focal_length +
torch.randn_like(cameras.focal_length) * 0.2)
# get the volumetric render
images = VolumeRenderer(raysampler=raysampler,
raymarcher=raymarcher,
sample_mode="bilinear")(
cameras=cameras_randomized,
volumes=volumes)[0][..., :3]
# instantiate the points renderer
point_radius = 6 * half_pix_size
points_renderer = PointsRenderer(
rasterizer=PointsRasterizer(
cameras=cameras_randomized,
raster_settings=PointsRasterizationSettings(
image_size=image_size,
radius=point_radius,
points_per_pixel=10),
),
compositor=AlphaCompositor(),
)
# get the point render
images_pts = points_renderer(pointclouds)
if shape == "sphere":
diff = (images - images_pts).abs().mean(dim=-1)
mu_diff = diff.mean(dim=(1, 2))
std_diff = diff.std(dim=(1, 2))
self.assertClose(mu_diff, torch.zeros_like(mu_diff), atol=3e-2)
self.assertClose(std_diff,
torch.zeros_like(std_diff),
atol=6e-2)
if DEBUG:
outdir = tempfile.gettempdir() + "/test_volume_vs_pts_renderer"
os.makedirs(outdir, exist_ok=True)
frames = []
for (image, image_pts) in zip(images, images_pts):
diff_image = (((image - image_pts) * 0.5 + 0.5).mean(
dim=2, keepdim=True).repeat(1, 1, 3))
image_pil = Image.fromarray(
(torch.cat((image, image_pts, diff_image),
dim=1).detach().cpu().numpy() *
255.0).astype(np.uint8))
frames.append(image_pil)
# export gif
outfile = os.path.join(outdir,
f"volume_vs_pts_render_{shape}.gif")
frames[0].save(
outfile,
save_all=True,
append_images=frames[1:],
duration=batch_size // 15,
loop=0,
)
print(f"exported {outfile}")
# export concatenated frames
outfile_cat = os.path.join(
outdir, f"volume_vs_pts_render_{shape}.png")
Image.fromarray(
np.concatenate([np.array(f) for f in frames],
axis=0)).save(outfile_cat)
print(f"exported {outfile_cat}")
def test_input_types(self, batch_size: int = 10):
"""
Check that ValueErrors are thrown where expected.
"""
# check the constructor
for bad_raysampler in (None, 5, []):
for bad_raymarcher in (None, 5, []):
with self.assertRaises(ValueError):
VolumeRenderer(raysampler=bad_raysampler,
raymarcher=bad_raymarcher)
raysampler = NDCMultinomialRaysampler(
image_width=100,
image_height=100,
n_pts_per_ray=10,
min_depth=0.1,
max_depth=1.0,
)
# init a trivial renderer
renderer = VolumeRenderer(raysampler=raysampler,
raymarcher=EmissionAbsorptionRaymarcher())
# get cameras
cameras = init_cameras(batch_size=batch_size)
# get volumes
volumes = init_boundary_volume(volume_size=(10, 10, 10),
batch_size=batch_size)[0]
# different batch sizes for cameras / volumes
with self.assertRaises(ValueError):
renderer(cameras=cameras, volumes=volumes[:-1])
# ray checks for VolumeSampler
volume_sampler = VolumeSampler(volumes=volumes)
n_rays = 100
for bad_ray_bundle in (
(
torch.rand(batch_size, n_rays, 3),
torch.rand(batch_size, n_rays + 1, 3),
torch.rand(batch_size, n_rays, 10),
),
(
torch.rand(batch_size + 1, n_rays, 3),
torch.rand(batch_size, n_rays, 3),
torch.rand(batch_size, n_rays, 10),
),
(
torch.rand(batch_size, n_rays, 3),
torch.rand(batch_size, n_rays, 2),
torch.rand(batch_size, n_rays, 10),
),
(
torch.rand(batch_size, n_rays, 3),
torch.rand(batch_size, n_rays, 3),
torch.rand(batch_size, n_rays),
),
):
ray_bundle = RayBundle(
**dict(
zip(
("origins", "directions", "lengths"),
[r.to(cameras.device) for r in bad_ray_bundle],
)),
xys=None,
)
with self.assertRaises(ValueError):
volume_sampler(ray_bundle)
# check also explicitly the ray bundle validation function
with self.assertRaises(ValueError):
_validate_ray_bundle_variables(*bad_ray_bundle)
args.volume_extent_world * np.sqrt(3) / 2,
max_depth=args.camera_radius +
args.volume_extent_world * np.sqrt(3) / 2,
)
# 2) Instantiate the raymarcher.
# Here, we use the standard EmissionAbsorptionRaymarcher
# which marches along each ray in order to render
# each ray into a single 3D color vector
# and an opacity scalar.
raymarcher = EmissionAbsorptionRaymarcher()
# Finally, instantiate the volumetric render
# with the raysampler and raymarcher objects.
renderer = VolumeRenderer(raysampler=raysampler,
raymarcher=raymarcher,
sample_mode=args.sample_mode)
## 5. Initialize the volumetric model.
if args.optimize_pixels:
volume_model = torch.rand((args.batch_size, 128, 128, 4),
device=device,
requires_grad=True)
else:
# Instantiate the volumetric model.
# We use a cubical volume with the size of
# one side = 128. The size of each voxel of the volume
# is set to args.volume_extent_world / volume_size s.t. the
# volume represents the space enclosed in a 3D bounding box
# centered at (0, 0, 0) with the size of each side equal to 3.
volume_model = VolumeModel(