Sm Jumpstart Foundation Sam 2 1 Video Segmentation
SageMaker JumpStart - deploy VS model
This notebook demonstrates how to use the SageMaker Python SDK to deploy a SageMaker JumpStart VS model and invoke the endpoint.
The model used in this notebook is Meta's SAM 2.1 model. This model is used in segmentation tasks over images and video. This notebook will be demonstrating the video segementation use case using the SAM 2.1 model. We will be using both point and box prompts in various contexts to infer segmentation masks over a sample video. Single object and combination prompts of points and a box is supported as well as the ability to correct masks and re-infer the masks.
NOTE: The masks generated are on a frame by frame basis using surrounding frame inferences to infer on each frame. This means that adding points after propgating the prompts across the video will only correct the alredy inferred masks. To remask a new inference sessions needs to be started.
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.
Select your desired model ID. You can search for available models in the Built-in Algorithms with pre-trained Model Table.
Deploy model
Using the model ID, define your model as a JumpStart model. You can deploy the model on other instance types by passing instance_type to JumpStartModel. See Deploy publicly available foundation models with the JumpStartModel class for more configuration options.
You can now deploy your JumpStart model. The deployment might take few minutes.
SAM2 Endpoint Testing
This notebook demonstrates testing of SAM2 endpoint functionality using modular test classes.
Stream Parser
Due to functionality of SAM2.1 generating masks we need a way to make sure that we return the large payload without reaching sagemaker limits. Therefore we leverage Sagemaker Streaming Responses to stream back chunks of our payload. We expect Json Line items from the endpoint so here we have a StreamParser class to parse the chunks we add to it and parse out complete Json Line items.
Util Funtions
Computer Vision tasks require that we transform videos into formats that are easily transmitable and visualizeable. The following functions are set up to help with those intermediate steps.
Functions:
- encode video (Mandatory)
- encode the binary data of an video into base64 for the endpoint to decode and infer upon.
- decompress_mask (Mandatory)
- as mentioned before we compress and stream the data back and to complement that we have a function to decompress the mask back into its original boolean array. The expected size is 1 channel of the same dimensions of the input image. The method of compresion is compressing a numpy array using zlib then compressing tha using base64. Decompression reverses that by decompressing using base64 and then decompressing using zlib. The resulitng data is loaded into a numpy array from buffer.
- save_visualization (Customizeable)
- we use opencv to apply the masks we get to each frame then stitch the new frames back together back into a video now with the masks.
- NOTE: this function, as is, requires FFmpeg and OpenCV.
Sagemaker Endpoint interaction abstraction
The following functions are completely optional when interacting with the Sagemaker endpoint. We abstracted repetitive code into the following functions to handle both streaming and non streaming requests. This endpoint has been built with Sticky(Stateful) Sessions in mind so operation of the endpoint is dependent on actions specified.
NOTE: The operations to start and end session should be non streaming requests. When starting a new session, the header X-Amzn-SageMaker-Session-Id or the parameter in the invoke_endpoint/invoke_endpoint_with_response_stream function, in the sagemaker runtime client, needs to be NEW_SESSION. Blog Post discussing Sagemaker Sticky Sessions.
SAM2.1 Video Predictor
Session Management
1. Start Session (start_session)
- Purpose: Initializes a new inference session
- Parameters:
path: Path to input image/videoinput_type: 'image' or 'video'session_id: Optional custom identifier
- Returns:
- Session ID
- Success status
- Image/video dimensions
2. Close Session (close_session)
- Purpose: Terminates an active session
- Parameter:
session_id: Session to close
- Returns:
- Success status
Video Segmentation Controls
3. Add Points (add_points)
- Purpose: Adds point prompts for video segmentation
- Parameters:
frame_index: Frame numberobject_id: Object tracking identifierpoints: Point coordinateslabels: Point labelsclear_old_points: Whether to clear existing points
- Returns:
- Frame index
- Object IDs list
- Predicted masks
4. Add Box (add_box)
- Purpose: Adds box prompt with optional points
- Parameters:
frame_index: Frame numberobject_id: Object tracking identifierbox: Box coordinates [x1,y1,x2,y2]points: Optional point coordinateslabels: Optional point labels
- Returns:
- Frame index
- Object IDs list
- Predicted masks
State Management
5. Clear Points in Frame (clear_points_in_frame)
- Purpose: Clears prompts for specific frame
- Parameters:
frame_index: Frame numberobject_id: Object tracking identifier
- Returns:
- Success status
6. Clear Points in Video (clear_points_in_video)
- Purpose: Clears all prompts in video
- Parameters: None
- Returns:
- Success status
7. Propagate in Video (propagate_in_video)
- Purpose: Propagates masks through video frames
- Parameter:
start_frame_index: Starting frame number
- Returns:
- Frame index
- Object IDs list
- Predicted masks
NOTE: When adding points and boxes it is recommended to do this before propogation to get full masks. Any points added after a propogation is called will only edit the masks previously generated. The edits may not be substantial and might only alter the masks minimally. If focus of prompt needs to change it is recommended to end session and start again.
Video Predictor Tests
Test video segmentation functionality with various tracking scenarios.
Download a sample image from jumpstart assets.
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.