import unittest
from towhee import pipeline
from PIL import Image
from config import DIMENSION, REPO_NAME, TEST_IMG


class TestImageEmbeddingPipelineClass(unittest.TestCase):
    test_img = Image.open(TEST_IMG)

    def test_image_embedding_resnet50(self):
        self.dimension = DIMENSION  #the dimension of image embedding
        self.repo_name = REPO_NAME
        embedding_pipeline = pipeline(self.repo_name)
        embedding = embedding_pipeline(self.test_img)
        assert (1, self.dimension)==op(img_tensor)[0].shape


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