diff --git a/README.md b/README.md index 7e9f1e0..b203d8e 100644 --- a/README.md +++ b/README.md @@ -14,13 +14,13 @@ This Operator generates feature vectors from the pytorch pretrained **Resnet50** **Args:** -​ model_name (`str`): +- model_name (`str`): -​ The model name for embedding, for example 'resnet50'. + The model name for embedding, for example 'resnet50'. -​ framework (`str`): +- framework (`str`): -​ The framework of the model, the default is 'pytorch'. + The framework of the model, the default is 'pytorch'. `__call__(self, img_tensor: torch.Tensor)` @@ -32,7 +32,7 @@ This Operator generates feature vectors from the pytorch pretrained **Resnet50** **Returns:** -​ (`Tuple[('cnn', numpy.ndarray)]`) +​ (`Tuple[('feature_vector', numpy.ndarray)]`) ​ The embedding of image. @@ -42,7 +42,7 @@ You can get the required python package by [requirements.txt](./requirements.txt ## How it works -The `towhee/resnet50-image-embedding` Operator implements the function of image embedding, which can add to the pipeline, for example, it's the key Operator named embedding_model within [image_embedding_resnet50](https://hub.towhee.io/towhee/image-embedding-resnet50) pipeline, and it is the red box in the picture below. +The `towhee/resnet50-image-embedding` Operator implements the function of image embedding, which can add to the pipeline, for example, it's the key Operator named embedding_model within [image-embedding-resnet50](https://hub.towhee.io/towhee/image-embedding-resnet50) pipeline, and it is the red box in the picture below. ![img](./readme_res/operator.png) diff --git a/resnet50_image_embedding.py b/resnet50_image_embedding.py index 8e9a7cf..a3401fd 100644 --- a/resnet50_image_embedding.py +++ b/resnet50_image_embedding.py @@ -33,7 +33,7 @@ class Resnet50ImageEmbedding(Operator): from pytorch.model import Model self.model = Model(model_name) - def __call__(self, img_tensor: torch.Tensor) -> NamedTuple('Outputs', [('cnn', numpy.ndarray)]): + def __call__(self, img_tensor: torch.Tensor) -> NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]): embedding = self.model(img_tensor) - Outputs = NamedTuple('Outputs', [('cnn', numpy.ndarray)]) + Outputs = NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]) return Outputs(embedding) diff --git a/resnet50_image_embedding.yaml b/resnet50_image_embedding.yaml index 7d1a2ba..ec63089 100644 --- a/resnet50_image_embedding.yaml +++ b/resnet50_image_embedding.yaml @@ -10,4 +10,4 @@ call: input: img_tensor: torch.Tensor output: - cnn: numpy.ndarray + feature_vector: numpy.ndarray