From baf284c0ee857dd51d709d89a4bc361b9b1bb8cc Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Fri, 24 Dec 2021 16:50:11 +0800 Subject: [PATCH] Update Signed-off-by: shiyu22 --- README.md | 6 +++--- image_embedding_operator_template.py | 4 ++-- image_embedding_operator_template.yaml | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index c44b17c..afe1cf5 100644 --- a/README.md +++ b/README.md @@ -32,9 +32,9 @@ __call__(self, img_tensor: torch.Tensor) **Args:** -- img_tensor: - - the input image tensor - - supported types: `torch.Tensor` +- img_path: + - path to the input image + - supported types: `str` **Returns:** diff --git a/image_embedding_operator_template.py b/image_embedding_operator_template.py index d10755d..bb13bb7 100644 --- a/image_embedding_operator_template.py +++ b/image_embedding_operator_template.py @@ -36,7 +36,7 @@ class ImageEmbeddingOperatorTemplate(Operator): # from pytorch import Model self.model = Model(model_name) - def __call__(self, img_tensor: torch.Tensor) -> NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]): + def __call__(self, img_path: str) -> NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]): """ Call it when use this class. @@ -48,6 +48,6 @@ class ImageEmbeddingOperatorTemplate(Operator): The image embedding. """ # call model with the img_tensor parameter - result = self.model(img_tensor) + result = self.model(img_path) Outputs = NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]) return Outputs(result) diff --git a/image_embedding_operator_template.yaml b/image_embedding_operator_template.yaml index 854f7fd..3232aa8 100644 --- a/image_embedding_operator_template.yaml +++ b/image_embedding_operator_template.yaml @@ -8,6 +8,6 @@ init: model_name: str call: input: - img_tensor: torch.Tensor + img_path: str output: feature_vector: numpy.ndarray \ No newline at end of file