Extract Image using Semantic Segmentation

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 pre-trained model
Run inference on the pre-trained model
Display and Extract Model Predicted Image

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 pre-trained model

Here, we download jumpstart model_manifest file from the jumpstart s3 bucket, filter out all the Semantic Segmentation 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 Semantic Segmentation models from the manifest list.
semseg_models = []
for model in model_list:
    model_id = model["model_id"]
    if "-semseg-" in model_id and model_id not in semseg_models:
        semseg_models.append(model_id)

print(f"\033[38;5;2mChose a model: \033[0;0m\n")

# display the model-ids in a dropdown to select a model for inference.
model_dropdown = Dropdown(
    options=semseg_models,
    value="mxnet-semseg-fcn-resnet101-coco",
    description="",
    style={"description_width": "initial"},
    layout={"width": "max-content"},
)
display(model_dropdown)

[ ]:

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

3. Run inference on the pre-trained model

Using JumpStart, we can perform inference on the pre-trained model, even without fine-tuning it first on a new dataset. ***

3.1. Retrieve JumpStart Artifacts & Deploy an Endpoint

We retrieve the deploy_image_uri, deploy_source_uri, and base_model_uri for the pre-trained model. To host the pre-trained base-model, we create an instance of `sagemaker.model.Model <https://sagemaker.readthedocs.io/en/stable/api/inference/model.html>`__ and deploy it. ***

[ ]:

from sagemaker import image_uris, model_uris, script_uris, hyperparameters
from sagemaker.model import Model
from sagemaker.predictor import Predictor
from sagemaker.utils import name_from_base


endpoint_name = name_from_base(f"jumpstart-example-infer-{model_id}")

inference_instance_type = "ml.m5.xlarge"

# Retrieve the inference docker container uri
deploy_image_uri = image_uris.retrieve(
    region=None,
    framework=None,  # automatically inferred from model_id
    image_scope="inference",
    model_id=model_id,
    model_version=model_version,
    instance_type=inference_instance_type,
)

# Retrieve the inference script uri. This includes scripts for model loading, inference handling etc.
deploy_source_uri = script_uris.retrieve(
    model_id=model_id, model_version=model_version, script_scope="inference"
)


# Retrieve the base model uri
base_model_uri = model_uris.retrieve(
    model_id=model_id, model_version=model_version, model_scope="inference"
)


# Create the SageMaker model instance
model = Model(
    image_uri=deploy_image_uri,
    source_dir=deploy_source_uri,
    model_data=base_model_uri,
    entry_point="inference.py",  # entry point file in source_dir and present in deploy_source_uri
    role=aws_role,
    predictor_cls=Predictor,
    name=endpoint_name,
)

# deploy the Model. Note that we need to pass Predictor class when we deploy model through Model class,
# for being able to run inference through the sagemaker API.
base_model_predictor = model.deploy(
    initial_instance_count=1,
    instance_type=inference_instance_type,
    predictor_cls=Predictor,
    endpoint_name=endpoint_name,
)

### 3.2. Download an example image for inference

We download an example image from the JumpStart assets S3 bucket.

[ ]:

jumpstart_assets_s3_bucket = f"jumpstart-cache-prod-{aws_region}"
pedestrian_img_key_prefix = "inference-notebook-assets"
pedestrian_img = "img_pedestrian.png"

boto3.client("s3").download_file(
    jumpstart_assets_s3_bucket, f"{pedestrian_img_key_prefix}/{pedestrian_img}", pedestrian_img
)

3.3. Query endpoint and parse response

[ ]:

import json


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;verbose",
        },
    )
    return query_response


def parse_response(query_response):
    """Parse response and return predictions as well as the set of all labels and object labels present in the image."""
    response_dict = json.loads(query_response)
    return response_dict["predictions"], response_dict["labels"], response_dict["image_labels"]

[ ]:

query_response = query(base_model_predictor, pedestrian_img)

predictions, labels, image_labels = parse_response(query_response)
print("Objects present in the picture:", image_labels)

### 3.4. Display model predictions

Next, we display the bounding boxes overlaid on the original image. To get color palette for visualization, we borrow the VOC palette implementation from GluonCV

[ ]:

import matplotlib.image as mpimg
from matplotlib import pyplot as plt
import numpy as np
from PIL import Image


def getvocpalette(num_cls):
    """Get a color palette."""

    n = num_cls
    palette = [0] * (n * 3)
    for j in range(0, n):
        lab = j
        palette[j * 3 + 0] = 0
        palette[j * 3 + 1] = 0
        palette[j * 3 + 2] = 0
        i = 0
        while lab > 0:
            palette[j * 3 + 0] |= ((lab >> 0) & 1) << (7 - i)
            palette[j * 3 + 1] |= ((lab >> 1) & 1) << (7 - i)
            palette[j * 3 + 2] |= ((lab >> 2) & 1) << (7 - i)
            i = i + 1
            lab >>= 3
    return palette


def display_predictions(predictions):
    """Display predictions with each pixel subsituted by the color of the corresponding label."""

    palette = getvocpalette(256)
    npimg = np.array(predictions)
    npimg[npimg == -1] = 255
    mask = Image.fromarray(npimg.astype("uint8"))

    mask.putpalette(palette)
    mask.save("Mask_putput.png")
    mmask = mpimg.imread("Mask_putput.png")
    plt.imshow(mask)

[ ]:

display_predictions(predictions)

4. Display and Extract Model Predicted Image

4.1. Extract Model Predicted Image

[ ]:

# Convert the pedestrian Image to JPEG so we get the pixel details
im1 = Image.open(r"img_pedestrian.png")
rgb_im1 = im1.convert("RGB")
rgb_im1.save("img_pedestrian-1.jpg")

# Convert the mask output to JPEG so we get the pixel details
im2 = Image.open(r"Mask_putput.png")
rgb_im2 = im2.convert("RGB")
rgb_im2.save("Mask_putput-1.jpg")

im1 = Image.open("img_pedestrian-1.jpg")
pixelMap1 = im1.load()
im2 = Image.open("Mask_putput-1.jpg")
pixelMap2 = im2.load()
width1, height1 = im1.size
width2, height2 = im2.size


# Set the height and width to the largest image

if width2 > width1:
    width = width2
else:
    width = width1

if height2 > height1:
    height = height2
else:
    height = height1


for i in range(width):  # for every col:
    for j in range(height):  # For every row
        R1, G1, B1 = pixelMap2[i, j]
        if R1 != 0 and G1 != 0 and B1 != 0:
            pixelMap2[i, j] = pixelMap1[i, j]
        else:
            pixelMap2[i, j] = (0, 0, 0)


im2 = im2.save("Final_Pedestrians.jpg")

Let’s display the original pedestrian image with the background included ___

[ ]:

im1.show()

### 4.2. Display Extracted Model Predicted Image

Display the pedestrian image after extraction ___

[ ]:

im3 = Image.open("Final_Pedestrians.jpg")
im3.show()

4.3. Clean up the endpoint

[ ]:

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

Notebook CI Test Results

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.