Browse Source
Update
Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
3 changed files with
6 additions and
6 deletions
-
README.md
-
image_embedding_operator_template.py
-
image_embedding_operator_template.yaml
|
|
@ -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:** |
|
|
|
|
|
|
|
|
|
@ -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) |
|
|
|
|
|
@ -8,6 +8,6 @@ init: |
|
|
|
model_name: str |
|
|
|
call: |
|
|
|
input: |
|
|
|
img_tensor: torch.Tensor |
|
|
|
img_path: str |
|
|
|
output: |
|
|
|
feature_vector: numpy.ndarray |