Image-Text Retrieval Embdding with ALBEF

author: David Wang

Description

This operator extracts features for image or text with ALBEF which can generate embeddings for text and image by jointly training an image encoder and text encoder to maximize the cosine similarity. This research introduced a contrastive loss to ALign the image and text representations BEfore Fusing (ALBEF) them through cross-modal attention, which enables more grounded vision and language representation learning. This repo is an adaptation from salesforce / ALBEF

Code Example

Load an image from path './teddy.jpg' to generate an image embedding.

Read the text 'A teddybear on a skateboard in Times Square.' to generate an text embedding.

Write the pipeline in simplified style:

from towhee import pipe, ops, DataCollection

img_pipe = (
    pipe.input('url')
    .map('url', 'img', ops.image_decode.cv2_rgb())
    .map('img', 'vec', ops.image_text_embedding.albef(model_name='albef_4m', modality='image')) 
    .output('img', 'vec')
)

text_pipe = (
    pipe.input('text')
    .map('text', 'vec', ops.image_text_embedding.albef(model_name='albef_4m', modality='text')) 
    .output('text', 'vec')
)

DataCollection(img_pipe('./teddy.jpg')).show()
DataCollection(text_pipe('A teddybear on a skateboard in Times Square.')).show()

Factory Constructor

Create the operator via the following factory method

albef(model_name, modality)

Parameters:

​ model_name: str

​ The model name of ALBEF. Supported model names:

• albef_4m
• albef_14m

​ modality: str

​ Which modality(image or text) is used to generate the embedding.

Interface

An image-text embedding operator takes a towhee image or string as input and generate an embedding in ndarray.

Parameters:

​ data: towhee.types.Image (a sub-class of numpy.ndarray) or str

​ The data (image or text based on specified modality) to generate embedding.

Returns: numpy.ndarray

​ The data embedding extracted by model.