Browse Source
Update pipelines with pydantic
Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
1 changed files with
13 additions and
12 deletions
-
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 |
|
|
|