Semantic Segmentation

In this example, we use CVCUDA to accelerate the pre and post processing pipelines in the deep learning inference use case involving a semantic segmentation model. The deep learning model can utilize either PyTorch or TensorRT to run the inference. The pre-processing pipeline converts the input into the format required by the input layer of the model whereas the post processing pipeline converts the output produced by the model into a visualization-friendly frame. We use the FCN ResNet101 model (from torchvision) to generate the predictions. This sample can work on a single image or a folder full of images or on a single video. All images have to be in the JPEG format and with the same dimensions unless run under the batch size of one. Video has to be in mp4 format with a fixed frame rate. We use the torchnvjpeg library to read the images and NVIDIA’s Video Processing Framework (VPF) to read/write videos.

The exact pre-processing operations are:

Decode Data -> Resize -> Convert Datatype(Float) -> Normalize (to 0-1 range, mean and stddev) -> convert to NCHW

The exact post-processing operations are:

Create Binary mask -> Upscale the mask -> Blur the input frames -> Joint Bilateral filter to smooth the mask -> Overlay the masks onto the original frame

Writing the Sample App

The segmentation sample app has been designed to be modular in all aspects. It imports and uses various modules such as data decoders, encoders, pipeline pre and post processors and the model inference. Some of these modules are defined in the same folder as the sample whereas the rest are defined in the common scripts folder for a wider re-use.

1. Modules used by this sample app that are defined in the common folder (i.e. not specific just to this sample) are the ImageBatchDecoder and ImageBatchEncoder for nvImageCodec based image decoding and encoding and VideoBatchDecoder and VideoBatchEncoder for PyNvVideoCodec based video decoding and encoding.

2. Modules specific to this sample (i.e. defined in the segmentation sample folder) are PreprocessorCvcuda and PostprocessorCvcuda for CVCUDA based pre and post processing pipelines and SegmentationPyTorch and SegmentationTensorRT for the model inference.

The first stage in our pipeline is importing all necessary python modules. Apart from the modules described above, this also includes modules such as torch and torchvision, torchnvjpeg, vpf and the main package of CVCUDA among others. Be sure to import pycuda.driver before importing any other GPU packages like torch or cvcuda to ensure a proper initialization.

# NOTE: One must import PyCuda driver first, before CVCUDA or VPF otherwise
# things may throw unexpected errors.
import pycuda.driver as cuda
import os
import sys
import logging
import cvcuda
import torch

# Bring the commons folder from the samples directory into our path so that
# we can import modules from it.
common_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
    "common",
    "python",
)
sys.path.insert(0, common_dir)

from perf_utils import (  # noqa: E402
    CvCudaPerf,
    get_default_arg_parser,
    parse_validate_default_args,
)

from nvcodec_utils import (  # noqa: E402
    VideoBatchDecoder,
    VideoBatchEncoder,
    ImageBatchDecoder,
    ImageBatchEncoder,
)

from pipelines import (  # noqa: E402
    PreprocessorCvcuda,
    PostprocessorCvcuda,
)

from model_inference import (  # noqa: E402
    SegmentationPyTorch,
    SegmentationTensorRT,
)

Then we define the main function which helps us parse various configuration parameters used throughout this sample as command line arguments. This sample allows configuring following parameters. All of them have their default values already set so that one can execute the sample without supplying any. Some of these arguments are shared across many other CVCUDA samples and hence come from the perf_utils.py class’s get_default_arg_parser() method.

1. -i, --input_path : Either a path to a JPEG image/MP4 video or a directory containing JPG images to use as input. When pointing to a directory, only JPG images will be read.

2. -o, --output_dir : The directory where the output segmentation overlay should be stored.

3. -c, --class_name : The segmentation class to visualize the results for.

4. -th, --target_img_height : The height to which you want to resize the input_image before running inference.

5. -tw, --target_img_width : The width to which you want to resize the input_image before running inference.

6. -b, --batch_size : The batch size used during inference. If only one image is used as input, the same input image will be read and used this many times. Useful for performance bench-marking.

7. -d, --device_id : The GPU device to use for this sample.

8. -bk, --backend : The inference backend to use. Currently supports PyTorch or TensorRT.

Once we are done parsing all the command line arguments, we would setup the CvCudaPerf object for any performance benchmarking needs and simply call the function run_sample with all those arguments.

cvcuda_perf = CvCudaPerf("segmentation_sample", default_args=args)
run_sample(
    args.input_path,
    args.output_dir,
    args.class_name,
    args.batch_size,
    args.target_img_height,
    args.target_img_width,
    args.device_id,
    args.backend,
    cvcuda_perf,
)

