mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
8219a52ccc
Summary: Fix a few linting warnings Reviewed By: nikhilaravi Differential Revision: D20720810 fbshipit-source-id: c5b6a25fdd7971cc8743b54bbe162464a874071d
229 lines
8.4 KiB
Python
229 lines
8.4 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
|
|
|
|
"""
|
|
Sanity checks for output images from the pointcloud renderer.
|
|
"""
|
|
import unittest
|
|
import warnings
|
|
from os import path
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from common_testing import TestCaseMixin, load_rgb_image
|
|
from PIL import Image
|
|
from pytorch3d.renderer.cameras import (
|
|
FoVOrthographicCameras,
|
|
FoVPerspectiveCameras,
|
|
look_at_view_transform,
|
|
)
|
|
from pytorch3d.renderer.compositing import alpha_composite, norm_weighted_sum
|
|
from pytorch3d.renderer.points import (
|
|
AlphaCompositor,
|
|
NormWeightedCompositor,
|
|
PointsRasterizationSettings,
|
|
PointsRasterizer,
|
|
PointsRenderer,
|
|
)
|
|
from pytorch3d.structures.pointclouds import Pointclouds
|
|
from pytorch3d.utils.ico_sphere import ico_sphere
|
|
|
|
|
|
# If DEBUG=True, save out images generated in the tests for debugging.
|
|
# All saved images have prefix DEBUG_
|
|
DEBUG = False
|
|
DATA_DIR = Path(__file__).resolve().parent / "data"
|
|
|
|
|
|
class TestRenderPoints(TestCaseMixin, unittest.TestCase):
|
|
def test_simple_sphere(self):
|
|
device = torch.device("cuda:0")
|
|
sphere_mesh = ico_sphere(1, device)
|
|
verts_padded = sphere_mesh.verts_padded()
|
|
# Shift vertices to check coordinate frames are correct.
|
|
verts_padded[..., 1] += 0.2
|
|
verts_padded[..., 0] += 0.2
|
|
pointclouds = Pointclouds(
|
|
points=verts_padded, features=torch.ones_like(verts_padded)
|
|
)
|
|
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=256, radius=5e-2, points_per_pixel=1
|
|
)
|
|
rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings)
|
|
compositor = NormWeightedCompositor()
|
|
renderer = PointsRenderer(rasterizer=rasterizer, compositor=compositor)
|
|
|
|
# Load reference image
|
|
filename = "simple_pointcloud_sphere.png"
|
|
image_ref = load_rgb_image("test_%s" % filename, 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:
|
|
filename = "DEBUG_%s" % filename
|
|
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
|
|
DATA_DIR / filename
|
|
)
|
|
self.assertClose(rgb, image_ref)
|
|
|
|
def test_pointcloud_with_features(self):
|
|
device = torch.device("cuda:0")
|
|
file_dir = Path(__file__).resolve().parent.parent / "docs/tutorials/data"
|
|
pointcloud_filename = file_dir / "PittsburghBridge/pointcloud.npz"
|
|
|
|
# Note, this file is too large to check in to the repo.
|
|
# Download the file to run the test locally.
|
|
if not path.exists(pointcloud_filename):
|
|
url = (
|
|
"https://dl.fbaipublicfiles.com/pytorch3d/data/"
|
|
"PittsburghBridge/pointcloud.npz"
|
|
)
|
|
msg = (
|
|
"pointcloud.npz not found, download from %s, save it at the path %s, and rerun"
|
|
% (url, pointcloud_filename)
|
|
)
|
|
warnings.warn(msg)
|
|
return True
|
|
|
|
# Load point cloud
|
|
pointcloud = np.load(pointcloud_filename)
|
|
verts = torch.Tensor(pointcloud["verts"]).to(device)
|
|
rgb_feats = torch.Tensor(pointcloud["rgb"]).to(device)
|
|
|
|
verts.requires_grad = True
|
|
rgb_feats.requires_grad = True
|
|
point_cloud = Pointclouds(points=[verts], features=[rgb_feats])
|
|
|
|
R, T = look_at_view_transform(20, 10, 0)
|
|
cameras = FoVOrthographicCameras(device=device, R=R, T=T, znear=0.01)
|
|
|
|
raster_settings = PointsRasterizationSettings(
|
|
# Set image_size so it is not a multiple of 16 (min bin_size)
|
|
# in order to confirm that there are no errors in coarse rasterization.
|
|
image_size=500,
|
|
radius=0.003,
|
|
points_per_pixel=10,
|
|
)
|
|
|
|
renderer = PointsRenderer(
|
|
rasterizer=PointsRasterizer(
|
|
cameras=cameras, raster_settings=raster_settings
|
|
),
|
|
compositor=AlphaCompositor(),
|
|
)
|
|
|
|
images = renderer(point_cloud)
|
|
|
|
# Load reference image
|
|
filename = "bridge_pointcloud.png"
|
|
image_ref = load_rgb_image("test_%s" % filename, 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(point_cloud)
|
|
rgb = images[0, ..., :3].squeeze().cpu()
|
|
if DEBUG:
|
|
filename = "DEBUG_%s" % filename
|
|
Image.fromarray((rgb.detach().numpy() * 255).astype(np.uint8)).save(
|
|
DATA_DIR / filename
|
|
)
|
|
self.assertClose(rgb, image_ref, atol=0.015)
|
|
|
|
# Check grad exists.
|
|
grad_images = torch.randn_like(images)
|
|
images.backward(grad_images)
|
|
self.assertIsNotNone(verts.grad)
|
|
self.assertIsNotNone(rgb_feats.grad)
|
|
|
|
def test_simple_sphere_batched(self):
|
|
device = torch.device("cuda:0")
|
|
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, features=torch.ones_like(verts_padded)
|
|
)
|
|
batch_size = 20
|
|
pointclouds = pointclouds.extend(batch_size)
|
|
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=256, radius=5e-2, points_per_pixel=1
|
|
)
|
|
rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings)
|
|
compositor = NormWeightedCompositor()
|
|
renderer = PointsRenderer(rasterizer=rasterizer, compositor=compositor)
|
|
|
|
# Load reference image
|
|
filename = "simple_pointcloud_sphere.png"
|
|
image_ref = load_rgb_image("test_%s" % filename, DATA_DIR)
|
|
|
|
images = renderer(pointclouds)
|
|
for i in range(batch_size):
|
|
rgb = images[i, ..., :3].squeeze().cpu()
|
|
if i == 0 and DEBUG:
|
|
filename = "DEBUG_%s" % filename
|
|
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
|
|
DATA_DIR / filename
|
|
)
|
|
self.assertClose(rgb, image_ref)
|
|
|
|
def test_compositor_background_color(self):
|
|
|
|
N, H, W, K, C, P = 1, 15, 15, 20, 4, 225
|
|
ptclds = torch.randn((C, P))
|
|
alphas = torch.rand((N, K, H, W))
|
|
pix_idxs = torch.randint(-1, 20, (N, K, H, W)) # 20 < P, large amount of -1
|
|
background_color = [0.5, 0, 1]
|
|
|
|
compositor_funcs = [
|
|
(NormWeightedCompositor, norm_weighted_sum),
|
|
(AlphaCompositor, alpha_composite),
|
|
]
|
|
|
|
for (compositor_class, composite_func) in compositor_funcs:
|
|
|
|
compositor = compositor_class(background_color)
|
|
|
|
# run the forward method to generate masked images
|
|
masked_images = compositor.forward(pix_idxs, alphas, ptclds)
|
|
|
|
# generate unmasked images for testing purposes
|
|
images = composite_func(pix_idxs, alphas, ptclds)
|
|
|
|
is_foreground = pix_idxs[:, 0] >= 0
|
|
|
|
# make sure foreground values are unchanged
|
|
self.assertClose(
|
|
torch.masked_select(masked_images, is_foreground[:, None]),
|
|
torch.masked_select(images, is_foreground[:, None]),
|
|
)
|
|
|
|
is_background = ~is_foreground[..., None].expand(-1, -1, -1, 4)
|
|
|
|
# permute masked_images to correctly get rgb values
|
|
masked_images = masked_images.permute(0, 2, 3, 1)
|
|
for i in range(3):
|
|
channel_color = background_color[i]
|
|
|
|
# check if background colors are properly changed
|
|
self.assertTrue(
|
|
masked_images[is_background]
|
|
.view(-1, 4)[..., i]
|
|
.eq(channel_color)
|
|
.all()
|
|
)
|
|
|
|
# check background color alpha values
|
|
self.assertTrue(
|
|
masked_images[is_background].view(-1, 4)[..., 3].eq(1).all()
|
|
)
|