docker-python/tests/test_kornia.py at master · HpcDataLab/docker-python · GitHub

19 lines (14 loc) · 508 Bytes

import unittest

import cv2

import torch

import kornia

class TestKornia(unittest.TestCase):

def test_imread_opencv(self):

img = cv2.imread('/input/tests/data/dot.png')

img_t = kornia.image_to_tensor(img)

self.assertEqual(img.shape, (1, 1, 3))

self.assertEqual(img_t.shape, (3, 1, 1))

def test_grayscale_torch(self):

img_rgb = torch.rand(2, 3, 4, 5)

img_gray = kornia.rgb_to_grayscale(img_rgb)

self.assertEqual(img_gray.shape, (2, 1, 4, 5))

Expand file tree