diff --git a/README.md b/README.md index 6f28e14..7e9f1e0 100644 --- a/README.md +++ b/README.md @@ -10,98 +10,41 @@ This Operator generates feature vectors from the pytorch pretrained **Resnet50** ## Interface -`Class Resnet50ImageEmbedding(Operator)` [\[source\]](./resnet50_image_embedding.py) +`__init__(self, model_name: str, framework: str = 'pytorch')` -`__init__(self, model_name: str)` +**Args:** -**params:** +​ model_name (`str`): -- model_name(str): the model name for embedding, like 'resnet50'. +​ The model name for embedding, for example 'resnet50'. + +​ framework (`str`): + +​ The framework of the model, the default is 'pytorch'. `__call__(self, img_tensor: torch.Tensor)` -**params:** +**Args:** -- img_tensor(torch.Tensor): the normalized image tensor. +​ img_tensor (`torch.Tensor`): -**return:** +​ The image tensor. -- cnn(numpy.ndarray): the embedding of image. +**Returns:** -## How to use +​ (`Tuple[('cnn', numpy.ndarray)]`) -### Requirements +​ The embedding of image. -You can get the required python package by [requirements.txt](./requirements.txt) and [pytorch/requirements.txt](./pytorch/requirements.txt). Towhee will automatically install these packages when you first load the Operator Repo, so you don't need to install them manually, here is just a list. +## Requirements -- towhee -- torch -- torchvision -- numpy +You can get the required python package by [requirements.txt](./requirements.txt). -### How it works +## 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. -![img](./pic/operator.png) - -When using this Operator to write Pipeline's Yaml file, you need to declare the following content according to the interface of Resnet50ImageEmbedding class: - -```yaml -operators: - - - name: 'embedding_model' - function: 'towhee/resnet50-image-embedding' - tag: 'main' - init_args: - model_name: 'resnet50' - inputs: - - - df: 'image_preproc' - name: 'img_tensor' - col: 0 - outputs: - - - df: 'embedding' - iter_info: - type: map -dataframes: - - - name: 'image_preproc' - columns: - - - name: 'img_transformed' - vtype: 'torch.Tensor' - - - name: 'embedding' - columns: - - - name: 'cnn' - vtype: 'numpy.ndarray' -``` - -We can see that in yaml, the **operator** part declares the `init_args` of the class and the `input` and `output` dataframe, and the **dataframe** declares the parameter `name` and `vtype`. - -### File Structure - -Here is the main file structure of the `resnet50-image-embedding` Operator. If you want to learn more about the source code or modify it yourself, you can learn from it. - -```bash -├── .gitattributes -├── .gitignore -├── README.md -├── __init__.py -├── requirements.txt #General python dependency package -├── resnet50_image_embedding.py #The python file for Towhee, it defines the interface of the system and usually does not need to be modified. -├── resnet50_image_embedding.yaml #The YAML file contains Operator information, such as model frame, input, and output. -├── pytorch #The directory of the pytorh -│   ├── __init__.py -│   ├── model #The directory of the pytorch model, which can store data such as weights. -│   ├── requirements.txt #The python dependency package for the pytorch model. -│   └── model.py #The code of the pytorch model, including the initialization model and prediction. -├── test_data/ #The directory of test data, including test.jpg -└── test_resnet50_image_embedding.py #The unittest file of this Operator. -``` +![img](./readme_res/operator.png) ## Reference diff --git a/pytorch/requirements.txt b/pytorch/requirements.txt deleted file mode 100644 index ac3ffc1..0000000 --- a/pytorch/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -torch>=1.2.0 -torchvision>=0.4.0 diff --git a/pic/operator.png b/readme_res/operator.png similarity index 100% rename from pic/operator.png rename to readme_res/operator.png diff --git a/test_data/test.jpg b/test_data/test.jpg deleted file mode 100755 index 8fdc2b3..0000000 Binary files a/test_data/test.jpg and /dev/null differ diff --git a/test_resnet50_image_embedding.py b/test_resnet50_image_embedding.py deleted file mode 100644 index ea6bc08..0000000 --- a/test_resnet50_image_embedding.py +++ /dev/null @@ -1,32 +0,0 @@ -import os -import unittest -from PIL import Image -from torchvision import transforms -from resnet50_image_embedding import Resnet50ImageEmbedding - - -class TestResnet50ImageEmbedding(unittest.TestCase): - """ - Simple operator test - """ - def test_image_embedding(self): - test_img = './test_data/test.jpg' - img = Image.open(test_img) - tfms = transforms.Compose( - [ - transforms.Resize(256), - transforms.CenterCrop(224), - transforms.ToTensor(), - transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), - ] - ) - img_tensor = tfms(img).unsqueeze(0) - model_name = 'resnet50' - dimension = 1000 - op = Resnet50ImageEmbedding(model_name) - print("The output shape of operator:", op(img_tensor)[0].shape) - self.assertEqual((1, dimension), op(img_tensor)[0].shape) - - -if __name__ == '__main__': - unittest.main() \ No newline at end of file