diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..ab2832a --- /dev/null +++ b/.gitattributes @@ -0,0 +1,35 @@ +# .gitattributes + +# Source files +# ============ +*.pxd text diff=python +*.py text diff=python +*.py3 text diff=python +*.pyw text diff=python +*.pyx text diff=python +*.pyz text diff=python +*.pyi text diff=python + +# Binary files +# ============ +*.db binary +*.p binary +*.pkl binary +*.pickle binary +*.pyc binary export-ignore +*.pyo binary export-ignore +*.pyd binary + +# Jupyter notebook +*.ipynb text + +# Model files +*.bin.* filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tar.gz filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..32030ff --- /dev/null +++ b/.gitignore @@ -0,0 +1,209 @@ +### Linux ### +*~ + +# temporary files which can be created if a process still has a handle open of a deleted file +.fuse_hidden* + +# KDE directory preferences +.directory + +# Linux trash folder which might appear on any partition or disk +.Trash-* + +# .nfs files are created when an open file is removed but is still being accessed +.nfs* + +### OSX ### +# General +.DS_Store +.AppleDouble +.LSOverride + +# Icon must end with two \r +Icon + + +# Thumbnails +._* + +# Files that might appear in the root of a volume +.DocumentRevisions-V100 +.fseventsd +.Spotlight-V100 +.TemporaryItems +.Trashes +.VolumeIcon.icns +.com.apple.timemachine.donotpresent + +# Directories potentially created on remote AFP share +.AppleDB +.AppleDesktop +Network Trash Folder +Temporary Items +.apdisk + +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +### Windows ### +# Windows thumbnail cache files +Thumbs.db +Thumbs.db:encryptable +ehthumbs.db +ehthumbs_vista.db + +# Dump file +*.stackdump + +# Folder config file +[Dd]esktop.ini + +# Recycle Bin used on file shares +$RECYCLE.BIN/ + +# Windows Installer files +*.cab +*.msi +*.msix +*.msm +*.msp + +# Windows shortcuts +*.lnk \ No newline at end of file diff --git a/README.md b/README.md index c17bb1c..c8483ae 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,57 @@ -# image-embedding-operator-template +# Template: Image Embedding Operator -This is another test repo \ No newline at end of file +Authors: + +## **Overview** + +> **Note:** this is just a **template**, not a runnable pipeline. + +This **class template for the image embedding operator** defines the image embedding functions, as well as the standard inputs and outputs. You can complete the operator by filling in the function(`__init__`, `__call__` ) in [image_embedding_operator_class.py](http://./image_embedding_operator_class.py) and update this README file. FYI, [image-embedding-resnet50](https://hub.towhee.io/towhee/image-embedding-resnet50) is based on this template. + +This Operator generates feature vectors from "someone" model, which is trained on "someone" dataset. + +## **Interface** + +```python +__init__(self, model_name: str, framework: str = 'pytorch') +``` + +**Args:** + +- model_name: + - the model name for embedding + - supported types: `str`, for example 'resnet50' + +- framework: + - the framework of the model + - supported types: `str`, default is 'pytorch' + +```python +__call__(self, img_tensor: torch.Tensor) +``` + +**Args:** + +- img_tensor: + - the input image tensor + - supported types: `torch.Tensor` + +**Returns:** + +The Operator returns a tuple `Tuple[('feature_vector', numpy.ndarray)]` containing following fields: + +- feature_vector: + - the embedding of the image + - data type: `numpy.ndarray` + +## **Requirements** + +You can get the required python package by [requirements.txt](https://zilliverse.feishu.cn/docs/requirements.txt). + +## **How it works** + +The `towhee/image-embedding-operator-template` 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-pipeline-template](https://hub.towhee.io/towhee/image-embedding-pipeline-template) pipeline, and it is the red box in the picture below. + +![img](./readme_res/operator.png) + +## **Reference** diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..37f5bd7 --- /dev/null +++ b/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2021 Zilliz. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/config.py b/config.py new file mode 100644 index 0000000..2cd63fa --- /dev/null +++ b/config.py @@ -0,0 +1,3 @@ +TEST_IMG = './test_data/test.jpg' +DIMENSION = 1000 +MODEL_NAME = 'resnet50' \ No newline at end of file diff --git a/image_embedding_operator_template.py b/image_embedding_operator_template.py new file mode 100644 index 0000000..d10755d --- /dev/null +++ b/image_embedding_operator_template.py @@ -0,0 +1,53 @@ +# Copyright 2021 Zilliz. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from typing import NamedTuple +import numpy +from towhee.operator import Operator + + +class ImageEmbeddingOperatorTemplate(Operator): + """ + An operator class that implements image embedding. + + Args: + model_name (`str`): + Embedding the image withe the specific model. + """ + def __init__(self, model_name: str, framework: str = 'pytorch') -> None: + """ + Initialize some parameters and load the model. + """ + super().__init__() + # please define Model class and declare model object, and initialize with the model_name + # if framework == 'pytorch': + # from pytorch import Model + self.model = Model(model_name) + + def __call__(self, img_tensor: torch.Tensor) -> NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]): + """ + Call it when use this class. + + Args: + img_tensor (`torch.Tensor`): + The image tensor, you can generate it with TransformImageOperatorClass. + Returns: + (`numpy.ndarray`) + The image embedding. + """ + # call model with the img_tensor parameter + result = self.model(img_tensor) + Outputs = NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]) + return Outputs(result) diff --git a/image_embedding_operator_template.yaml b/image_embedding_operator_template.yaml new file mode 100644 index 0000000..854f7fd --- /dev/null +++ b/image_embedding_operator_template.yaml @@ -0,0 +1,13 @@ +name: 'image-embedding-operator-template' +labels: + recommended_framework: 'pytorch1.2.0' #the recommend framework, such as + class: 'image-embedding' + others: 'resnet50' #your model name +operator: 'towhee/image-embedding-operator-template' #your operator repo name +init: + model_name: str +call: + input: + img_tensor: torch.Tensor + output: + feature_vector: numpy.ndarray \ No newline at end of file diff --git a/readme_res/operator.png b/readme_res/operator.png new file mode 100644 index 0000000..ca5fff6 Binary files /dev/null and b/readme_res/operator.png differ diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e69de29 diff --git a/test_data/test.jpg b/test_data/test.jpg new file mode 100755 index 0000000..8fdc2b3 Binary files /dev/null and b/test_data/test.jpg differ diff --git a/test_image_embedding_operator_template.py b/test_image_embedding_operator_template.py new file mode 100644 index 0000000..41f7303 --- /dev/null +++ b/test_image_embedding_operator_template.py @@ -0,0 +1,32 @@ +import unittest +from PIL import Image +from torchvision import transforms +from image_embedding_operator_template import ImageEmbeddingOperatorTemplate +from config import DIMENSION, MODEL_NAME, TEST_IMG + + +class TestImageEmbeddingOperatorTemplate(unittest.TestCase): + """ + Simple operator test + """ + 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) + + def test_image_embedding(self): + self.model_name = MODEL_NAME + self.dimension = DIMENSION + op = ImageEmbeddingOperatorTemplate(self.model_name) + print("The output shape of operator:", op(self.img_tensor)[0].shape) + assert (1, self.dimension)==op(self.img_tensor)[0].shape + + +if __name__ == '__main__': + unittest.main()