diff --git a/.README.md.swp b/.README.md.swp new file mode 100644 index 0000000..f0871f0 Binary files /dev/null and b/.README.md.swp differ diff --git a/README.md b/README.md index fa47a5d..379ec2b 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,68 @@ -# inceptionresnetv1 +# Inception-ResNet v1 Face Embedding Operator + +*author: David Wang* + + + +## Desription + +This operator extracts embedding vector from facial image using [Inception-ResNet](https://arxiv.org/pdf/1602.07261.pdf). The implementation is an adaptation from [timesler/facenet-pytorch](https://github.com/timesler/facenet-pytorch). + +## Code Example + +Extract face image embedding from './img.png'. + + *Write the pipeline in simplified style*: + +```python +import towhee + +towhee.glob('./img.png') \ + .image_decode.cv2() \ + .face_embedding.inceptionresnetv1() \ + .tolist() +``` + +*Write a same pipeline with explicit inputs/outputs name specifications:* + +```python +import towhee + +towhee.glob['path']('./img.png') \ + .image_decode.cv2['path', 'img']() \ + .face_embedding.inceptionresnetv1['img', 'vec']() \ + .select('img','vec') \ + .show() +``` +result + + +## Factory Constructor + +Create the operator via the following factory method + +***face_embedding.inceptionresnetv1(image_size = 160)*** + +**Parameters:** + +***image_size***: *int* + +Scaled input image size to extract embedding. The higher resolution would generate the more discriminateive feature but cost more time to calculate. + +supported types: `int`, default is 160. + + +## Interface + +A face embedding operator takes a face image as input. It extracts the embedding back to ndarray. + +**Parameters:** + +​ ***img***: *towhee.types.Image (a sub-class of numpy.ndarray)* + +​ The input image. + +**Returns**: *numpy.ndarray* + +​ The extracted image embedding. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ee3d38b --- /dev/null +++ b/__init__.py @@ -0,0 +1,26 @@ +# 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 os + +# # For requirements. +# try: +# import efficientnet_pytorch +# except ModuleNotFoundError: +# os.system('pip install efficientnet_pytorch') +from .inceptionresnetv1 import Inceptionresnetv1 + +def inceptionresnetv1(image_size = 160 ): + return Inceptionresnetv1(image_size) + diff --git a/inceptionresnetv1.py b/inceptionresnetv1.py new file mode 100644 index 0000000..ceddb4a --- /dev/null +++ b/inceptionresnetv1.py @@ -0,0 +1,63 @@ +# 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 as np +import torch +import cv2 +from facenet_pytorch import InceptionResnetV1 + +from towhee import register +from towhee.operator import NNOperator +from towhee.types.image_utils import to_pil +from towhee._types import Image +from towhee.types import arg, to_image_color + +@register(output_schema=['vec']) +class Inceptionresnetv1(NNOperator): + """ + comment placeholder + """ + def __init__(self, image_size = 160): + self.image_size = image_size + self._model = InceptionResnetV1(pretrained='vggface2') + self._model.eval() + + @arg(1, to_image_color('RGB') ) + def __call__(self, img: Image) -> np.ndarray: + img = self.preprocess(img) + embs = self._model(torch.FloatTensor(img).permute(0,3,1,2)).detach().numpy() + return embs + + def preprocess(self, img: Image): + #img shape expected to be [n, h, w, c=3] or [h, w, c=3]. + if len(img.shape) == 3: + img = cv2.resize(img, (self.image_size, self.image_size)) + img = np.expand_dims(img, 0) + elif len(img.shape) == 4: + pass + else: + raise ValueError('unknown tensor shape, need to be [n, h, w, c=3] or [h, w, c=3].') + img = self._fixed_image_standardization(img) + return img + + def _fixed_image_standardization(self, image_tensor): + processed_tensor = (image_tensor - 127.5) / 128.0 + return processed_tensor + + def train(self): + """ + For training model + """ + pass diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c8c292b --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +opencv-python>=4.1.2 +torch>=1.7.0 +torchvision>=0.8.1 +numpy +facenet-pytorch>=2.5.2 diff --git a/result.png b/result.png new file mode 100644 index 0000000..ebd27d5 Binary files /dev/null and b/result.png differ