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 logging
import numpy import numpy
import towhee
from towhee.operator.base import NNOperator, OperatorFlag 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.types.arg import arg, to_image_color
from towhee import register from towhee import register
@ -59,7 +59,7 @@ class TimmImage(NNOperator):
self.skip_tfms = skip_preprocess self.skip_tfms = skip_preprocess
@arg(1, to_image_color('RGB')) @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') img = PILImage.fromarray(img.astype('uint8'), 'RGB')
if not self.skip_tfms: if not self.skip_tfms:
img = self.tfms(img).unsqueeze(0) img = self.tfms(img).unsqueeze(0)

Loading…
Cancel
Save