How to use torchgeometry - 10 common examples

def test_jit(self):
        batch_size, channels, height, width = 2, 3, 64, 64
        img = torch.ones(batch_size, channels, height, width)
        gray = color.RgbToGrayscale()
        gray_traced = torch.jit.trace(color.RgbToGrayscale(), img)
        assert_allclose(gray(img), gray_traced(img))

def test_jit(self):
        batch_size, channels, height, width = 2, 3, 64, 64
        img = torch.ones(batch_size, channels, height, width)
        gray = color.RgbToGrayscale()
        gray_traced = torch.jit.trace(color.RgbToGrayscale(), img)
        assert_allclose(gray(img), gray_traced(img))

def test_smoke(self):
        height, width, channels = 4, 5, 3 
        img = np.ones((height, width, channels))

        transforms = nn.Sequential(
            taug.ToTensor(),
            taug.Grayscale(),
        )
        assert transforms(img).shape == (1, height, width)

def test_rgb_to_tensor(self):
        height, width, channels = 4, 5, 3 
        img = np.ones((height, width, channels))
        assert taug.ToTensor()(img).shape == (channels, height, width)

def test_rgb_to_tensor_batch(self):
        batch_size, height, width, channels = 2, 4, 5, 3 
        img = np.ones((batch_size, height, width, channels))
        assert taug.ToTensor()(img).shape == \
            (batch_size, channels, height, width)

def test_mono_to_tensor(self):
        height, width = 4, 5
        img = np.ones((height, width))
        assert taug.ToTensor()(img).shape == (1, height, width)

def test_smoke(self):
        assert str(taug.ToTensor()) == 'ToTensor()'

def test_gray_to_tensor(self):
        height, width, channels = 4, 5, 1
        img = np.ones((height, width, channels))
        assert taug.ToTensor()(img).shape == (1, height, width)

def test_smoke_batch(self):
        batch_size, height, width, channels = 2, 4, 5, 3 
        img = np.ones((batch_size, height, width, channels))

        transforms = nn.Sequential(
            taug.ToTensor(),
            taug.Grayscale(),
        )
        assert transforms(img).shape == (batch_size, 1, height, width)

def test_smoke_batch(self):
        batch_size, height, width, channels = 2, 4, 5, 3 
        img = np.ones((batch_size, height, width, channels))

        transforms = torchvision.transforms.Compose([
            taug.ToTensor(),
            taug.Grayscale(),
        ])
        assert transforms(img).shape == (batch_size, 1, height, width)