diff --git a/image_embedding_resnet50.py b/image_embedding_resnet50.py new file mode 100644 index 0000000..27de7a7 --- /dev/null +++ b/image_embedding_resnet50.py @@ -0,0 +1,21 @@ +from towhee import pipe, ops, AutoPipes, AutoConfig + +@AutoConfig.register +class ImageEmbeddingConfig: + """ + For UT + """ + def __init__(self): + self.model = 'resnet50' + +@AutoPipes.register +def ImageEmbedding(config): + """ + For UT + """ + return ( + pipe.input('path') + .map('path', 'img', ops.image_decode()) + .map('img', 'embedding', ops.image_embedding.timm(model_name=config.model)) + .output('embedding') + ) \ No newline at end of file