logo
Browse Source

Update pipelines with pydantic

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 12 months ago
parent
commit
18eab743cc
  1. 25
      image_embedding.py

25
image_embedding.py

@ -1,19 +1,20 @@
from typing import Optional
from pydantic import BaseModel
from towhee import pipe, ops, AutoPipes, AutoConfig
@AutoConfig.register
class ImageEmbeddingConfig:
def __init__(self):
# config for audio_decode.ffmpeg
self.mode = 'BGR'
# config for audio_embedding.vggish
self.model_name = 'resnet50'
self.num_classes: int = 1000
self.skip_preprocess: bool = False
# config for triton
self.device = -1
class ImageEmbeddingConfig(BaseModel):
# config for image_decode
mode: Optional[str] = 'BGR'
# config for image_embedding
model_name: Optional[str] = 'resnet50'
num_classes: Optional[int] = 1000
skip_preprocess: Optional[bool] = False
# config for triton
device: Optional[int] = -1
@AutoPipes.register

Loading…
Cancel
Save