diff --git a/image_embedding_resnet50.py b/image_embedding_resnet50.py index a863783..63ad3e4 100644 --- a/image_embedding_resnet50.py +++ b/image_embedding_resnet50.py @@ -1,15 +1,10 @@ -from towhee import pipe, ops, AutoPipes, AutoConfig - -@AutoConfig.register -class ImageEmbeddingConfig: - def __init__(self): - self.model = 'resnet50' +from towhee import pipe, ops, AutoPipes @AutoPipes.register -def ImageEmbedding(config): +def ImageEmbedding(): return ( pipe.input('path') .map('path', 'img', ops.image_decode()) - .map('img', 'embedding', ops.image_embedding.timm(model_name=config.model)) + .map('img', 'embedding', ops.image_embedding.timm(model_name='resnet50')) .output('embedding') ) \ No newline at end of file