def test_ensure_4d_3_spatial_dims(self): tensor = torch.rand(4, 5, 6) assert io.ensure_4d(tensor, num_spatial_dims=3).shape == (1, 4, 5, 6)
def test_ensure_4d_nd_not_supported(self): tensor = torch.rand(1, 2, 3, 4, 5) with self.assertRaises(ValueError): io.ensure_4d(tensor)
def test_ensure_4d_2d_3dims_rgb_last(self): tensor = torch.rand(4, 5, 3) assert io.ensure_4d(tensor).shape == (3, 4, 5, 1)
def test_ensure_4d_3d(self): tensor = torch.rand(4, 5, 6) assert io.ensure_4d(tensor).shape == (1, 4, 5, 6)
def test_ensure_4d_5d_t_gt_1(self): tensor = torch.rand(3, 4, 5, 2, 2) with self.assertRaises(ValueError): io.ensure_4d(tensor)
def test_ensure_4d_5d(self): tensor = torch.rand(3, 4, 5, 1, 2) assert io.ensure_4d(tensor).shape == (2, 3, 4, 5)
def test_ensure_4d_2_spatial_dims(self): tensor = torch.rand(4, 5, 6) assert ensure_4d(tensor, num_spatial_dims=2).shape == (4, 5, 6, 1)
def test_ensure_4d_2d_3dims_rgb_first(self): tensor = torch.rand(3, 4, 5) assert ensure_4d(tensor).shape == (3, 4, 5, 1)
def test_ensure_4d_2d(self): tensor = torch.rand(4, 5) assert ensure_4d(tensor).shape == (1, 4, 5, 1)