From 01494e429f420d33378840ddef38dcdb3defa892 Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Wed, 22 Dec 2021 18:27:23 +0800 Subject: [PATCH] Update the test script Signed-off-by: shiyu22 --- config.py | 3 --- test_image_embedding_operator_template.py | 24 ++++++++--------------- 2 files changed, 8 insertions(+), 19 deletions(-) delete mode 100644 config.py diff --git a/config.py b/config.py deleted file mode 100644 index 2cd63fa..0000000 --- a/config.py +++ /dev/null @@ -1,3 +0,0 @@ -TEST_IMG = './test_data/test.jpg' -DIMENSION = 1000 -MODEL_NAME = 'resnet50' \ No newline at end of file diff --git a/test_image_embedding_operator_template.py b/test_image_embedding_operator_template.py index 41f7303..d5d4970 100644 --- a/test_image_embedding_operator_template.py +++ b/test_image_embedding_operator_template.py @@ -2,14 +2,10 @@ import unittest from PIL import Image from torchvision import transforms from image_embedding_operator_template import ImageEmbeddingOperatorTemplate -from config import DIMENSION, MODEL_NAME, TEST_IMG -class TestImageEmbeddingOperatorTemplate(unittest.TestCase): - """ - Simple operator test - """ - img = Image.open(TEST_IMG) +def get_transformered_img(img_path): + img = Image.open(img_path) tfms = transforms.Compose( [ transforms.Resize(256), @@ -19,14 +15,10 @@ class TestImageEmbeddingOperatorTemplate(unittest.TestCase): ] ) img_tensor = tfms(img).unsqueeze(0) + return img_tensor - def test_image_embedding(self): - self.model_name = MODEL_NAME - self.dimension = DIMENSION - op = ImageEmbeddingOperatorTemplate(self.model_name) - print("The output shape of operator:", op(self.img_tensor)[0].shape) - assert (1, self.dimension)==op(self.img_tensor)[0].shape - - -if __name__ == '__main__': - unittest.main() +def test_image_embedding_operator(model_name='resnet50', img_path='./test_data/test.jpg', dimension=2048): + op = ImageEmbeddingOperatorTemplate(model_name) + img_tensor = get_transformered_img(img_path) + embedding = op(img_tensor) + assert (1, dimension)==embedding[0].shape