def save(
self,
data: Pointclouds,
path: Union[str, Path],
path_manager: PathManager,
binary: Optional[bool],
decimal_places: Optional[int] = None,
**kwargs,
) -> bool:
if not endswith(path, self.known_suffixes):
return False
points = data.points_list()[0]
features = data.features_list()[0]
with _open_file(path, path_manager, "wb") as f:
_save_ply(
f=f,
verts=points,
verts_colors=features,
verts_normals=torch.FloatTensor([]),
faces=None,
ascii=binary is False,
decimal_places=decimal_places,
)
return True
def forward(self, pts, features):
# Make sure pts and features are equal
assert pts.size(2) == 3
assert pts.size(1) == features.size(1)
pts[:, :, 0] = -pts[:, :, 0]
pts[:, :, 1] = -pts[:, :, 1]
radius = float(self.radius) / float(self.img_size) * 2.0
params = compositing.CompositeParams(radius=radius)
pointcloud = Pointclouds(points=pts, features=features)
points_idx, _, dist = rasterize_points(pointcloud, self.img_size,
radius, self.points_per_pixel)
dist = dist / pow(radius, 2)
alphas = (1 - dist.clamp(max=1, min=1e-3).pow(0.5).pow(
self.gamma).permute(0, 3, 1, 2))
transformed_feature_alphas = compositing.alpha_composite(
points_idx.permute(0, 3, 1, 2).long(), alphas,
pointcloud.features_packed().permute(1, 0), params)
return transformed_feature_alphas
def forward(self, pts3D, src):
bs = src.size(0)
if len(src.size()) > 3:
bs, c, w, _ = src.size()
image_size = w
pts3D = pts3D.permute(0, 2, 1)
src = src.unsqueeze(2).repeat(1, 1, w, 1, 1).view(bs, c, -1)
else:
bs = src.size(0)
image_size = self.size
# Make sure these have been arranged in the same way
assert pts3D.size(2) == 3
assert pts3D.size(1) == src.size(2)
pts3D[:, :, 1] = -pts3D[:, :, 1]
pts3D[:, :, 0] = -pts3D[:, :, 0]
# Add on the default feature to the end of the src
# src = torch.cat((src, self.default_feature.repeat(bs, 1, 1)), 2)
radius = float(self.radius) / float(image_size) * 2.0
pts3D = Pointclouds(points=pts3D, features=src.permute(0, 2, 1))
points_idx, _, dist = rasterize_points(pts3D, image_size, radius,
self.points_per_pixel)
if os.environ["DEBUG"]:
print("Max dist: ", dist.max(), pow(radius, self.opts.rad_pow))
dist = dist / pow(radius, self.opts.rad_pow)
if os.environ["DEBUG"]:
print("Max dist: ", dist.max())
alphas = ((1 - dist.clamp(max=1, min=1e-3).pow(0.5)).pow(
self.opts.tau).permute(0, 3, 1, 2))
if self.opts.accumulation == 'alphacomposite':
transformed_src_alphas = compositing.alpha_composite(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(1, 0),
)
elif self.opts.accumulation == 'wsum':
transformed_src_alphas = compositing.weighted_sum(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(1, 0),
)
elif self.opts.accumulation == 'wsumnorm':
transformed_src_alphas = compositing.weighted_sum_norm(
points_idx.permute(0, 3, 1, 2).long(),
alphas,
pts3D.features_packed().permute(1, 0),
)
return transformed_src_alphas
def _clone_pointcloud(verts0, device, batch_size):
"""
Helper function to detach and clone the verts.
This is needed in order to set up the tensors for
gradient computation in different tests.
"""
verts = verts0.detach().clone()
verts.requires_grad = True
pointclouds = Pointclouds(points=[verts])
pointclouds = pointclouds.to(device).extend(batch_size)
return verts, pointclouds
def __call__(self, points, translate=True):
assert len(points.shape) == 3 and points.shape[-1] == 3
bs = points.size(0)
rgb = torch.ones((bs, points.size(1), 3),
device=points.device) * self.max_brightness
if translate:
tm = points.mean(dim=-2, keepdim=False)
T = T3.Translate(-tm, device=points.device)
points = T.transform_points(points)
# There's error on normals
# Probably not needed on just translation
# normals = T.transform_normals(normals)
cloud = Pointclouds(points=points, features=rgb)
gamma = [self.opt.pulsar_gamma for _ in range(bs)]
cameras = self.renderer.rasterizer.cameras
R = cameras.R.expand(bs, -1, -1)
T = cameras.T.expand(bs, -1)
znear = cameras.znear.expand(bs)
zfar = cameras.zfar.expand(bs)
fov = cameras.fov.expand(bs)
aspect_ratio = cameras.aspect_ratio.expand(bs)
return self.renderer(cloud,
gamma=gamma,
R=R,
T=T,
znear=znear,
zfar=zfar,
fov=fov,
aspect_ratio=aspect_ratio)
def test_points_image_size_arg(self):
points = Pointclouds([verts0])
with self.assertRaises(ValueError) as cm:
rasterize_points(
points,
(100, 200, 3),
0.0001,
points_per_pixel=1,
)
self.assertTrue("tuple/list of (H, W)" in cm.msg)
with self.assertRaises(ValueError) as cm:
rasterize_points(
points,
(0, 10),
0.0001,
points_per_pixel=1,
)
self.assertTrue("sizes must be positive" in cm.msg)
with self.assertRaises(ValueError) as cm:
rasterize_points(
points,
(100.5, 120.5),
0.0001,
points_per_pixel=1,
)
self.assertTrue("sizes must be integers" in cm.msg)
def visualize_nerf_outputs(nerf_out: dict, output_cache: List, viz: Visdom,
visdom_env: str):
"""
Visualizes the outputs of the `RadianceFieldRenderer`.
Args:
nerf_out: An output of the validation rendering pass.
output_cache: A list with outputs of several training render passes.
viz: A visdom connection object.
visdom_env: The name of visdom environment for visualization.
"""
# Show the training images.
ims = torch.stack([o["image"] for o in output_cache])
ims = torch.cat(list(ims), dim=1)
viz.image(
ims.permute(2, 0, 1),
env=visdom_env,
win="images",
opts={"title": "train_images"},
)
# Show the coarse and fine renders together with the ground truth images.
ims_full = torch.cat(
[
nerf_out[imvar][0].permute(2, 0, 1).detach().cpu().clamp(0.0, 1.0)
for imvar in ("rgb_coarse", "rgb_fine", "rgb_gt")
],
dim=2,
)
viz.image(
ims_full,
env=visdom_env,
win="images_full",
opts={"title": "coarse | fine | target"},
)
# Make a 3D plot of training cameras and their emitted rays.
camera_trace = {
f"camera_{ci:03d}": o["camera"].cpu()
for ci, o in enumerate(output_cache)
}
ray_pts_trace = {
f"ray_pts_{ci:03d}": Pointclouds(
ray_bundle_to_ray_points(
o["coarse_ray_bundle"]).detach().cpu().view(1, -1, 3))
for ci, o in enumerate(output_cache)
}
plotly_plot = plot_scene(
{
"training_scene": {
**camera_trace,
**ray_pts_trace,
},
},
pointcloud_max_points=5000,
pointcloud_marker_size=1,
camera_scale=0.3,
)
viz.plotlyplot(plotly_plot, env=visdom_env, win="scenes")
def test_render_pointcloud(self):
"""
Test a textured point cloud is rendered correctly in a non square image.
"""
device = torch.device("cuda:0")
pointclouds = Pointclouds(
points=[torus_points * 2.0],
features=torch.ones_like(torus_points[None, ...]),
).to(device)
R, T = look_at_view_transform(2.7, 0.0, 0.0)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
raster_settings = PointsRasterizationSettings(
image_size=(512, 1024), radius=5e-2, points_per_pixel=1
)
rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings)
compositor = AlphaCompositor()
renderer = PointsRenderer(rasterizer=rasterizer, compositor=compositor)
# Load reference image
image_ref = load_rgb_image("test_pointcloud_rectangle_image.png", DATA_DIR)
for bin_size in [0, None]:
# Check both naive and coarse to fine produce the same output.
renderer.rasterizer.raster_settings.bin_size = bin_size
images = renderer(pointclouds)
rgb = images[0, ..., :3].squeeze().cpu()
if DEBUG:
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_pointcloud_rectangle_image.png"
)
# NOTE some pixels can be flaky
cond1 = torch.allclose(rgb, image_ref, atol=0.05)
self.assertTrue(cond1)
def test_texture_sampling_cow(self):
# test texture sampling for the cow example by converting
# the cow mesh and its texture uv to a pointcloud with texture
device = torch.device("cuda:0")
obj_dir = get_pytorch3d_dir() / "docs/tutorials/data"
obj_filename = obj_dir / "cow_mesh/cow.obj"
for text_type in ("uv", "atlas"):
# Load mesh + texture
if text_type == "uv":
mesh = load_objs_as_meshes(
[obj_filename], device=device, load_textures=True, texture_wrap=None
)
elif text_type == "atlas":
mesh = load_objs_as_meshes(
[obj_filename],
device=device,
load_textures=True,
create_texture_atlas=True,
texture_atlas_size=8,
texture_wrap=None,
)
points, normals, textures = sample_points_from_meshes(
mesh, num_samples=50000, return_normals=True, return_textures=True
)
pointclouds = Pointclouds(points, normals=normals, features=textures)
for pos in ("front", "back"):
# Init rasterizer settings
if pos == "back":
azim = 0.0
elif pos == "front":
azim = 180
R, T = look_at_view_transform(2.7, 0, azim)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
raster_settings = PointsRasterizationSettings(
image_size=512, radius=1e-2, points_per_pixel=1
)
rasterizer = PointsRasterizer(
cameras=cameras, raster_settings=raster_settings
)
compositor = NormWeightedCompositor()
renderer = PointsRenderer(rasterizer=rasterizer, compositor=compositor)
images = renderer(pointclouds)
rgb = images[0, ..., :3].squeeze().cpu()
if DEBUG:
filename = "DEBUG_cow_mesh_to_pointcloud_%s_%s.png" % (
text_type,
pos,
)
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / filename
)
def read(
self,
path: Union[str, Path],
device,
path_manager: PathManager,
**kwargs,
) -> Optional[Pointclouds]:
if not endswith(path, self.known_suffixes):
return None
verts, faces, features = _load_ply(
f=path, path_manager=path_manager, return_vertex_colors=True
)
verts = verts.to(device)
if features is None:
pointcloud = Pointclouds(points=[verts])
else:
pointcloud = Pointclouds(points=[verts], features=[features.to(device)])
return pointcloud
def cli():
"""
Basic example for the pulsar sphere renderer using the PyTorch3D interface.
Writes to `basic-pt3d.png`.
"""
LOGGER.info("Rendering on GPU...")
torch.manual_seed(1)
n_points = 10
width = 1_000
height = 1_000
device = torch.device("cuda")
# Generate sample data.
vert_pos = torch.rand(n_points, 3, dtype=torch.float32,
device=device) * 10.0
vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0
vert_col = torch.rand(n_points, 3, dtype=torch.float32, device=device)
pcl = Pointclouds(points=vert_pos[None, ...], features=vert_col[None, ...])
# Alternatively, you can also use the look_at_view_transform to get R and T:
# R, T = look_at_view_transform(
# dist=30.0, elev=0.0, azim=180.0, at=((0.0, 0.0, 30.0),), up=((0, 1, 0),),
# )
cameras = PerspectiveCameras(
# The focal length must be double the size for PyTorch3D because of the NDC
# coordinates spanning a range of two - and they must be normalized by the
# sensor width (see the pulsar example). This means we need here
# 5.0 * 2.0 / 2.0 to get the equivalent results as in pulsar.
focal_length=(5.0 * 2.0 / 2.0, ),
R=torch.eye(3, dtype=torch.float32, device=device)[None, ...],
T=torch.zeros((1, 3), dtype=torch.float32, device=device),
image_size=((width, height), ),
device=device,
)
vert_rad = torch.rand(n_points, dtype=torch.float32, device=device)
raster_settings = PointsRasterizationSettings(
image_size=(width, height),
radius=vert_rad,
)
rasterizer = PointsRasterizer(cameras=cameras,
raster_settings=raster_settings)
renderer = PulsarPointsRenderer(rasterizer=rasterizer).to(device)
# Render.
image = renderer(
pcl,
gamma=(1.0e-1, ), # Renderer blending parameter gamma, in [1., 1e-5].
znear=(1.0, ),
zfar=(45.0, ),
radius_world=True,
bg_col=torch.ones((3, ), dtype=torch.float32, device=device),
)[0]
LOGGER.info("Writing image to `%s`.", path.abspath("basic-pt3d.png"))
imageio.imsave("basic-pt3d.png",
(image.cpu().detach() * 255.0).to(torch.uint8).numpy())
LOGGER.info("Done.")
def test_points_renderer_to(self):
"""
Test moving all the tensors in the points renderer to a new device.
"""
device1 = torch.device("cpu")
R, T = look_at_view_transform(1500, 0.0, 0.0)
raster_settings = PointsRasterizationSettings(image_size=256,
radius=0.001,
points_per_pixel=1)
cameras = FoVPerspectiveCameras(device=device1,
R=R,
T=T,
aspect_ratio=1.0,
fov=60.0,
zfar=100)
rasterizer = PointsRasterizer(cameras=cameras,
raster_settings=raster_settings)
renderer = PointsRenderer(rasterizer=rasterizer,
compositor=AlphaCompositor())
mesh = ico_sphere(2, device1)
verts_padded = mesh.verts_padded()
pointclouds = Pointclouds(points=verts_padded,
features=torch.randn_like(verts_padded))
self._check_points_renderer_props_on_device(renderer, device1)
# Test rendering on cpu
output_images = renderer(pointclouds)
self.assertEqual(output_images.device, device1)
# Move renderer and pointclouds to another device and re render
device2 = torch.device("cuda:0")
renderer = renderer.to(device2)
pointclouds = pointclouds.to(device2)
self._check_points_renderer_props_on_device(renderer, device2)
output_images = renderer(pointclouds)
self.assertEqual(output_images.device, device2)
def _get_cube_cloud(self):
"""Get vertices of a unit-cube, with colors assigned according to
vertex coordinates."""
vertices = list(
itertools.product(*zip([-0.5, -0.5, -0.5], [0.5, 0.5, 0.5])))
vertices = np.insert(vertices, 0, [0, 0, 0], axis=0)
vertices = th.as_tensor(vertices, dtype=th.float32, device=self.device)
# Map vertices to colors. =RGB(0.25~0.75)
colors = (0.5 + 0.5 * vertices)
cloud = Pointclouds(points=vertices[None], features=colors[None])
return cloud
def test_save_load_with_normals(self):
points = torch.tensor([[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0]],
dtype=torch.float32)
normals = torch.tensor([[0, 1, 0], [1, 0, 0], [1, 4, 1], [1, 0, 0]],
dtype=torch.float32)
features = torch.rand_like(points)
for do_features, do_normals in itertools.product([True, False],
[True, False]):
cloud = Pointclouds(
points=[points],
features=[features] if do_features else None,
normals=[normals] if do_normals else None,
)
device = torch.device("cuda:0")
io = IO()
with NamedTemporaryFile(mode="w", suffix=".ply") as f:
io.save_pointcloud(cloud.cuda(), f.name)
f.flush()
cloud2 = io.load_pointcloud(f.name, device=device)
self.assertEqual(cloud2.device, device)
cloud2 = cloud2.cpu()
self.assertClose(cloud2.points_padded(), cloud.points_padded())
if do_normals:
self.assertClose(cloud2.normals_padded(),
cloud.normals_padded())
else:
self.assertIsNone(cloud.normals_padded())
self.assertIsNone(cloud2.normals_padded())
if do_features:
self.assertClose(cloud2.features_packed(), features)
else:
self.assertIsNone(cloud2.features_packed())
def test_pointcloud(self):
data = _CommonData()
clouds = Pointclouds(points=torch.tensor([[data.point]])).extend(2)
colorful_cloud = Pointclouds(points=torch.tensor([[data.point]]),
features=torch.ones(1, 1, 3)).extend(2)
points_per_pixel = 2
# for camera in [data.camera_screen]:
for camera in (data.camera_ndc, data.camera_screen):
rasterizer = PointsRasterizer(
cameras=camera,
raster_settings=PointsRasterizationSettings(
image_size=data.image_size,
radius=0.0001,
points_per_pixel=points_per_pixel,
),
)
# when rasterizing we expect only one pixel to be occupied
rasterizer_output = rasterizer(clouds).idx
self.assertTupleEqual(rasterizer_output.shape, (2, ) +
data.image_size + (points_per_pixel, ))
found = torch.nonzero(rasterizer_output != -1)
self.assertTupleEqual(found.shape, (2, 4))
self.assertListEqual(found[0].tolist(), [0, data.y, data.x, 0])
self.assertListEqual(found[1].tolist(), [1, data.y, data.x, 0])
if camera.in_ndc():
# Pulsar not currently working in screen space.
pulsar_renderer = PulsarPointsRenderer(rasterizer=rasterizer)
pulsar_output = pulsar_renderer(colorful_cloud,
gamma=(0.1, 0.1),
znear=(0.1, 0.1),
zfar=(70, 70))
self.assertTupleEqual(pulsar_output.shape,
(2, ) + data.image_size + (3, ))
# Look for points rendered in the red channel only, expecting our one.
# Check the first batch element only.
# TODO: Something is odd with the second.
found = torch.nonzero(pulsar_output[0, :, :, 0])
self.assertTupleEqual(found.shape, (1, 2))
self.assertListEqual(found[0].tolist(), [data.y, data.x])
def init_meshes_clouds(
batch_size: int = 10,
num_verts: int = 1000,
num_faces: int = 3000,
num_points: int = 3000,
device: str = "cuda:0",
):
device = torch.device(device)
nump = torch.randint(low=1, high=num_points, size=(batch_size, ))
numv = torch.randint(low=3, high=num_verts, size=(batch_size, ))
numf = torch.randint(low=1, high=num_faces, size=(batch_size, ))
verts_list = []
faces_list = []
points_list = []
for i in range(batch_size):
# Randomly choose vertices
verts = torch.rand((numv[i], 3),
dtype=torch.float32,
device=device)
verts.requires_grad_(True)
# Randomly choose faces. Our tests below compare argmin indices
# over faces and edges. Argmin is sensitive even to small numeral variations
# thus we make sure that faces are valid
# i.e. a face f = (i0, i1, i2) s.t. i0 != i1 != i2,
# otherwise argmin due to numeral sensitivities cannot be resolved
faces, allf = [], 0
validf = numv[i].item() - numv[i].item() % 3
while allf < numf[i]:
ff = torch.randperm(numv[i],
device=device)[:validf].view(-1, 3)
faces.append(ff)
allf += ff.shape[0]
faces = torch.cat(faces, 0)
if faces.shape[0] > numf[i]:
faces = faces[:numf[i]]
verts_list.append(verts)
faces_list.append(faces)
# Randomly choose points
points = torch.rand((nump[i], 3),
dtype=torch.float32,
device=device)
points.requires_grad_(True)
points_list.append(points)
meshes = Meshes(verts_list, faces_list)
pcls = Pointclouds(points_list)
return meshes, pcls
def warp_canon_depth(self, canon_depth):
b, h, w = canon_depth.shape
# grid_3d supposed to be the 3D mesh.
grid_3d = self.get_warped_3d_grid(canon_depth).reshape(b, -1, 3)
faces = get_face_idx(b, h, w).to(canon_depth.device)
# fixme: https://github.com/facebookresearch/pytorch3d/issues/202
grid_3d = Pointclouds(points=grid_3d, features=faces)
warped_depth = self.renderer(grid_3d).squeeze()
# allow some margin out of valid range
margin = (self.max_depth - self.min_depth) / 2
warped_depth = warped_depth.clamp(min=self.min_depth - margin,
max=self.max_depth + margin)
return warped_depth
def get_per_point_visibility_mask(pointclouds: Pointclouds,
fragments: NamedTuple) -> Pointclouds:
"""
compute per-point visibility (0/1), append value to pointclouds features
Returns:
boolean mask for packed tensors (P_total,)
"""
P_total = pointclouds.num_points_per_cloud().sum().item()
P_max = pointclouds.num_points_per_cloud().max().item()
try:
mask = fragments.occupancy.bool() # float
except:
mask = fragments.idx[..., 0] >= 0 # bool
pts_visibility = torch.full((P_total, ),
False,
dtype=torch.bool,
device=pointclouds.device)
# all rendered points (indices in packed points)
visible_idx = fragments.idx[mask].unique().long().view(-1)
visible_idx = visible_idx[visible_idx >= 0]
pts_visibility[visible_idx] = True
return pts_visibility
def collate_fn(batch):
out = dict()
batch_keys = batch[0].keys()
skip_keys = ['scan_dict'] # These will be manually collated
# For each not in skip_keys, use default torch collator
for key in [k for k in batch_keys if k not in skip_keys]:
out[key] = torch.utils.data._utils.collate.default_collate(
[d[key] for d in batch])
scan_all = [
torch.Tensor(sample['scan_dict']['points']) for sample in batch
]
out['scan'] = Pointclouds(points=scan_all)
return out
def rasterize_mc_samples(
xys: torch.Tensor,
feats: torch.Tensor,
image_size_hw: Tuple[int, int],
radius: float = 0.03,
topk: int = 5,
masks: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Rasterizes Monte-Carlo sampled features back onto the image.
Specifically, the code uses the PyTorch3D point rasterizer to render
a z-flat point cloud composed of the xy MC locations and their features.
Args:
xys: B x N x 2 2D point locations in PyTorch3D NDC convention
feats: B x N x dim tensor containing per-point rendered features.
image_size_hw: Tuple[image_height, image_width] containing
the size of rasterized image.
radius: Rasterization point radius.
topk: The maximum z-buffer size for the PyTorch3D point cloud rasterizer.
masks: B x N x 1 tensor containing the alpha mask of the
rendered features.
"""
if masks is None:
masks = torch.ones_like(xys[..., :1])
feats = torch.cat((feats, masks), dim=-1)
pointclouds = Pointclouds(
points=torch.cat([xys, torch.ones_like(xys[..., :1])], dim=-1),
features=feats,
)
data_rendered, render_mask, _ = render_point_cloud_pytorch3d(
PerspectiveCameras(device=feats.device),
pointclouds,
render_size=image_size_hw,
point_radius=radius,
topk=topk,
)
data_rendered, masks_pt = data_rendered.split(
[data_rendered.shape[1] - 1, 1], dim=1)
render_mask = masks_pt * render_mask
return data_rendered, render_mask
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 read(
self,
path: Union[str, Path],
device,
path_manager: PathManager,
**kwargs,
) -> Optional[Pointclouds]:
if not endswith(path, self.known_suffixes):
return None
verts, faces, features, normals = _load_ply(f=path, path_manager=path_manager)
verts = verts.to(device)
if features is not None:
features = [features.to(device)]
if normals is not None:
normals = [normals.to(device)]
pointcloud = Pointclouds(points=[verts], features=features, normals=normals)
return pointcloud
def __init__(self, dir):
print("Loading scene ", dir)
self.ReadConfig(dir)
a = numpy.loadtxt(dir + "poses_view_matrix.txt")
self.num_images = a.shape[0] // 4
self.view_matrices = []
for i in range(self.num_images):
V = a[i * 4:(i + 1) * 4, :]
self.view_matrices.append(V)
# convert view matrix to R,t
self.camera_rs = []
self.camera_ts = []
for V in self.view_matrices:
R = V[0:3, 0:3]
t = V[0:3, 3:4]
self.camera_rs.append(R)
self.camera_ts.append(t.transpose())
self.R = torch.tensor(self.camera_rs).float().to(device)
self.T = torch.tensor(self.camera_ts).squeeze(1).float().to(device)
self.image_size = (self.h, self.w)
self.image_size_tensor = torch.Tensor(
self.image_size).unsqueeze(0).to(device).int()
pc = load_ply(
"/home/dari/Projects/pointrendering2/Code/scenes/tt_train_colmap/point_cloud.ply"
)
verts = torch.Tensor(pc[0]).to(device).unsqueeze(0)
rgb = verts.clone()
rgb = torch.ones_like(verts)
# PrintTensorInfo(verts)
self.point_cloud = Pointclouds(points=verts, features=rgb)
print("Scene loaded: R, T , K , size")
PrintTensorInfo(self.R)
PrintTensorInfo(self.T)
PrintTensorInfo(self.K)
PrintTensorInfo(self.image_size_tensor)
def test_small_faces_case(self):
for device in [torch.device("cpu"), torch.device("cuda:0")]:
mesh_vertices = torch.tensor(
[
[-0.0021, -0.3769, 0.7146],
[-0.0161, -0.3771, 0.7146],
[-0.0021, -0.3771, 0.7147],
],
dtype=torch.float32,
device=device,
)
mesh1_faces = torch.tensor([[0, 2, 1]], device=device)
mesh2_faces = torch.tensor([[2, 0, 1]], device=device)
pcd_points = torch.tensor([[-0.3623, -0.5340, 0.7727]],
device=device)
mesh1 = Meshes(verts=[mesh_vertices], faces=[mesh1_faces])
mesh2 = Meshes(verts=[mesh_vertices], faces=[mesh2_faces])
pcd = Pointclouds(points=[pcd_points])
loss1 = point_mesh_face_distance(mesh1, pcd)
loss2 = point_mesh_face_distance(mesh2, pcd)
self.assertClose(loss1, loss2)
def test_save_pointcloud(self):
header = "\n".join([
"ply",
"format binary_little_endian 1.0",
"element vertex 8",
"property float x",
"property float y",
"property float z",
"property float red",
"property float green",
"property float blue",
"end_header",
"",
]).encode("ascii")
data = struct.pack("<" + "f" * 48, *range(48))
points = torch.FloatTensor([0, 1, 2]) + 6 * torch.arange(8)[:, None]
features = torch.FloatTensor([3, 4, 5]) + 6 * torch.arange(8)[:, None]
pointcloud = Pointclouds(points=[points], features=[features])
io = IO()
with NamedTemporaryFile(mode="rb", suffix=".ply") as f:
io.save_pointcloud(data=pointcloud, path=f.name)
f.flush()
f.seek(0)
actual_data = f.read()
reloaded_pointcloud = io.load_pointcloud(f.name)
self.assertEqual(header + data, actual_data)
self.assertClose(reloaded_pointcloud.points_list()[0], points)
self.assertClose(reloaded_pointcloud.features_list()[0], features)
with NamedTemporaryFile(mode="r", suffix=".ply") as f:
io.save_pointcloud(data=pointcloud, path=f.name, binary=False)
reloaded_pointcloud2 = io.load_pointcloud(f.name)
self.assertEqual(f.readline(), "ply\n")
self.assertEqual(f.readline(), "format ascii 1.0\n")
self.assertClose(reloaded_pointcloud2.points_list()[0], points)
self.assertClose(reloaded_pointcloud2.features_list()[0], features)
def eval_one_dir(exp_dir, n_pts=50000):
"""
Function for one directory
"""
device = torch.device('cuda:0')
cfg = config.load_config(os.path.join(exp_dir, 'config.yaml'))
dataset = config.create_dataset(cfg.data, mode='val')
meshes_gt = dataset.get_meshes().to(device)
val_gt_pts_file = os.path.join(cfg.data.data_dir, 'val%d.ply' % n_pts)
if os.path.isfile(val_gt_pts_file):
points, normals = np.split(read_ply(val_gt_pts_file), 2, axis=1)
pcl_gt = Pointclouds(
torch.from_numpy(points[None, ...]).float(),
torch.from_numpy(normals[None, ...]).float()).to(device)
else:
pcl_gt = dataset.get_pointclouds(n_pts).to(device)
trimesh.Trimesh(pcl_gt.points_packed().cpu().numpy(),
vertex_normals=pcl_gt.normals_packed().cpu().numpy(),
process=False).export(val_gt_pts_file,
vertex_normal=True)
# load vis directories
vis_dir = os.path.join(exp_dir, 'vis')
vis_files = sorted(get_filenames(vis_dir, '_mesh.ply'))
iters = [int(os.path.basename(v).split('_')[0]) for v in vis_files]
best_dict = defaultdict(lambda: float('inf'))
vis_eval_csv = os.path.join(vis_dir, "evaluation_n%d.csv" % n_pts)
if not os.path.isfile(vis_eval_csv):
with open(os.path.join(vis_dir, "evaluation_n%d.csv" % n_pts),
"w") as f:
fieldnames = ['mtime', 'it', 'chamfer_p', 'chamfer_n', 'pf_dist']
writer = csv.DictWriter(f,
fieldnames=fieldnames,
restval="-",
extrasaction="ignore")
writer.writeheader()
mtime0 = None
for it, vis_file in zip(iters, vis_files):
eval_dict = OrderedDict()
mtime = os.path.getmtime(vis_file)
if mtime0 is None:
mtime0 = mtime
eval_dict['it'] = it
eval_dict['mtime'] = mtime - mtime0
val_pts_file = os.path.join(
vis_dir,
os.path.basename(vis_file).replace('_mesh',
'_val%d' % n_pts))
if os.path.isfile(val_pts_file):
points, normals = np.split(read_ply(val_pts_file),
2,
axis=1)
points = torch.from_numpy(points).float().to(
device=device).view(1, -1, 3)
normals = torch.from_numpy(normals).float().to(
device=device).view(1, -1, 3)
else:
mesh = trimesh.load(vis_file, process=False)
# points, normals = pcu.sample_mesh_poisson_disk(
# mesh.vertices, mesh.faces,
# mesh.vertex_normals.ravel().reshape(-1, 3), n_pts, use_geodesic_distance=True)
# p_idx = np.random.permutation(points.shape[0])[:n_pts]
# points = points[p_idx, ...]
# normals = normals[p_idx, ...]
# points = torch.from_numpy(points).float().to(
# device=device).view(1, -1, 3)
# normals = torch.from_numpy(normals).float().to(
# device=device).view(1, -1, 3)
meshes = Meshes(
torch.from_numpy(mesh.vertices[None, ...]).float(),
torch.from_numpy(mesh.faces[None,
...]).float()).to(device)
points, normals = sample_points_from_meshes(
meshes, n_pts, return_normals=True)
trimesh.Trimesh(points.cpu().numpy()[0],
vertex_normals=normals.cpu().numpy()[0],
process=False).export(val_pts_file,
vertex_normal=True)
pcl = Pointclouds(points, normals)
chamfer_p, chamfer_n = chamfer_distance(
points,
pcl_gt.points_padded(),
x_normals=normals,
y_normals=pcl_gt.normals_padded(),
)
eval_dict['chamfer_p'] = chamfer_p.item()
eval_dict['chamfer_n'] = chamfer_n.item()
pf_dist = point_mesh_face_distance(meshes_gt, pcl)
eval_dict['pf_dist'] = pf_dist.item()
writer.writerow(eval_dict)
for k, v in eval_dict.items():
if v < best_dict[k]:
best_dict[k] = v
print('best {} so far ({}): {:.4g}'.format(
k, vis_file, v))
# generation dictories
gen_dir = os.path.join(exp_dir, 'generation')
if not os.path.isdir(gen_dir):
return
final_file = os.path.join(gen_dir, 'mesh.ply')
val_pts_file = final_file[:-4] + '_val%d' % n_pts + '.ply'
if not os.path.isfile(final_file):
return
gen_file_csv = os.path.join(gen_dir, "evaluation_n%d.csv" % n_pts)
if not os.path.isfile(gen_file_csv):
with open(os.path.join(gen_dir, "evaluation_n%d.csv" % n_pts),
"w") as f:
fieldnames = ['chamfer_p', 'chamfer_n', 'pf_dist']
writer = csv.DictWriter(f,
fieldnames=fieldnames,
restval="-",
extrasaction="ignore")
writer.writeheader()
eval_dict = OrderedDict()
mesh = trimesh.load(final_file)
# points, normals = pcu.sample_mesh_poisson_disk(
# mesh.vertices, mesh.faces,
# mesh.vertex_normals.ravel().reshape(-1, 3), n_pts, use_geodesic_distance=True)
# p_idx = np.random.permutation(points.shape[0])[:n_pts]
# points = points[p_idx, ...]
# normals = normals[p_idx, ...]
# points = torch.from_numpy(points).float().to(
# device=device).view(1, -1, 3)
# normals = torch.from_numpy(normals).float().to(
# device=device).view(1, -1, 3)
meshes = Meshes(
torch.from_numpy(mesh.vertices[None, ...]).float(),
torch.from_numpy(mesh.faces[None, ...]).float()).to(device)
points, normals = sample_points_from_meshes(meshes,
n_pts,
return_normals=True)
trimesh.Trimesh(points.cpu().numpy()[0],
vertex_normals=normals.cpu().numpy()[0],
process=False).export(val_pts_file,
vertex_normal=True)
pcl = Pointclouds(points, normals)
chamfer_p, chamfer_n = chamfer_distance(
points,
pcl_gt.points_padded(),
x_normals=normals,
y_normals=pcl_gt.normals_padded(),
)
eval_dict['chamfer_p'] = chamfer_p.item()
eval_dict['chamfer_n'] = chamfer_n.item()
pf_dist = point_mesh_face_distance(meshes_gt, pcl)
eval_dict['pf_dist'] = pf_dist.item()
writer.writerow(eval_dict)
for k, v in eval_dict.items():
if v < best_dict[k]:
best_dict[k] = v
print('best {} so far ({}): {:.4g}'.format(
k, final_file, v))
def render_yaw(self,
im,
depth,
v_before=None,
v_after=None,
rotations=None,
maxr=90,
nsample=9,
crop_mesh=None):
b, c, h, w = im.shape
grid_3d = self.depth_to_3d_grid(depth)
if crop_mesh is not None:
top, bottom, left, right = crop_mesh # pixels from border to be cropped
if top > 0:
grid_3d[:, :top, :, 1] = grid_3d[:, top:top + 1, :,
1].repeat(1, top, 1)
grid_3d[:, :top, :, 2] = grid_3d[:, top:top + 1, :,
2].repeat(1, top, 1)
if bottom > 0:
grid_3d[:, -bottom:, :, 1] = grid_3d[:, -bottom - 1:-bottom, :,
1].repeat(1, bottom, 1)
grid_3d[:, -bottom:, :, 2] = grid_3d[:, -bottom - 1:-bottom, :,
2].repeat(1, bottom, 1)
if left > 0:
grid_3d[:, :, :left, 0] = grid_3d[:, :, left:left + 1,
0].repeat(1, 1, left)
grid_3d[:, :, :left, 2] = grid_3d[:, :, left:left + 1,
2].repeat(1, 1, left)
if right > 0:
grid_3d[:, :, -right:, 0] = grid_3d[:, :, -right - 1:-right,
0].repeat(1, 1, right)
grid_3d[:, :, -right:, 2] = grid_3d[:, :, -right - 1:-right,
2].repeat(1, 1, right)
grid_3d = grid_3d.reshape(b, -1, 3)
im_trans = []
# inverse warp
if v_before is not None:
rot_mat, trans_xyz = get_transform_matrices(v_before)
grid_3d = self.translate_pts(grid_3d, -trans_xyz)
grid_3d = self.rotate_pts(grid_3d, rot_mat.transpose(2, 1))
if rotations is None:
rotations = torch.linspace(-math.pi / 180 * maxr,
math.pi / 180 * maxr, nsample)
for i, ri in enumerate(rotations):
ri = torch.FloatTensor([0, ri, 0]).to(im.device).view(1, 3)
rot_mat_i, _ = get_transform_matrices(ri)
grid_3d_i = self.rotate_pts(grid_3d, rot_mat_i.repeat(b, 1, 1))
if v_after is not None:
if len(v_after.shape) == 3:
v_after_i = v_after[i]
else:
v_after_i = v_after
rot_mat, trans_xyz = get_transform_matrices(v_after_i)
grid_3d_i = self.rotate_pts(grid_3d_i, rot_mat)
grid_3d_i = self.translate_pts(grid_3d_i, trans_xyz)
textures = get_textures_from_im(im, tx_size=self.tex_cube_size)
grid_3d_i = Pointclouds(points=grid_3d_i, features=textures)
warped_images = self.renderer(grid_3d_i).clamp(min=-1., max=1.)
# print("扭过来的图像大小", warped_images.shape) torch.Size([8, 64, 64, 3])
im_trans += [warped_images]
return torch.stack(im_trans, 1) # b x t x c x h x w
def _add_pointcloud_trace(
fig: go.Figure,
pointclouds: Pointclouds,
trace_name: str,
subplot_idx: int,
ncols: int,
max_points_per_pointcloud: int,
marker_size: int,
): # pragma: no cover
"""
Adds a trace rendering a Pointclouds object to the passed in figure, with
a given name and in a specific subplot.
Args:
fig: plotly figure to add the trace within.
pointclouds: Pointclouds object to render. It can be batched.
trace_name: name to label the trace with.
subplot_idx: identifies the subplot, with 0 being the top left.
ncols: the number of subplots per row.
max_points_per_pointcloud: the number of points to render, which are randomly sampled.
marker_size: the size of the rendered points
"""
pointclouds = pointclouds.detach().cpu().subsample(
max_points_per_pointcloud)
verts = pointclouds.points_packed()
features = pointclouds.features_packed()
color = None
if features is not None:
if features.shape[1] == 4: # rgba
template = "rgb(%d, %d, %d, %f)"
rgb = (features[:, :3].clamp(0.0, 1.0) * 255).int()
color = [
template % (*rgb_, a_)
for rgb_, a_ in zip(rgb, features[:, 3])
]
if features.shape[1] == 3:
template = "rgb(%d, %d, %d)"
rgb = (features.clamp(0.0, 1.0) * 255).int()
color = [template % (r, g, b) for r, g, b in rgb]
row = subplot_idx // ncols + 1
col = subplot_idx % ncols + 1
fig.add_trace(
go.Scatter3d(
x=verts[:, 0],
y=verts[:, 1],
z=verts[:, 2],
marker={
"color": color,
"size": marker_size
},
mode="markers",
name=trace_name,
),
row=row,
col=col,
)
# Access the current subplot's scene configuration
plot_scene = "scene" + str(subplot_idx + 1)
current_layout = fig["layout"][plot_scene]
# update the bounds of the axes for the current trace
verts_center = verts.mean(0)
max_expand = (verts.max(0)[0] - verts.min(0)[0]).max()
_update_axes_bounds(verts_center, max_expand, current_layout)
def test_simple_sphere(self):
device = torch.device("cuda:0")
# Load reference image
ref_filename = "test_simple_pointcloud_sphere.png"
image_ref_filename = DATA_DIR / ref_filename
# Rescale image_ref to the 0 - 1 range and convert to a binary mask.
image_ref = convert_image_to_binary_mask(image_ref_filename).to(
torch.int32)
sphere_mesh = ico_sphere(1, device)
verts_padded = sphere_mesh.verts_padded()
verts_padded[..., 1] += 0.2
verts_padded[..., 0] += 0.2
pointclouds = Pointclouds(points=verts_padded)
R, T = look_at_view_transform(2.7, 0.0, 0.0)
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)
raster_settings = PointsRasterizationSettings(image_size=256,
radius=5e-2,
points_per_pixel=1)
#################################
# 1. Test init without cameras.
##################################
# Initialize without passing in the cameras
rasterizer = PointsRasterizer()
# Check that omitting the cameras in both initialization
# and the forward pass throws an error:
with self.assertRaisesRegex(ValueError, "Cameras must be specified"):
rasterizer(pointclouds)
##########################################
# 2. Test rasterizing a single pointcloud
##########################################
fragments = rasterizer(pointclouds,
cameras=cameras,
raster_settings=raster_settings)
# Convert idx to a binary mask
image = fragments.idx[0, ..., 0].squeeze().cpu()
image[image >= 0] = 1.0
image[image < 0] = 0.0
if DEBUG:
Image.fromarray((image.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_test_rasterized_sphere_points.png")
self.assertTrue(torch.allclose(image, image_ref[..., 0]))
########################################
# 3. Test with a batch of pointclouds
########################################
batch_size = 10
pointclouds = pointclouds.extend(batch_size)
fragments = rasterizer(pointclouds,
cameras=cameras,
raster_settings=raster_settings)
for i in range(batch_size):
image = fragments.idx[i, ..., 0].squeeze().cpu()
image[image >= 0] = 1.0
image[image < 0] = 0.0
self.assertTrue(torch.allclose(image, image_ref[..., 0]))
def point_mesh_face_distance(meshes: Meshes, pcls: Pointclouds):
"""
Computes the distance between a pointcloud and a mesh within a batch.
Given a pair `(mesh, pcl)` in the batch, we define the distance to be the
sum of two distances, namely `point_face(mesh, pcl) + face_point(mesh, pcl)`
`point_face(mesh, pcl)`: Computes the squared distance of each point p in pcl
to the closest triangular face in mesh and averages across all points in pcl
`face_point(mesh, pcl)`: Computes the squared distance of each triangular face in
mesh to the closest point in pcl and averages across all faces in mesh.
The above distance functions are applied for all `(mesh, pcl)` pairs in the batch
and then averaged across the batch.
Args:
meshes: A Meshes data structure containing N meshes
pcls: A Pointclouds data structure containing N pointclouds
Returns:
loss: The `point_face(mesh, pcl) + face_point(mesh, pcl)` distance
between all `(mesh, pcl)` in a batch averaged across the batch.
"""
if len(meshes) != len(pcls):
raise ValueError("meshes and pointclouds must be equal sized batches")
N = len(meshes)
# packed representation for pointclouds
points = pcls.points_packed() # (P, 3)
points_first_idx = pcls.cloud_to_packed_first_idx()
max_points = pcls.num_points_per_cloud().max().item()
# packed representation for faces
verts_packed = meshes.verts_packed()
faces_packed = meshes.faces_packed()
tris = verts_packed[faces_packed] # (T, 3, 3)
tris_first_idx = meshes.mesh_to_faces_packed_first_idx()
max_tris = meshes.num_faces_per_mesh().max().item()
# point to face distance: shape (P,)
point_to_face = point_face_distance(points, points_first_idx, tris,
tris_first_idx, max_points)
# weight each example by the inverse of number of points in the example
point_to_cloud_idx = pcls.packed_to_cloud_idx() # (sum(P_i),)
num_points_per_cloud = pcls.num_points_per_cloud() # (N,)
weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx)
weights_p = 1.0 / weights_p.float()
point_to_face = point_to_face * weights_p
point_dist = point_to_face.sum() / N
# face to point distance: shape (T,)
face_to_point = face_point_distance(points, points_first_idx, tris,
tris_first_idx, max_tris)
# weight each example by the inverse of number of faces in the example
tri_to_mesh_idx = meshes.faces_packed_to_mesh_idx() # (sum(T_n),)
num_tris_per_mesh = meshes.num_faces_per_mesh() # (N, )
weights_t = num_tris_per_mesh.gather(0, tri_to_mesh_idx)
weights_t = 1.0 / weights_t.float()
face_to_point = face_to_point * weights_t
face_dist = face_to_point.sum() / N
return point_dist + face_dist
