diff --git a/README.md b/README.md index 70dfc7d..16787d1 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,88 @@ -# isc +# Image Embedding with ISC + +*author: Jael Gu* + +
+ +## 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). + +
+ +## 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() +``` + + +*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() +``` + + +
+ +## 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. + +
+ +## 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. + + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..57f9766 --- /dev/null +++ b/__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) diff --git a/checkpoints/tf_efficientnetv2_m_in21ft1k.pth b/checkpoints/tf_efficientnetv2_m_in21ft1k.pth new file mode 100644 index 0000000..0e11a40 --- /dev/null +++ b/checkpoints/tf_efficientnetv2_m_in21ft1k.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:512390a15ab52f8ab5adc14a315b157a5d6e732e3361c7a0dce5be6e95a86d30 +size 420418996 diff --git a/isc.py b/isc.py new file mode 100644 index 0000000..101d4a7 --- /dev/null +++ b/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,) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e26e6c2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +numpy +torchvision +timm>=0.5.4