52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
|
|
import torch
|
|
import unittest
|
|
|
|
from cotracker.models.core.model_utils import bilinear_sampler
|
|
|
|
|
|
class TestBilinearSampler(unittest.TestCase):
|
|
# Sample from an image (4d)
|
|
def _test4d(self, align_corners):
|
|
H, W = 4, 5
|
|
# Construct a grid to obtain indentity sampling
|
|
input = torch.randn(H * W).view(1, 1, H, W).float()
|
|
coords = torch.meshgrid(torch.arange(H), torch.arange(W))
|
|
coords = torch.stack(coords[::-1], dim=-1).float()[None]
|
|
if not align_corners:
|
|
coords = coords + 0.5
|
|
sampled_input = bilinear_sampler(input, coords, align_corners=align_corners)
|
|
torch.testing.assert_close(input, sampled_input)
|
|
|
|
# Sample from a video (5d)
|
|
def _test5d(self, align_corners):
|
|
T, H, W = 3, 4, 5
|
|
# Construct a grid to obtain indentity sampling
|
|
input = torch.randn(H * W).view(1, 1, H, W).float()
|
|
input = torch.stack([input, input + 1, input + 2], dim=2)
|
|
coords = torch.meshgrid(torch.arange(T), torch.arange(W), torch.arange(H))
|
|
coords = torch.stack(coords, dim=-1).float().permute(0, 2, 1, 3)[None]
|
|
|
|
if not align_corners:
|
|
coords = coords + 0.5
|
|
sampled_input = bilinear_sampler(input, coords, align_corners=align_corners)
|
|
torch.testing.assert_close(input, sampled_input)
|
|
|
|
def test4d(self):
|
|
self._test4d(align_corners=True)
|
|
self._test4d(align_corners=False)
|
|
|
|
def test5d(self):
|
|
self._test5d(align_corners=True)
|
|
self._test5d(align_corners=False)
|
|
|
|
|
|
# run the test
|
|
unittest.main()
|