> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.stereolabs.com/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.stereolabs.com/_mcp/server.

# How to Use PyTorch with ZED

## Introduction

The ZED SDK can be interfaced with a PyTorch project to add 3D localization of objects detected with a custom neural network. In this tutorial, we will combine Mask R-CNN with the ZED SDK to detect, segment, classify and locate objects in 3D using a ZED stereo camera and PyTorch.

![image](/_fern-img/5063ae45c93d608e6cd65e358474dc209dc4b316239983161454dbf641625f7b.webp)

## Installation

The [Mask R-CNN 3D](https://github.com/stereolabs/zed-pytorch/) project depends on the following libraries:

* ZED SDK and Python API
* PyTorch (with cuDNN)
* OpenCV
* CUDA
* Python 3
* Apex

> **Warning**
>
> This sample has not been updated since 2022 and targets an older software stack: CUDA 10.0, PyTorch 1.4 and the `maskrcnn-benchmark` project, which is no longer maintained. Current ZED SDK versions require a more recent CUDA version, so these instructions may need to be adapted.

### Get the sample

Clone the sample project from [GitHub](https://github.com/stereolabs/zed-pytorch/). The `requirements.txt` file used below is located in this directory.

```bash
git clone https://github.com/stereolabs/zed-pytorch.git
cd zed-pytorch
```

### ZED SDK

Install the [ZED SDK](https://www.stereolabs.com/developers/release/) and [Python API](/docs/development/api-languages/python).

### PyTorch Installation

#### Using Conda (recommended)

A dedicated environment can be created to set up PyTorch. Keep your environment activated while installing the following packages.

```bash
conda create --name pytorch1 -y
conda activate pytorch1
```

When installing PyTorch, the selected CUDA version must match the one used by the ZED SDK. Here, we use CUDA version **10.0**.

```bash
conda install pytorch torchvision cudatoolkit=10.0 -c pytorch
conda install -c conda-forge --yes --file requirements.txt
```

> **Note**
>
> Do not forget to install [Python API](/docs/development/api-languages/python) inside your current environment.

#### Using Pip

```bash
pip3 install torch torchvision
pip3 install -r requirements.txt
```

For more information, please refer to the [PyTorch](https://pytorch.org/get-started/locally/) setup page.

### Apex Installation

We make use of NVIDIA's [Apex](https://github.com/NVIDIA/apex) API. To install it, run the following:

```bash
git clone https://github.com/NVIDIA/apex
cd apex
python3 setup.py install
```

### Mask R-CNN Installation

Setup Mask R-CNN. If you're using a conda environment, make sure it is still active before running the following commands.

```bash
git clone https://github.com/facebookresearch/maskrcnn-benchmark.git
cd maskrcnn-benchmark
python3 setup.py install
```

## Running Mask R-CNN 3D

The next commands are launched from the `zed-pytorch` sample directory.

Run the code with Python 3. You should be detecting objects captured by your ZED camera using the Mask R-CNN ResNet 50 model and localizing them in 3D.

```bash
python zed_object_detection.py --config-file configs/caffe2/e2e_mask_rcnn_R_50_C4_1x_caffe2.yaml --min-image-size 256
```

## Testing Other Models

Pre-trained models can be found in [MODEL\_ZOO.md](https://github.com/stereolabs/zed-pytorch/blob/master/other/MODEL_ZOO.md). Selected models are downloaded automatically. Here we test Mask R-CNN with ResNet 101.

```bash
python zed_object_detection.py --config-file configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.yaml --min-image-size 300
```

Now let's test 3D key points extraction:

```bash
python zed_object_detection.py --config-file configs/caffe2/e2e_keypoint_rcnn_R_50_FPN_1x_caffe2.yaml --min-image-size 300
```

## Other Options

You can launch object segmentation on recorded videos in SVO format using the following command:

```bash
python zed_object_detection.py --svo-filename path/to/svo_file.svo
```

Best accuracy can be obtained using `--min-image-size 800` (with reduced FPS).

```bash
python zed_object_detection.py --min-image-size 800
```

To display heatmaps, use `--show-mask-heatmaps`.

```bash
python zed_object_detection.py --min-image-size 300 --show-mask-heatmaps
```

Finally, to run the model on the CPU, use `MODEL.DEVICE cpu`.

```bash
python zed_object_detection.py --min-image-size 300 MODEL.DEVICE cpu
```