The run_sample function is the primary function that runs this sample. It sets up the requested CUDA device, CUDA context and CUDA stream. CUDA streams help us execute CUDA operations on a non-default stream and enhances the overall performance. Additionally, NVTX markers are used throughout this sample (via CvCudaPerf) to facilitate performance bench-marking using NVIDIA NSIGHT systems and benchmark.py. In order to keep things simple, we are only creating one CUDA stream to run all the stages of this sample. The same stream is available in CVCUDA, PyTorch and TensorRT.

cvcuda_perf.push_range("run_sample")

# Define the objects that handle the pipeline stages
image_size = (target_img_width, target_img_height)

# Define the cuda device, context and streams.
cuda_device = cuda.Device(device_id)
cuda_ctx = cuda_device.retain_primary_context()
cuda_ctx.push()
# Use the the default stream for cvcuda and torch
# Since we never created a stream current will be the CUDA default stream
cvcuda_stream = cvcuda.Stream().current
torch_stream = torch.cuda.default_stream(device=cuda_device)

Next, we instantiate various classes to help us run the sample. These classes are:

1. PreprocessorCvcuda : A CVCUDA based pre-processing pipeline for semantic segmentation.

2. ImageBatchDecoder : A nvImageCodec based image decoder to read the images.

3. ImageBatchEncoder : A nvImageCodec based image encoder to write the images.

4. VideoBatchDecoder : A PyNvVideoCodec based video decoder to read the video.

5. VideoBatchEncoder : A PyNvVideoCodec based video encoder to write the video.

6. PostprocessorCvcuda : A CVCUDA based post-processing pipeline for semantic segmentation.

7. SegmentationPyTorch : A PyTorch based semantic segmentation model to execute inference.

8. SegmentationTensorRT : A TensorRT based semantic segmentation model to execute inference.

These classes are defined in modular fashion and exposes a uniform interface which allows easy plug-and-play in appropriate places. For example, one can use the same API to decode/encode images using PyTorch as that of decode/encode videos using VPF. Similarly, one can invoke the inference in the exact same way with PyTorch as with TensorRT.

Additionally, the encoder and decoder interfaces also exposes start and join methods, making it easy to upgrade them to a multi-threading environment (if needed.) Such multi-threading capabilities are slated for a future release.

# Now define the object that will handle pre-processing
preprocess = PreprocessorCvcuda(device_id, cvcuda_perf)

if os.path.splitext(input_path)[1] == ".jpg" or os.path.isdir(input_path):
    # Treat this as data modality of images
    decoder = ImageBatchDecoder(
        input_path,
        batch_size,
        device_id,
        cuda_ctx,
        cvcuda_stream,
        cvcuda_perf,
    )

    encoder = ImageBatchEncoder(
        output_dir,
        device_id=device_id,
        cvcuda_perf=cvcuda_perf,
    )
else:
    # Treat this as data modality of videos
    decoder = VideoBatchDecoder(
        input_path,
        batch_size,
        device_id,
        cuda_ctx,
        cvcuda_stream,
        cvcuda_perf,
    )

    encoder = VideoBatchEncoder(
        output_dir,
        decoder.fps,
        device_id,
        cuda_ctx,
        cvcuda_stream,
        cvcuda_perf,
    )

# Define the post-processor
postprocess = PostprocessorCvcuda(
    encoder.input_layout,
    encoder.gpu_input,
    device_id,
    cvcuda_perf,
)

# Setup the segmentation models.
if backend == "pytorch":
    inference = SegmentationPyTorch(
        output_dir,
        class_name,
        batch_size,
        image_size,
        device_id,
        cvcuda_perf,
    )
elif backend == "tensorrt":
    inference = SegmentationTensorRT(
        output_dir,
        class_name,
        batch_size,
        image_size,
        device_id,
        cvcuda_perf,
    )
else:
    raise ValueError("Unknown backend: %s" % backend)

With all of these components initialized, the overall data flow per a data batch looks like the following:

Decode batch -> Preprocess Batch -> Run Inference -> Post Process Batch -> Encode batch

# Define and execute the processing pipeline ------------
cvcuda_perf.push_range("pipeline")

# Fire up encoder/decoder
decoder.start()
encoder.start()

