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.
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
[ ]:
from sagemaker.jumpstart.model import JumpStartModel
model = JumpStartModel(model_id=model_id)
model_predictor = model.deploy()
### 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 = query_response
return response_dict["predictions"], response_dict["labels"], response_dict["image_labels"]
[ ]:
query_response = query(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
model_predictor.delete_model()
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.
