From e3e02e24ac9efb965a548c51e0882e345770cca8 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Fri, 1 Apr 2022 16:53:50 +0800 Subject: [PATCH] Optimize image type Signed-off-by: Jael Gu --- timm_image.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/timm_image.py b/timm_image.py index 3d2c84d..32ed98d 100644 --- a/timm_image.py +++ b/timm_image.py @@ -15,8 +15,8 @@ import logging import numpy +import towhee from towhee.operator.base import NNOperator, OperatorFlag -from towhee.types import Image as towheeImage from towhee.types.arg import arg, to_image_color from towhee import register @@ -59,7 +59,7 @@ class TimmImage(NNOperator): self.skip_tfms = skip_preprocess @arg(1, to_image_color('RGB')) - def __call__(self, img: towheeImage) -> numpy.ndarray: + def __call__(self, img: towhee._types.Image) -> numpy.ndarray: img = PILImage.fromarray(img.astype('uint8'), 'RGB') if not self.skip_tfms: img = self.tfms(img).unsqueeze(0)