import os
import unittest
from towhee import pipeline
from PIL import Image


class TestImageEmbeddingResnet50(unittest.TestCase):
    def test_image_embedding_resnet50(self):
        img_src = './test_data/test.jpg'
        test_img = Image.open(img_src)
        dimension = 1000
        embedding_pipeline = pipeline('towhee/image-embedding-resnet50')
        embedding = embedding_pipeline(test_img)
        self.assertEqual((1, dimension), embedding[0][0].shape)


if __name__ == '__main__':
    unittest.main()