Skip to content

stereolabs/zed-pytorch

Repository files navigation

3D Mask R-CNN using the ZED and Pytorch

The ZED SDK can be interfaced with Pytorch for adding 3D localization of custom objects detected with MaskRCNN. In this Python 3 sample, we will show you how to detect, segmente, classify and locate objects in 3D space using the ZED stereo camera and Pytorch.

Installation

Dependencies

The 3D Object Detection project depends on the following libraries:

  • Python 3
  • CUDA
  • ZED SDK (Python API)
  • Pytorch
  • OpenCV
  • Apex

Getting Started

ZED SDK Installation

Install the ZED SDK and the ZED Python API.

Pytorch Installation

Using Conda (recommended)

The CUDA version must match the one used for the ZED SDK, in that case CUDA 10.0. A dedicated environment can be created to setup Pytorch, but don't forget to activate it, especially when installing MaskRCNN.

conda create --name pytorch1 -y
conda activate pytorch1
conda install pytorch torchvision cudatoolkit=10.0 -c pytorch
conda install --yes --file requirements.txt

Note: Do not forget to install Python API inside your current environment.

Using Pip

pip3 install torch==1.4.0+cu100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html
pip3 install -r requirements.txt

For more information please refer to this page https://pytorch.org/get-started/locally/.

Apex Installation

We make use of NVIDIA's Apex API. To install it, run the following:

$ 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.

$ git clone https://github.com/facebookresearch/maskrcnn-benchmark.git
$ cd maskrcnn-benchmark
$ python setup.py install

Running the sample

Download the sample code from this repository. We provide a simple webcam demo that illustrates how you can use maskrcnn for inference :

python zed_object_detection.py --config-file configs/caffe2/e2e_mask_rcnn_R_50_C4_1x_caffe2.yaml --min-image-size 256
# A SVO file can be loaded
python zed_object_detection.py --svo-filename path/to/svo_file.svo

# for best results, use min-image-size 800
python zed_object_detection.py --min-image-size 800

# or change the model that you want to use
python zed_object_detection.py --config-file configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.yaml --min-image-size 300

# in order to see the probability heatmaps, pass --show-mask-heatmaps
python zed_object_detection.py --min-image-size 300 --show-mask-heatmaps

# for the keypoint demo
python zed_object_detection.py --config-file configs/caffe2/e2e_keypoint_rcnn_R_50_FPN_1x_caffe2.yaml --min-image-size 300

# can also run it on the CPU
python zed_object_detection.py --min-image-size 300 MODEL.DEVICE cpu

The model are downloaded automatically

Model Zoo and Baselines

Pre-trained models can be found in other/MODEL_ZOO.md