From b84b7e54fc5b0e6260cce48621cc38f16ea87d55 Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Wed, 15 Dec 2021 22:19:01 +0800 Subject: [PATCH] Update Signed-off-by: shiyu22 --- __init__.py | 12 ++++++------ vit_image_embedding.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/__init__.py b/__init__.py index 8516979..9379ed6 100644 --- a/__init__.py +++ b/__init__.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os +# import os -# For requirements. -try: - import timm -except ModuleNotFoundError: - os.system('pip install timm') +# # For requirements. +# try: +# import timm +# except ModuleNotFoundError: +# os.system('pip install timm') diff --git a/vit_image_embedding.py b/vit_image_embedding.py index a416bba..ed88e2b 100644 --- a/vit_image_embedding.py +++ b/vit_image_embedding.py @@ -18,8 +18,6 @@ from typing import NamedTuple from pathlib import Path from PIL import Image import torch -from timm.data import resolve_data_config -from timm.data.transforms_factory import create_transform import numpy from towhee.operator import Operator @@ -41,6 +39,8 @@ class VitImageEmbedding(Operator): sys.path.append(str(Path(__file__).parent)) if framework == 'pytorch': from pytorch.model import Model + from timm.data import resolve_data_config + from timm.data.transforms_factory import create_transform self.model = Model(model_name, weights_path) config = resolve_data_config({}, model=self.model._model) self.tfms = create_transform(**config)