From 3a85c5d920f7c0e7ab01f5f2c901a3ed98855804 Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Wed, 22 Dec 2021 18:30:32 +0800 Subject: [PATCH] Update the test script Signed-off-by: shiyu22 --- config.py | 3 --- test_image_embedding_pipeline_yaml.py | 15 ++++----------- 2 files changed, 4 insertions(+), 14 deletions(-) delete mode 100644 config.py diff --git a/config.py b/config.py deleted file mode 100644 index 6e65711..0000000 --- a/config.py +++ /dev/null @@ -1,3 +0,0 @@ -TEST_IMG = './test_data/test.jpg' -DIMENSION = 1000 -REPO_NAME = 'towhee/image-embedding-resnet' \ No newline at end of file diff --git a/test_image_embedding_pipeline_yaml.py b/test_image_embedding_pipeline_yaml.py index cac3ed3..f0d3c2e 100644 --- a/test_image_embedding_pipeline_yaml.py +++ b/test_image_embedding_pipeline_yaml.py @@ -3,14 +3,7 @@ from towhee import pipeline from config import DIMENSION, REPO_NAME, TEST_IMG -class TestImageEmbeddingPipelineClass(unittest.TestCase): - 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(TEST_IMG) - assert (1, self.dimension)==op(img_tensor)[0].shape - - -if __name__ == '__main__': - unittest.main() +def test_image_embedding_resnet50(repo_name='towhee/image-embedding-pipeline-template', img_path='./test_data/test.jpg', dimension=2048): + embedding_pipeline = pipeline(repo_name) + embedding = embedding_pipeline(img_path) + assert (1, dimension)==embedding[0].shape