logo
Browse Source

Optimize image type

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
test
Jael Gu 3 years ago
parent
commit
e3e02e24ac
  1. 4
      timm_image.py

4
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)

Loading…
Cancel
Save