mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-07-31 10:52:50 +08:00
83bacda8fb
Summary: Fix recent flake complaints Reviewed By: MichaelRamamonjisoa Differential Revision: D51811912 fbshipit-source-id: 65183f5bc7058da910e4d5a63b2250ce8637f1cc
66 lines
2.5 KiB
Python
66 lines
2.5 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.implicitron.models.utils import preprocess_input, weighted_sum_losses
|
|
|
|
|
|
class TestUtils(unittest.TestCase):
|
|
def test_prepare_inputs_wrong_num_dim(self):
|
|
img = torch.randn(3, 3, 3)
|
|
text = (
|
|
"Model received unbatched inputs. "
|
|
+ "Perhaps they came from a FrameData which had not been collated."
|
|
)
|
|
with self.assertRaisesRegex(ValueError, text):
|
|
img, fg_prob, depth_map = preprocess_input(
|
|
img, None, None, True, True, 0.5, (0.0, 0.0, 0.0)
|
|
)
|
|
|
|
def test_prepare_inputs_mask_image_true(self):
|
|
batch, channels, height, width = 2, 3, 10, 10
|
|
img = torch.ones(batch, channels, height, width)
|
|
# Create a mask on the lower triangular matrix
|
|
fg_prob = torch.tril(torch.ones(batch, 1, height, width)) * 0.3
|
|
|
|
out_img, out_fg_prob, out_depth_map = preprocess_input(
|
|
img, fg_prob, None, True, False, 0.3, (0.0, 0.0, 0.0)
|
|
)
|
|
|
|
self.assertTrue(torch.equal(out_img, torch.tril(img)))
|
|
self.assertTrue(torch.equal(out_fg_prob, fg_prob >= 0.3))
|
|
self.assertIsNone(out_depth_map)
|
|
|
|
def test_prepare_inputs_mask_depth_true(self):
|
|
batch, channels, height, width = 2, 3, 10, 10
|
|
img = torch.ones(batch, channels, height, width)
|
|
depth_map = torch.randn(batch, channels, height, width)
|
|
# Create a mask on the lower triangular matrix
|
|
fg_prob = torch.tril(torch.ones(batch, 1, height, width)) * 0.3
|
|
|
|
out_img, out_fg_prob, out_depth_map = preprocess_input(
|
|
img, fg_prob, depth_map, False, True, 0.3, (0.0, 0.0, 0.0)
|
|
)
|
|
|
|
self.assertTrue(torch.equal(out_img, img))
|
|
self.assertTrue(torch.equal(out_fg_prob, fg_prob >= 0.3))
|
|
self.assertTrue(torch.equal(out_depth_map, torch.tril(depth_map)))
|
|
|
|
def test_weighted_sum_losses(self):
|
|
preds = {"a": torch.tensor(2), "b": torch.tensor(2)}
|
|
weights = {"a": 2.0, "b": 0.0}
|
|
loss = weighted_sum_losses(preds, weights)
|
|
self.assertEqual(loss, 4.0)
|
|
|
|
def test_weighted_sum_losses_raise_warning(self):
|
|
preds = {"a": torch.tensor(2), "b": torch.tensor(2)}
|
|
weights = {"c": 2.0, "d": 2.0}
|
|
self.assertIsNone(weighted_sum_losses(preds, weights))
|