mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-07-31 10:52:50 +08:00
36edf2b302
Summary: Needed to properly change devices during OpenGL rasterization. Reviewed By: jcjohnson Differential Revision: D37698568 fbshipit-source-id: 38968149d577322e662d3b5d04880204b0a7be29
126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import unittest
|
|
|
|
import torch
|
|
from pytorch3d.renderer.cameras import look_at_view_transform, PerspectiveCameras
|
|
from pytorch3d.renderer.mesh.rasterizer import Fragments
|
|
from pytorch3d.renderer.mesh.shader import (
|
|
HardDepthShader,
|
|
HardFlatShader,
|
|
HardGouraudShader,
|
|
HardPhongShader,
|
|
SoftDepthShader,
|
|
SoftPhongShader,
|
|
SplatterPhongShader,
|
|
)
|
|
from pytorch3d.structures.meshes import Meshes
|
|
|
|
from .common_testing import TestCaseMixin
|
|
|
|
|
|
class TestShader(TestCaseMixin, unittest.TestCase):
|
|
def setUp(self):
|
|
self.shader_classes = [
|
|
HardDepthShader,
|
|
HardFlatShader,
|
|
HardGouraudShader,
|
|
HardPhongShader,
|
|
SoftDepthShader,
|
|
SoftPhongShader,
|
|
SplatterPhongShader,
|
|
]
|
|
|
|
def test_to(self):
|
|
cpu_device = torch.device("cpu")
|
|
cuda_device = torch.device("cuda:0")
|
|
|
|
R, T = look_at_view_transform()
|
|
|
|
for shader_class in self.shader_classes:
|
|
for cameras_class in (None, PerspectiveCameras):
|
|
if cameras_class is None:
|
|
cameras = None
|
|
else:
|
|
cameras = PerspectiveCameras(device=cpu_device, R=R, T=T)
|
|
|
|
cpu_shader = shader_class(device=cpu_device, cameras=cameras)
|
|
if cameras is None:
|
|
self.assertIsNone(cpu_shader.cameras)
|
|
else:
|
|
self.assertEqual(cpu_device, cpu_shader.cameras.device)
|
|
self.assertEqual(cpu_device, cpu_shader.materials.device)
|
|
self.assertEqual(cpu_device, cpu_shader.lights.device)
|
|
|
|
cuda_shader = cpu_shader.to(cuda_device)
|
|
self.assertIs(cpu_shader, cuda_shader)
|
|
if cameras is None:
|
|
self.assertIsNone(cuda_shader.cameras)
|
|
with self.assertRaisesRegex(ValueError, "Cameras must be"):
|
|
cuda_shader._get_cameras()
|
|
else:
|
|
self.assertEqual(cuda_device, cuda_shader.cameras.device)
|
|
self.assertIsInstance(cuda_shader._get_cameras(), cameras_class)
|
|
self.assertEqual(cuda_device, cuda_shader.materials.device)
|
|
self.assertEqual(cuda_device, cuda_shader.lights.device)
|
|
|
|
def test_cameras_check(self):
|
|
verts = torch.tensor(
|
|
[[-1, -1, 0], [1, -1, 1], [1, 1, 0], [-1, 1, 1]], dtype=torch.float32
|
|
)
|
|
faces = torch.tensor([[0, 1, 2], [2, 3, 0]], dtype=torch.int64)
|
|
meshes = Meshes(verts=[verts], faces=[faces])
|
|
|
|
pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
|
|
barycentric_coords = torch.tensor(
|
|
[[0.1, 0.2, 0.7], [0.3, 0.5, 0.2]], dtype=torch.float32
|
|
).view(1, 1, 1, 2, -1)
|
|
fragments = Fragments(
|
|
pix_to_face=pix_to_face,
|
|
bary_coords=barycentric_coords,
|
|
zbuf=torch.ones_like(pix_to_face),
|
|
dists=torch.ones_like(pix_to_face),
|
|
)
|
|
|
|
for shader_class in self.shader_classes:
|
|
shader = shader_class()
|
|
|
|
with self.assertRaises(ValueError):
|
|
shader(fragments, meshes)
|
|
|
|
def test_depth_shader(self):
|
|
shader_classes = [
|
|
HardDepthShader,
|
|
SoftDepthShader,
|
|
]
|
|
|
|
verts = torch.tensor(
|
|
[[-1, -1, 0], [1, -1, 1], [1, 1, 0], [-1, 1, 1]], dtype=torch.float32
|
|
)
|
|
faces = torch.tensor([[0, 1, 2], [2, 3, 0]], dtype=torch.int64)
|
|
meshes = Meshes(verts=[verts], faces=[faces])
|
|
|
|
pix_to_face = torch.tensor([0, 1], dtype=torch.int64).view(1, 1, 1, 2)
|
|
barycentric_coords = torch.tensor(
|
|
[[0.1, 0.2, 0.7], [0.3, 0.5, 0.2]], dtype=torch.float32
|
|
).view(1, 1, 1, 2, -1)
|
|
for faces_per_pixel in [1, 2]:
|
|
fragments = Fragments(
|
|
pix_to_face=pix_to_face[:, :, :, :faces_per_pixel],
|
|
bary_coords=barycentric_coords[:, :, :, :faces_per_pixel],
|
|
zbuf=torch.ones_like(pix_to_face),
|
|
dists=torch.ones_like(pix_to_face),
|
|
)
|
|
R, T = look_at_view_transform()
|
|
cameras = PerspectiveCameras(R=R, T=T)
|
|
|
|
for shader_class in shader_classes:
|
|
shader = shader_class()
|
|
|
|
out = shader(fragments, meshes, cameras=cameras)
|
|
self.assertEqual(out.shape, (1, 1, 1, 1))
|