Add files

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>

README.md

@@ -1,2 +1,88 @@
-# isc
+# Image Embedding with ISC
+
+*author: Jael Gu*
+
+<br />
+
+## Desription
+
+An image embedding operator generates a vector given an image.
+This operator extracts features for image top ranked models from
+[Image Similarity Challenge 2021](https://github.com/facebookresearch/isc2021) - Descriptor Track.
+The default pretrained model weights are from [The 1st Place Solution of ISC21 (Descriptor Track)](https://github.com/lyakaap/ISC21-Descriptor-Track-1st).
+
+<br />
+
+## Code Example
+
+Load an image from path './towhee.jpg'
+and use the pretrained ISC model ('resnet50') to generate an image embedding. 
+
+ *Write the pipeline in simplified style:*
+
+```python
+import towhee
+
+towhee.glob('./towhee.jpg') \
+      .image_decode() \
+      .image_embedding.isc() \
+      .show()
+```
+<img src="./result1.png" height="50px"/>
+
+*Write a same pipeline with explicit inputs/outputs name specifications:*
+
+```python
+import towhee
+
+towhee.glob['path']('./towhee.jpg') \
+      .image_decode['path', 'img']() \
+      .image_embedding.isc['img', 'vec']() \
+      .select['img', 'vec']() \
+      .show()
+```
+<img src="./result2.png" height="150px"/>
+
+<br />
+
+## Factory Constructor
+
+Create the operator via the following factory method
+
+***image_embedding.isc(skip_preprocess=False, device=None)***
+
+**Parameters:**
+
+***skip_preprocess:*** *bool*
+
+The flag to control whether to skip image preprocess.
+The default value is False.
+If set to True, it will skip image preprocessing steps (transforms).
+In this case, input image data must be prepared in advance in order to properly fit the model.
+
+***device:*** *str*
+
+The device to run this operator, defaults to None.
+When it is None, 'cuda' will be used if it is available, otherwise 'cpu' is used.
+
+<br />
+
+## Interface
+
+An image embedding operator takes a towhee image as input.
+It uses the pre-trained model specified by model name to generate an image embedding in ndarray.
+
+**Parameters:**
+
+***img:*** *towhee.types.Image (a sub-class of numpy.ndarray)*
+
+The decoded image data in numpy.ndarray.
+
+
+
+**Returns:** *numpy.ndarray*
+
+The image embedding extracted by model.
+
+
 

__init__.py

@@ -0,0 +1,19 @@
+# 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 .isc import Isc
+
+
+def isc(**kwargs):
+    return Isc(**kwargs)

isc.py

@@ -0,0 +1,108 @@
+# 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.
+
+import logging
+import os
+import numpy
+from typing import Union, List
+from pathlib import Path
+
+import towhee
+from towhee.operator.base import NNOperator, OperatorFlag
+from towhee.types.arg import arg, to_image_color
+from towhee import register
+from towhee.models import isc
+
+import torch
+from torch import nn
+from torchvision import transforms
+from PIL import Image as PILImage
+import timm
+
+import warnings
+
+warnings.filterwarnings('ignore')
+log = logging.getLogger()
+
+
+@register(output_schema=['vec'])
+class Isc(NNOperator):
+    """
+    The operator uses pretrained ISC model to extract features for an image input.
+
+    Args:
+        skip_preprocess (`bool = False`):
+            Whether skip image transforms.
+    """
+
+    def __init__(self, timm_backbone: str = 'tf_efficientnetv2_m_in21ft1k',
+                 skip_preprocess: bool = False, checkpoint_path: str = None, device: str = None) -> None:
+        super().__init__()
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        self.skip_tfms = skip_preprocess
+        if checkpoint_path is None:
+            checkpoint_path = os.path.join(str(Path(__file__).parent), 'checkpoints', timm_backbone + '.pth')
+
+        backbone = timm.create_model(timm_backbone, features_only=True, pretrained=True)
+        self.model = isc.create_model(pretrained=True, checkpoint_path=checkpoint_path, device=self.device,
+                                      backbone=backbone, p=3.0, eval_p=1.0)
+        self.model.eval()
+
+        self.tfms = transforms.Compose([
+                transforms.Resize((512, 512)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=backbone.default_cfg['mean'],
+                                     std=backbone.default_cfg['std'])
+        ])
+
+    def __call__(self, data: Union[List[towhee._types.Image], towhee._types.Image]):
+        if not isinstance(data, list):
+            imgs = [data]
+        else:
+            imgs = data
+        img_list = []
+        for img in imgs:
+            img = self.convert_img(img)
+            img = img if self.skip_tfms else self.tfms(img)
+            img_list.append(img)
+        inputs = torch.stack(img_list)
+        inputs = inputs.to(self.device)
+        features = self.model(inputs)
+        features = features.to('cpu').flatten(1)
+
+        if isinstance(data, list):
+            vecs = list(features.detach().numpy())
+        else:
+            vecs = features.squeeze(0).detach().numpy()
+        return vecs
+
+    @arg(1, to_image_color('RGB'))
+    def convert_img(self, img: towhee._types.Image):
+        img = PILImage.fromarray(img.astype('uint8'), 'RGB')
+        return img
+
+
+# if __name__ == '__main__':
+#     from towhee import ops
+#
+#     path = 'https://github.com/towhee-io/towhee/raw/main/towhee_logo.png'
+#
+#     decoder = ops.image_decode.cv2()
+#     img = decoder(path)
+#
+#     op = Isc()
+#     out = op(img)
+#     assert out.shape == (256,)

requirements.txt

@@ -0,0 +1,3 @@
+numpy
+torchvision
+timm>=0.5.4