Browse Source
Optimize image type
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
test
1 changed files with
2 additions and
2 deletions
-
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) |
|
|
|
