Add a new python unit test to test creating mediapipe Image from cvmat.

PiperOrigin-RevId: 495655719
This commit is contained in:
Jiuqiang Tang 2022-12-15 11:49:06 -08:00 committed by Copybara-Service
parent 299aa03302
commit fd50b6aa2f

View File

@ -28,6 +28,8 @@ import PIL.Image
from mediapipe.python._framework_bindings import image
from mediapipe.python._framework_bindings import image_frame
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
Image = image.Image
ImageFormat = image_frame.ImageFormat
@ -187,5 +189,26 @@ class ImageTest(absltest.TestCase):
gc.collect()
self.assertEqual(sys.getrefcount(rgb_image), initial_ref_count)
def test_image_create_from_cvmat(self):
image_path = os.path.join(os.path.dirname(__file__),
'solutions/testdata/hands.jpg')
mat = cv2.imread(image_path).astype(np.uint8)
mat = cv2.cvtColor(mat, cv2.COLOR_BGR2RGB)
rgb_image = Image(image_format=ImageFormat.SRGB, data=mat)
self.assertEqual(rgb_image.width, 720)
self.assertEqual(rgb_image.height, 382)
self.assertEqual(rgb_image.channels, 3)
self.assertEqual(rgb_image.image_format, ImageFormat.SRGB)
self.assertTrue(np.array_equal(mat, rgb_image.numpy_view()))
def test_image_create_from_file(self):
image_path = os.path.join(os.path.dirname(__file__),
'solutions/testdata/hands.jpg')
loaded_image = Image.create_from_file(image_path)
self.assertEqual(loaded_image.width, 720)
self.assertEqual(loaded_image.height, 382)
self.assertEqual(loaded_image.channels, 3)
self.assertEqual(loaded_image.image_format, ImageFormat.SRGB)
if __name__ == '__main__':
absltest.main()