Introduction to JumpStart - Image Embedding

This notebook’s CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

Set Up
Select a model
Retrieve JumpStart Artifacts & Deploy an Endpoint
Download example images for inference
Query endpoint and parse response
Clean up the endpoint

Note: This notebook was tested on ml.t3.medium instance in Amazon SageMaker Studio with Python 3 (Data Science) kernel and in Amazon SageMaker Notebook instance with conda_python3 kernel.

1. Set Up

[ ]:

! pip install sagemaker ipywidgets --upgrade --quiet

Permissions and environment variables

[ ]:

import sagemaker, boto3, json
from sagemaker import get_execution_role

aws_role = get_execution_role()
aws_region = boto3.Session().region_name
sess = sagemaker.Session()

2. Select a model

Here, we download jumpstart model_manifest file from the jumpstart s3 bucket, filter-out all the Image Embedding models and select a model for inference. ***

[ ]:

from ipywidgets import Dropdown

# download JumpStart model_manifest file.
boto3.client("s3").download_file(
    f"jumpstart-cache-prod-{aws_region}", "models_manifest.json", "models_manifest.json"
)
with open("models_manifest.json", "rb") as json_file:
    model_list = json.load(json_file)

# filter-out all the Image Embedding models from the manifest list.
image_embedding_models = []
for model in model_list:
    model_id = model["model_id"]
    if "-icembedding-" in model_id and model_id not in image_embedding_models:
        image_embedding_models.append(model_id)

# display the model-ids in a dropdown to select a model for inference.
model_dropdown = Dropdown(
    options=image_embedding_models,
    value="tensorflow-icembedding-bit-m-r101x1-ilsvrc2012-featurevector-1",
    description="Select a model",
    style={"description_width": "initial"},
    layout={"width": "max-content"},
)

Chose a model for Inference

[ ]:

display(model_dropdown)

[ ]:

# model_version="*" fetches the latest version of the model
model_id, model_version = model_dropdown.value, "*"

3. Retrieve JumpStart Artifacts & Deploy an Endpoint

[ ]:

from sagemaker.jumpstart.model import JumpStartModel

model = JumpStartModel(model_id=model_id)
model_predictor = model.deploy()

We download example images from the JumpStart assets S3 bucket.

[ ]:

region = boto3.Session().region_name
s3_bucket = f"jumpstart-cache-prod-{region}"
key_prefix = "inference-notebook-assets"
s3 = boto3.client("s3")


def download_from_s3(key_filenames):
    for key_filename in key_filenames:
        s3.download_file(s3_bucket, f"{key_prefix}/{key_filename}", key_filename)


cat_jpg, dog_jpg = "cat.jpg", "dog.jpg"
download_from_s3(key_filenames=[cat_jpg, dog_jpg])

5. Query endpoint and parse response

[ ]:

def query(model_predictor, image_file_name):
    """Query the model predictor."""

    with open(image_file_name, "rb") as file:
        input_img_rb = file.read()

    query_response = model_predictor.predict(
        input_img_rb,
        {
            "ContentType": "application/x-image",
            "Accept": "application/json",
        },
    )
    return query_response


def parse_response(query_response):
    """Parse response and return the embedding."""

    model_predictions = query_response
    translation_text = model_predictions["embedding"]
    return translation_text

[ ]:

from IPython.core.display import HTML

for img_name in [cat_jpg, dog_jpg]:
    query_response = query(model_predictor, img_name)
    embedding = parse_response(query_response)
    first_5element_embeddings = "{" + ", ".join([str(id) for id in embedding[:5]]) + "}"
    display(
        HTML(
            f'<img src={img_name} alt={img_name} align="left" style="width: 250px;"/>'
            f"<figcaption>First-5 elements of the feature vector (embedding) are: {first_5element_embeddings}</figcaption>"
            f"<figcaption>Total length of the feature vector (embedding) is: {len(embedding)}</figcaption>"
        )
    )

[ ]:

# Delete the SageMaker endpoint
model_predictor.delete_model()
model_predictor.delete_endpoint()

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.