pytorch3d/tests/pulsar/test_depth.py
Thomas Polasek 055ab3a2e3 Convert directory fbcode/vision to use the Ruff Formatter
Summary:
Converts the directory specified to use the Ruff formatter in pyfmt

ruff_dog

If this diff causes merge conflicts when rebasing, please run
`hg status -n -0 --change . -I '**/*.{py,pyi}' | xargs -0 arc pyfmt`
on your diff, and amend any changes before rebasing onto latest.
That should help reduce or eliminate any merge conflicts.

allow-large-files

Reviewed By: bottler

Differential Revision: D66472063

fbshipit-source-id: 35841cb397e4f8e066e2159550d2f56b403b1bef
2024-11-26 02:38:20 -08:00

96 lines
3.4 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.
"""Test the sorting of the closest spheres."""
import logging
import os
import sys
import unittest
from os import path
import imageio
import numpy as np
import torch
from ..common_testing import TestCaseMixin
# Making sure you can run this, even if pulsar hasn't been installed yet.
sys.path.insert(0, path.join(path.dirname(__file__), "..", ".."))
devices = [torch.device("cuda"), torch.device("cpu")]
IN_REF_FP = path.join(path.dirname(__file__), "reference", "nr0000-in.pth")
OUT_REF_FP = path.join(path.dirname(__file__), "reference", "nr0000-out.pth")
class TestDepth(TestCaseMixin, unittest.TestCase):
"""Test different numbers of channels."""
def test_basic(self):
from pytorch3d.renderer.points.pulsar import Renderer
for device in devices:
gamma = 1e-5
max_depth = 15.0
min_depth = 5.0
renderer = Renderer(
256,
256,
10000,
orthogonal_projection=True,
right_handed_system=False,
n_channels=1,
).to(device)
data = torch.load(IN_REF_FP, map_location="cpu")
# For creating the reference files.
# Use in case of updates.
# data["pos"] = torch.rand_like(data["pos"])
# data["pos"][:, 0] = data["pos"][:, 0] * 2. - 1.
# data["pos"][:, 1] = data["pos"][:, 1] * 2. - 1.
# data["pos"][:, 2] = data["pos"][:, 2] + 9.5
result, result_info = renderer.forward(
data["pos"].to(device),
data["col"].to(device),
data["rad"].to(device),
data["cam_params"].to(device),
gamma,
min_depth=min_depth,
max_depth=max_depth,
return_forward_info=True,
bg_col=torch.zeros(1, device=device, dtype=torch.float32),
percent_allowed_difference=0.01,
)
depth_map = Renderer.depth_map_from_result_info_nograd(result_info)
depth_vis = (depth_map - depth_map[depth_map > 0].min()) * 200 / (
depth_map.max() - depth_map[depth_map > 0.0].min()
) + 50
if not os.environ.get("FB_TEST", False):
imageio.imwrite(
path.join(
path.dirname(__file__),
"test_out",
"test_depth_test_basic_depth.png",
),
depth_vis.cpu().numpy().astype(np.uint8),
)
# For creating the reference files.
# Use in case of updates.
# torch.save(
# data, path.join(path.dirname(__file__), "reference", "nr0000-in.pth")
# )
# torch.save(
# {"sphere_ids": sphere_ids, "depth_map": depth_map},
# path.join(path.dirname(__file__), "reference", "nr0000-out.pth"),
# )
# sys.exit(0)
reference = torch.load(OUT_REF_FP, map_location="cpu")
self.assertClose(reference["depth_map"].to(device), depth_map)
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
unittest.main()