Step-by-step tutorial on how to integrate custom video object segmentation neural network into Supervisely platform on the example of XMem.
Introduction
In this tutorial you will learn how to integrate your video object segmentation model into Supervisely Ecosystem. Supervisely Python SDK allows to integrate models for numerous video object tracking tasks, such as tracking of bounding boxes, masks, keypoints, polylines, etc. This tutorial takes XMem video object segmentation model as an example and provides a complete instruction to integrate it as an application into Supervisely Ecosystem. You can find and try XMem Supervisely integration here.
Implementation details
To integrate your custom video object segmentation model, you need to subclass sly.nn.inference.MaskTracking and implement 2 methods:
load_on_device method for loading the weights and initializing the model on a specific device. Takes a model_dir argument, which is a directory for all model files (like configs, weights, etc.), and a device argument - a torch.device like cuda:0, cpu.
predict method for model inference. It takes a frames argument - a list of numpy arrays, which represents a set of video frames, and an input_mask agrument - a mask with the objects in the first frame of the video. These objects will be tracked on all input frames. It should be a numpy array of shape (H, W), where 0 values represent the background, and other numbers represent the target objects (for example, if you have 2 target objects, than input_mask array will consist of 0, 1 and 2 values).
Overall structure
The overall structure of the class we will implement looks like this:
import supervisely as slyimport torchimport numpy as npclassMyModel(sly.nn.inference.MaskTracking):defload_on_device(self,model_dir:str,device: Literal["cpu","cuda","cuda:0","cuda:1","cuda:2","cuda:3"]="cpu",):# initialize model, load weights, load model on devicepassdefpredict(self,frames: List[np.ndarray],input_mask: np.ndarray,)-> List[np.ndarray]:# a simple code example# disable gradient calculation torch.set_grad_enabled(False) results =[]# pass input mask to your model, run it on given list of frames (frame-by-frame)for frame in frames: prediction =self.model(input_mask, frame)# save predictions to a list results.append(prediction)# update progress bar on each iterationself.video_interface._notify(task="mask tracking")return results
The superclass has a serve method. For running the code on the Supervisely platform, serve method should be executed:
The serve method deploys your model as a REST API service on the Supervisely platform. It means that other applications are able to send requests to your model and get predictions from it.
XMem video object segmentation model
Now let's implement the class specifically for XMem.
Getting started
Step 1. Prepare ~/supervisely.env file with credentials. Learn more here
Step 4. Open the repository directory in Visual Studio Code.
Step-by-step implementation
Creating necessary files and directories
After cloning original repo we will create supervisely_integration folder, where all code for integration will be stored. There will be 2 directories - docker (we will put our Dockerfile here) and serve (app directory). Inside serve directory we will create src subdirectory and put main.py file there. Inside serve folder we will also create debug.env file - it will contain your team id:
We will also create requirements.txt file, where all app dependencies will be stored:
Now we can start coding our main.py file.
Defining imports and global variables
1. load_on_device
The following code creates XMem model with default hyperparameters recommended by original repository and defines resolution to which input video will be resized (we will use 480 as in original work). Also load_on_device will keep the model as a self.model and the device as self.device for further use:
For local debug we can load model weights from local storage, but in production we recommend to save weights to a Docker image.
2. predict
The core method for model inference. Here we are disabling gradient calculation, resizing input mask and frames via interpolation, inference XMem model frame-by-frame, saving postprocessed predictions to a list and updating progress bar on every iteration.
The method must return a list of numpy arrays with a length equal to the number of input frames. Each array is a predicted mask of shape (H, W), which represents the objects in one frame. In other words it should have format similar to the input_mask. For instance, if you're tracking two objects over 20 frames, your input frames variable will be a list of 20 numpy arrays, the input_mask will be a numpy array with shape (H, W), containing values of 0, 1, and 2. Similarly, the results variable will contain a list of 20 numpy arrays, with each individual array also shaped (H, W) and filled with 0, 1, and 2 values.
In the end of each iteration we update a progress bar via self.video_interface._notify(task="mask tracking") - it is necessary for app UI to look correctly:
It is crucial to disable gradient calculation in predict method, not in load_on_device, because these methods are being executed in different threads, so if you try disabling gradient calculation in load_on_device method, then it will have no effect during inference, which can significantly increase GPU memory consumption.
When load_on_device and predict methods are implemented, it is necessary to initialize our model class and execute serve method:
Debug in Supervisely platform
Once the code is written, it's time to test it right in the Supervisely platform as a debugging app.
First of all it is necessary to create .vscode folder and launch.json file inside this folder. Your launch.json file should contain the following:
If you develop in a Docker container, you should run the container with --cap_add=NET_ADMIN option.
Install sudo apt-get install wireguard iproute2 or brew install wireguard-tools for Mac.
Define your TEAM_ID in the debug.env file. *Actually there are other env variables that is needed, but they are already provided in ./vscode/launch.json for you.
Switch the launch.json config to the Advanced debug in Supervisely platform:
Advanced Debug in Supervisely
Run the code.
✅ It will deploy the model in the Supervisely platform as a REST API.
Here is how advanced debug mode launch looks like:
After advanced debug launch you must be able to debug your app via Develop & Debug app:
supervisely_integration/serve/src/main.py - main inference script
supervisely_integration/serve/README.md - readme of your application, it is the main page of an application in Ecosystem with some images, videos, and how-to-use guides
supervisely_integration/serve/config.json - configuration of the Supervisely application, which defines the name and description of the app, its context menu, icon, poster, and running settings
supervisely_integration/serve/requirements.txt - all packages needed for debugging
supervisely_integration/serve/debug.env - file with variables used for debugging
supervisely_integration/docker - directory with the custom Dockerfile for this application and the script that builds it and publishes it to the docker registry
App configuration
App configuration is stored in config.json file. A detailed explanation of all possible fields is covered in this Configuration Tutorial. Let's check the config for our current app:
Here is the explanation for the fields:
type - type of the module in Supervisely Ecosystem
version - version of Supervisely App Engine. Just keep it by default
name - the name of the application
description - the description of the application
categories - these tags are used to place the application in the correct category in Ecosystem.
session_tags - these tags will be assigned to every running session of the application. They can be used by other apps to find and filter all running sessions
need_gpu: true - should be true if you want to use any cuda devices
gpu: required - app can be runned on both CPU and GPU devices, but it is recommended to use GPU for higher inference speed
community_agent: false - this means that this app can not be run on the agents started by Supervisely team, so users have to connect their own computers and run the app only on their own agents. Only applicable in Community Edition. Enterprise customers use their private instances so they can ignore the current option
docker_image - Docker container will be started from the defined Docker image, github repository will be downloaded and mounted inside the container.
entrypoint - the command that starts our application in a container
port - port inside the container
headless: true means that the app has no User Interface
allowed_shapes - shapes can be tracked with this model. In Supervisely masks can be represented by bitmap and polygon geometries.
App release
Once you've tested the code, it's time to release it into the platform. It can be released as an App that is shared with the all Supervisely community, or as your own private App.
import supervisely as sly
import os
from dotenv import load_dotenv
from typing_extensions import Literal
from typing import List
import numpy as np
import torch
from model.network import XMem
from inference.inference_core import InferenceCore
from dataset.range_transform import im_normalization
from inference.interact.interactive_utils import index_numpy_to_one_hot_torch
# for debug, has no effect in production
load_dotenv("supervisely_integration/serve/debug.env")
load_dotenv(os.path.expanduser("~/supervisely.env"))
weights_location_path = "/weights/XMem.pth"
class XMemTracker(sly.nn.inference.MaskTracking):
def load_on_device(
self,
model_dir: str,
device: Literal["cpu", "cuda", "cuda:0", "cuda:1", "cuda:2", "cuda:3"] = "cpu",
):
# define model configuration (default hyperparameters)
self.config = {
"top_k": 30,
"mem_every": 5,
"deep_update_every": -1,
"enable_long_term": True,
"enable_long_term_count_usage": True,
"num_prototypes": 128,
"min_mid_term_frames": 5,
"max_mid_term_frames": 10,
"max_long_term_elements": 10000,
}
# define resolution to which input video will be resized (was taken from original repository)
self.resolution = 480
# build model
self.device = torch.device(device)
self.model = XMem(self.config, weights_location_path, map_location=self.device).eval()
self.model = self.model.to(self.device)
def predict(
self,
frames: List[np.ndarray],
input_mask: np.ndarray,
) -> List[np.ndarray]:
# disable gradient calculation
torch.set_grad_enabled(False)
# empty cache
if torch.cuda.is_available():
torch.cuda.empty_cache()
# object IDs should be consecutive and start from 1 (0 represents the background)
num_objects = len(np.unique(input_mask)) - 1
# load processor
processor = InferenceCore(self.model, config=self.config)
processor.set_all_labels(range(1, num_objects + 1))
# resize input mask
original_width, original_height = input_mask.shape[1], input_mask.shape[0]
scaler = min(original_width, original_height) / self.resolution
resized_width = int(original_width / scaler)
resized_height = int(original_height / scaler)
input_mask = torch.from_numpy(input_mask)
input_mask = input_mask.view(1, 1, input_mask.shape[0], input_mask.shape[1])
input_mask = torch.nn.functional.interpolate(input_mask, (resized_height, resized_width), mode="nearest")
input_mask = input_mask.squeeze().numpy()
results = []
# track input objects' masks
with torch.cuda.amp.autocast(enabled=True):
for i, frame in enumerate(frames):
# preprocess frame
frame = frame.transpose(2, 0, 1)
frame = torch.from_numpy(frame)
frame = torch.unsqueeze(frame, 0)
frame = torch.nn.functional.interpolate(frame, (resized_height, resized_width), mode="nearest")
frame = frame.squeeze()
frame = frame.float().to(self.device) / 255
frame = im_normalization(frame)
# inference model on a specific frame
if i == 0:
# preprocess input mask
input_mask = index_numpy_to_one_hot_torch(input_mask, num_objects + 1)
# the background mask is not fed into the model
input_mask = input_mask[1:]
input_mask = input_mask.to(self.device)
prediction = processor.step(frame, input_mask)
else:
prediction = processor.step(frame)
# postprocess prediction
prediction = torch.argmax(prediction, dim=0)
prediction = prediction.cpu().to(torch.uint8)
prediction = prediction.view(1, 1, prediction.shape[0], prediction.shape[1])
prediction = torch.nn.functional.interpolate(prediction, (original_height, original_width), mode="nearest")
prediction = prediction.squeeze().numpy()
# save predicted mask
results.append(prediction)
# update progress bar
self.video_interface._notify(task="mask tracking")
return results
