logo
Browse Source

add schema info

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
main
junjie.jiang 2 years ago
parent
commit
2989fee8f3
  1. 10
      timm_image.py

10
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)]

Loading…
Cancel
Save