Browse Source
Update the test script
Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
2 changed files with
4 additions and
14 deletions
-
config.py
-
test_image_embedding_pipeline_yaml.py
|
|
@ -1,3 +0,0 @@ |
|
|
TEST_IMG = './test_data/test.jpg' |
|
|
|
|
|
DIMENSION = 1000 |
|
|
|
|
|
REPO_NAME = 'towhee/image-embedding-resnet' |
|
|
|
|
|
@ -3,14 +3,7 @@ from towhee import pipeline |
|
|
from config import DIMENSION, REPO_NAME, TEST_IMG |
|
|
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 |
|
|
|
