Update

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>

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
 

pytorch/requirements.txt

@@ -1,2 +0,0 @@
-torch>=1.2.0
-torchvision>=0.4.0

test_resnet50_image_embedding.py

@@ -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()