Browse Source
add schema info
Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
main
1 changed files with
10 additions and
0 deletions
-
timm_image.py
|
@ -22,6 +22,7 @@ import towhee |
|
|
from towhee.operator.base import NNOperator, OperatorFlag |
|
|
from towhee.operator.base import NNOperator, OperatorFlag |
|
|
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 |
|
|
|
|
|
from towhee.types import Image |
|
|
|
|
|
|
|
|
import torch |
|
|
import torch |
|
|
from torch import nn |
|
|
from torch import nn |
|
@ -144,3 +145,12 @@ class TimmImage(NNOperator): |
|
|
else: |
|
|
else: |
|
|
log.error(f'Invalid format "{format}". Currently supported formats: "pytorch".') |
|
|
log.error(f'Invalid format "{format}". Currently supported formats: "pytorch".') |
|
|
return model_list |
|
|
return model_list |
|
|
|
|
|
|
|
|
|
|
|
def input_schema(self): |
|
|
|
|
|
return [(Image, (-1, -1, 3))] |
|
|
|
|
|
|
|
|
|
|
|
def output_schema(self): |
|
|
|
|
|
image = Image(numpy.random.randn(480, 480, 3), "RGB") |
|
|
|
|
|
ret = self(image) |
|
|
|
|
|
data_type = type(ret.reshape(-1)[0]) |
|
|
|
|
|
return [(data_type, ret.shape)] |
|
|