# Loop through all input frames
batch_idx = 0
while True:
    cvcuda_perf.push_range("batch", batch_idx=batch_idx)
    # Make sure that cvcuda and torch are using the same stream
    with cvcuda_stream, torch.cuda.stream(torch_stream):
        # Stage 1: decode
        batch = decoder()
        if batch is None:
            cvcuda_perf.pop_range(total_items=0)  # for batch
            break  # No more frames to decode
        assert batch_idx == batch.batch_idx

        logger.info("Processing batch %d" % batch_idx)

        # Stage 2: pre-processing
        orig_tensor, resized_tensor, normalized_tensor = preprocess(
            batch.data,
            out_size=image_size,
        )

        # Stage 3: inference
        probabilities = inference(normalized_tensor)

        # Stage 4: post-processing
        blurred_frame = postprocess(
            probabilities,
            orig_tensor,
            resized_tensor,
            inference.class_index,
        )

        # Stage 5: encode
        batch.data = blurred_frame
        encoder(batch)

        batch_idx += 1

    cvcuda_perf.pop_range(total_items=batch.data.shape[0])  # for batch

# Make sure encoder finishes any outstanding work
encoder.join()

cvcuda_perf.pop_range()  # for pipeline

cuda_ctx.pop()

That’s it for the semantic segmentation sample. To understand more about how each stage in the pipeline works, please explore the following sections:

• PreprocessorCvcuda
• PostprocessorCvcuda
• ImageBatchDecoder
• ImageBatchEncoder
• VideoBatchDecoder
• VideoBatchEncoder
• SegmentationPyTorch
• SegmentationTensorRT

Running the Sample

This sample can be invoked without any command-line arguments like the following. In that case it will use the default values. It uses TensorRT as the inference engine, Weimaraner.jpg as the input image, writes the output overlay for the background class to the /tmp directory with batch size of 4. Upon the first run, generating serialized TensorRT model may take some time for a given batch size.

python3 segmentation/python/main.py

To run it on a single image with batch size 5 for the background class writing the output to a specific directory:

python3 segmentation/python/main.py -i assets/images/tabby_tiger_cat.jpg -o /tmp -b 5 -c __background__

To run it on a folder worth of images with batch size 5 for the background class writing the output to a specific directory:

python3 segmentation/python/main.py -i assets/images/ -o /tmp -b 5 -c __background__

To run on a single image with custom resized input given to the model for the dog class with batch size of 1:

python3 segmentation/python/main.py -i assets/images/Weimaraner.jpg -o /tmp -b 1 -c dog -th 224 -tw 224

To run on a single image with custom resized input given to the model for the dog class with batch size of 1 using the PyTorch backend instead of TensorRT:

python3 segmentation/python/main.py -i assets/images/Weimaraner.jpg -o /tmp -b 1 -c dog -th 224 -tw 224 -bk pytorch

To run on a single video with for the background class with batch size of 5:

python segmentation/python/main.py -i assets/videos/pexels-ilimdar-avgezer-7081456.mp4 -b 5 -c __background__

Understanding the Output

This sample takes as input the one or more images or one video and generates the semantic segmentation mask overlay on the inputs corresponding to a class of your choice and saves it in a directory. Since this sample works on batches, sometimes the batch size and the number of images read may not be a perfect multiple. In such cases, the last batch may have a smaller batch size. If the batch size to anything above 1 for one image input case, the same image is fed in the entire batch and identical image masks are generated and saved for all of them.

user@machine:~/cvcuda/samples$ python3 segmentation/python/main.py -b 5 -c __background__ -o /tmp -i assets/images/
[perf_utils:85] 2023-07-27 23:17:49 WARNING perf_utils is used without benchmark.py. Benchmarking mode is turned off.
[perf_utils:89] 2023-07-27 23:17:49 INFO   Using CV-CUDA version: 0.6.0-beta
[pipelines:35] 2023-07-27 23:17:50 INFO   Using CVCUDA as preprocessor.
[torch_utils:60] 2023-07-27 23:17:50 INFO   Found a total of 3 JPEG images.
[torch_utils:77] 2023-07-27 23:17:50 INFO   Using torchnvjpeg as decoder.
[torch_utils:151] 2023-07-27 23:17:50 INFO   Using PyTorch/PIL as encoder.
[pipelines:124] 2023-07-27 23:17:50 INFO   Using CVCUDA as post-processor.
[model_inference:242] 2023-07-27 23:17:50 INFO   Using TensorRT as the inference engine.
[segmentation:183] 2023-07-27 23:17:50 INFO   Processing batch 0
[torch_utils:165] 2023-07-27 23:17:50 INFO   Saving the overlay result to: /tmp/out_peoplenet.jpg
[torch_utils:165] 2023-07-27 23:17:50 INFO   Saving the overlay result to: /tmp/out_tabby_tiger_cat.jpg
[torch_utils:165] 2023-07-27 23:17:50 INFO   Saving the overlay result to: /tmp/out_Weimaraner.jpg