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')
    )