From 1186063eeda0175dd1b283b0c9538e5bd98f0797 Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Fri, 24 Mar 2023 15:18:33 +0800 Subject: [PATCH] Add image_embedding_resnet50.py Signed-off-by: shiyu22 Add image_embedding_resnet50.py Signed-off-by: shiyu22 Add image_embedding_resnet50.py Signed-off-by: shiyu22 --- image_embedding_resnet50.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 image_embedding_resnet50.py 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