def test_raysampler(self):
data = _CommonData()
gridsampler = NDCMultinomialRaysampler(
image_width=data.W,
image_height=data.H,
n_pts_per_ray=2,
min_depth=1.0,
max_depth=2.0,
)
for camera in (data.camera_ndc, data.camera_screen):
bundle = gridsampler(camera)
self.assertTupleEqual(bundle.xys.shape,
(2, ) + data.image_size + (2, ))
self.assertTupleEqual(bundle.directions.shape,
(2, ) + data.image_size + (3, ))
self.assertClose(
bundle.xys[:, data.y, data.x],
torch.tensor(data.point[:2]).expand(2, -1),
)
# We check only the first batch element.
# Second element varies because of camera location.
self.assertClose(
bundle.directions[0, data.y, data.x],
torch.tensor(data.point),
)
def test_input_types(self):
"""
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):
ImplicitRenderer(raysampler=bad_raysampler,
raymarcher=bad_raymarcher)
# init a trivial renderer
renderer = ImplicitRenderer(
raysampler=NDCMultinomialRaysampler(
image_width=100,
image_height=100,
n_pts_per_ray=10,
min_depth=0.1,
max_depth=1.0,
),
raymarcher=EmissionAbsorptionRaymarcher(),
)
# get default cameras
cameras = init_cameras()
for bad_volumetric_function in (None, 5, []):
with self.assertRaises(ValueError):
renderer(cameras=cameras,
volumetric_function=bad_volumetric_function)
def __init__(
self,
n_pts_per_ray: int,
min_depth: float,
max_depth: float,
n_rays_per_image: int,
image_width: int,
image_height: int,
stratified: bool = False,
stratified_test: bool = False,
):
"""
Args:
n_pts_per_ray: The number of points sampled along each ray.
min_depth: The minimum depth of a ray-point.
max_depth: The maximum depth of a ray-point.
n_rays_per_image: Number of Monte Carlo ray samples when training
(`self.training==True`).
image_width: The horizontal size of the image grid.
image_height: The vertical size of the image grid.
stratified: If `True`, stratifies (=randomly offsets) the depths
of each ray point during training (`self.training==True`).
stratified_test: If `True`, stratifies (=randomly offsets) the depths
of each ray point during evaluation (`self.training==False`).
"""
super().__init__()
self._stratified = stratified
self._stratified_test = stratified_test
# Initialize the grid ray sampler.
self._grid_raysampler = NDCMultinomialRaysampler(
image_width=image_width,
image_height=image_height,
n_pts_per_ray=n_pts_per_ray,
min_depth=min_depth,
max_depth=max_depth,
)
# Initialize the Monte Carlo ray sampler.
self._mc_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=min_depth,
max_depth=max_depth,
)
# create empty ray cache
self._ray_cache = {}
def __post_init__(self):
super().__init__()
self.scene_center = torch.FloatTensor(self.scene_center)
self._sampling_mode = {
EvaluationMode.TRAINING:
RenderSamplingMode(self.sampling_mode_training),
EvaluationMode.EVALUATION:
RenderSamplingMode(self.sampling_mode_evaluation),
}
self._raysamplers = {
EvaluationMode.TRAINING:
NDCMultinomialRaysampler(
image_width=self.image_width,
image_height=self.image_height,
n_pts_per_ray=self.n_pts_per_ray_training,
min_depth=self.min_depth,
max_depth=self.max_depth,
n_rays_per_image=self.n_rays_per_image_sampled_from_mask
if self._sampling_mode[EvaluationMode.TRAINING]
== RenderSamplingMode.MASK_SAMPLE else None,
unit_directions=True,
stratified_sampling=self.stratified_point_sampling_training,
),
EvaluationMode.EVALUATION:
NDCMultinomialRaysampler(
image_width=self.image_width,
image_height=self.image_height,
n_pts_per_ray=self.n_pts_per_ray_evaluation,
min_depth=self.min_depth,
max_depth=self.max_depth,
n_rays_per_image=self.n_rays_per_image_sampled_from_mask
if self._sampling_mode[EvaluationMode.EVALUATION]
== RenderSamplingMode.MASK_SAMPLE else None,
unit_directions=True,
stratified_sampling=self.stratified_point_sampling_evaluation,
),
}
def get_rgbd_point_cloud(
camera: CamerasBase,
image_rgb: torch.Tensor,
depth_map: torch.Tensor,
mask: Optional[torch.Tensor] = None,
mask_thr: float = 0.5,
mask_points: bool = True,
) -> Pointclouds:
"""
Given a batch of images, depths, masks and cameras, generate a colored
point cloud by unprojecting depth maps to the and coloring with the source
pixel colors.
"""
imh, imw = image_rgb.shape[2:]
# convert the depth maps to point clouds using the grid ray sampler
pts_3d = ray_bundle_to_ray_points(
NDCMultinomialRaysampler(
image_width=imw,
image_height=imh,
n_pts_per_ray=1,
min_depth=1.0,
max_depth=1.0,
)(camera)._replace(lengths=depth_map[:, 0, ..., None])
)
pts_mask = depth_map > 0.0
if mask is not None:
pts_mask *= mask > mask_thr
pts_mask = pts_mask.reshape(-1)
pts_3d = pts_3d.reshape(-1, 3)[pts_mask]
pts_colors = torch.nn.functional.interpolate(
image_rgb,
# pyre-fixme[6]: Expected `Optional[int]` for 2nd param but got
# `List[typing.Any]`.
size=[imh, imw],
mode="bilinear",
align_corners=False,
)
pts_colors = pts_colors.permute(0, 2, 3, 1).reshape(-1, 3)[pts_mask]
return Pointclouds(points=pts_3d[None], features=pts_colors[None])
def _compare_with_meshes_renderer(self,
image_size,
batch_size=11,
sphere_diameter=0.6):
"""
Generate a spherical RGB volumetric function and its corresponding mesh
and check whether MeshesRenderer returns the same images as the
corresponding ImplicitRenderer.
"""
# generate NDC camera extrinsics and intrinsics
cameras = init_cameras(batch_size, image_size=image_size, ndc=True)
# get rand offset of the volume
sphere_centroid = torch.randn(batch_size, 3,
device=cameras.device) * 0.1
sphere_centroid.requires_grad = True
# init the grid raysampler with the ndc grid
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)
# the list of differentiable camera vars
cam_vars = ("R", "T", "focal_length", "principal_point")
# enable the gradient caching for the camera variables
for cam_var in cam_vars:
getattr(cameras_randomized, cam_var).requires_grad = True
# get the implicit renderer
images_opacities = ImplicitRenderer(
raysampler=raysampler, raymarcher=raymarcher)(
cameras=cameras_randomized,
volumetric_function=spherical_volumetric_function,
sphere_centroid=sphere_centroid,
sphere_diameter=sphere_diameter,
)[0]
# check that the renderer does not erase gradients
loss = images_opacities.sum()
loss.backward()
for check_var in (
*[
getattr(cameras_randomized, cam_var)
for cam_var in cam_vars
],
sphere_centroid,
):
self.assertIsNotNone(check_var.grad)
# instantiate the corresponding spherical mesh
ico = ico_sphere(level=4, device=cameras.device).extend(batch_size)
verts = (torch.nn.functional.normalize(ico.verts_padded(), dim=-1) *
sphere_diameter + sphere_centroid[:, None])
meshes = Meshes(
verts=verts,
faces=ico.faces_padded(),
textures=TexturesVertex(verts_features=(
torch.nn.functional.normalize(verts, dim=-1) * 0.5 + 0.5)),
)
# instantiate the corresponding mesh renderer
lights = PointLights(device=cameras.device, location=[[0.0, 0.0, 0.0]])
renderer_textured = MeshRenderer(
rasterizer=MeshRasterizer(
cameras=cameras_randomized,
raster_settings=RasterizationSettings(
image_size=image_size,
blur_radius=1e-3,
faces_per_pixel=10,
z_clip_value=None,
perspective_correct=False,
),
),
shader=SoftPhongShader(
device=cameras.device,
cameras=cameras_randomized,
lights=lights,
materials=Materials(
ambient_color=((2.0, 2.0, 2.0), ),
diffuse_color=((0.0, 0.0, 0.0), ),
specular_color=((0.0, 0.0, 0.0), ),
shininess=64,
device=cameras.device,
),
blend_params=BlendParams(sigma=1e-3,
gamma=1e-4,
background_color=(0.0, 0.0, 0.0)),
),
)
# get the mesh render
images_opacities_meshes = renderer_textured(meshes,
cameras=cameras_randomized,
lights=lights)
if DEBUG:
outdir = tempfile.gettempdir() + "/test_implicit_vs_mesh_renderer"
os.makedirs(outdir, exist_ok=True)
frames = []
for (image_opacity,
image_opacity_mesh) in zip(images_opacities,
images_opacities_meshes):
image, opacity = image_opacity.split([3, 1], dim=-1)
image_mesh, opacity_mesh = image_opacity_mesh.split([3, 1],
dim=-1)
diff_image = (((image - image_mesh) * 0.5 + 0.5).mean(
dim=2, keepdim=True).repeat(1, 1, 3))
image_pil = Image.fromarray((torch.cat(
(
image,
image_mesh,
diff_image,
opacity.repeat(1, 1, 3),
opacity_mesh.repeat(1, 1, 3),
),
dim=1,
).detach().cpu().numpy() * 255.0).astype(np.uint8))
frames.append(image_pil)
# export gif
outfile = os.path.join(outdir, "implicit_vs_mesh_render.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, "implicit_vs_mesh_render.png")
Image.fromarray(
np.concatenate([np.array(f) for f in frames],
axis=0)).save(outfile_cat)
print(f"exported {outfile_cat}")
# compare the renders
diff = (images_opacities - images_opacities_meshes).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=5e-2)
self.assertClose(std_diff, torch.zeros_like(std_diff), atol=6e-2)
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)