# 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 warnings 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 sys try: from towhee import accelerate except: def accelerate(func): return func import torch import timm from torch import nn from PIL import Image as PILImage from timm.data import create_transform from timm import create_model try: from timm.models import get_pretrained_cfg except ImportError: from timm.models.registry import _model_default_cfgs def get_pretrained_cfg(model_name): return _model_default_cfgs[model_name] warnings.filterwarnings('ignore') log = logging.getLogger('isc_op') _ = sys.modules[__name__] @accelerate class Model: def __init__(self, timm_backbone, checkpoint_path, device): self.device = device self.model = isc.create_model(pretrained=True, checkpoint_path=checkpoint_path, device=self.device, timm_backbone=timm_backbone, p=1.0, eval_p=1.0) self.model.eval() def __call__(self, x): x = x.to(self.device) return self.model(x) @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', img_size: int = 512, checkpoint_path: str = None, skip_preprocess: bool = False, device: str = None) -> None: super().__init__() if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu' self.device = device if isinstance(device, str) else 'cpu' if device < 0 else torch.device(device) self.skip_tfms = skip_preprocess self.timm_backbone = timm_backbone if timm.__version__ < '0.9.0' else 'tf_efficientnetv2_m.in21k_ft_in1k' if checkpoint_path is None: checkpoint_path = os.path.join(str(Path(__file__).parent), 'checkpoints', timm_backbone + '.pth') self.model = Model(self.timm_backbone, checkpoint_path, self.device) self.tfms = create_transform( input_size=img_size, interpolation=self.config['interpolation'], mean=self.config['mean'], std=self.config['std'], crop_pct=self.config['crop_pct'] ) 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 features = self.model(inputs) features = features.to('cpu') if isinstance(data, list): vecs = list(features.detach().numpy()) else: vecs = features.squeeze(0).detach().numpy() return vecs @property def _model(self): return self.model.model @property def config(self): config = get_pretrained_cfg(self.timm_backbone) if hasattr(config, 'to_dict'): config = config.to_dict() return config def save_model(self, format: str = 'pytorch', path: str = 'default'): if path == 'default': path = str(Path(__file__).parent) path = os.path.join(path, 'saved', format) os.makedirs(path, exist_ok=True) name = self.timm_backbone.replace('/', '-') path = os.path.join(path, name) if format in ['pytorch', 'torchscript']: path = path + '.pt' elif format == 'onnx': path = path + '.onnx' else: raise ValueError(f'Invalid format {format}.') dummy_input = torch.rand(1, 3, 224, 224) if format == 'pytorch': torch.save(self._model, path) elif format == 'torchscript': try: try: jit_model = torch.jit.script(self._model) except Exception: jit_model = torch.jit.trace(self._model, dummy_input, strict=False) torch.jit.save(jit_model, path) except Exception as e: log.error(f'Fail to save as torchscript: {e}.') raise RuntimeError(f'Fail to save as torchscript: {e}.') elif format == 'onnx': try: torch.onnx.export(self._model.to('cpu'), dummy_input, path, input_names=['input_0'], output_names=['output_0'], opset_version=14, dynamic_axes={ 'input_0': {0: 'batch_size', 2: 'height', 3: 'width'}, 'output_0': {0: 'batch_size', 1: 'dim'} }, do_constant_folding=True ) except Exception as e: log.error(f'Fail to save as onnx: {e}.') raise RuntimeError(f'Fail to save as onnx: {e}.') # todo: elif format == 'tensorrt': else: log.error(f'Unsupported format "{format}".') return path @arg(1, to_image_color('RGB')) def convert_img(self, img: 'towhee.types.Image'): img = PILImage.fromarray(img.astype('uint8'), 'RGB') return img @property def supported_formats(self): return ['onnx'] def train(self, training_config=None, train_dataset=None, eval_dataset=None, resume_checkpoint_path=None, **kwargs): from .train_isc import train_isc training_args = kwargs.pop('training_args', None) train_isc(self._model, training_args) # 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,)