Unofficial Implementation of CondInst for TensorRT

This repository is an unofficial implementation of CondInst: Conditional Convolutions for Instance Segmentation.

This implementation utilizes heatmap-based object detection. Please see CornerNet: Detecting Objects as Paired Keypoints and CenterNet: Objects as Points for details.

Benchmark

Backbone	Dataset	Input image size	Output mask size	mAP	FP16 inference time Jetson Nano 4GB	FP16 inference time Jetson Xavier NX	Download link
ResNet50-FPN	COCO2017	640x480	160x120	20.1	406 ms	57.8 ms	weight, onnx

Setup

• Build docker image.

  $ cd docker
  $ docker build -t condinst_tensorrt -f Dockerfile .

• To build docker image on Jetson, run the following command. This image only works with JetPack<=4.6.

  $ cd docker
  $ docker build -t condinst_tensorrt -f Dockerfile.jetson .

Usage

• Run docker container.

  $ cd docker
  $ ./run.sh <path/to/this/repository> <path/to/dataset/directory>

• To run docker container on Jetson, run the following command instead.

  $ cd docker
  $ ./run_jetson.sh <path/to/this/repository> <path/to/dataset/directory>

Train

• Train on COCO dataset.

  $ python3 main.py \
      --input_width 640 \
      --input_height 480 \
      --num_classes 80 \
      --topk 40 \
      --mixed_precision True \
      train \
      --train_dir <path/to/train/image/directory> \
      --train_ann <path/to/train/annotation/file> \
      --val_dir <path/to/validation/image/directory> \
      --val_ann <path/to/validation/annotation/file> \
      --pretrained_model <path/to/pretrained/model/if/available> \
      --batch_size 8 \
      --accumulate_grad_batches 16 \
      --num_workers 4 \
      --max_epochs 10 \
      --learning_rate 0.01 \
      --save_model <path/to/model.pt>

  • You might need to increase shared memory size for docker container in order to increase num_workers or input_width or input_height.
  • If you experience NaN losses, try lowering learning_rate or disabling mixed_precision

• Visualize loss on Tensorboard

  $ tensorboard --logdir=runs

Evaluation

• Evaluate model. This feature is experimental and may not produce accurate results.

  $ python3 main.py \
      --input_width 640 \
      --input_height 480 \
      --num_classes 80 \
      --topk 40 \
      --mixed_precision True \
      eval \
      --val_dir <path/to/validation/image/directory> \
      --val_ann <path/to/validation/annotation/file> \
      --batch_size 8 \
      --num_workers 4 \
      --load_model <path/to/model.pt>

Test

• Perform inference and visualize results.

  python3 main.py \
      --input_width 640 \
      --input_height 480 \
      --num_classes 80 \
      --topk 40 \ # Max number of detection
      --mixed_precision True \
      test \
      --test_image_dir <path/to/image/directory> \
      --test_output_dir <path/to/output/directory> \
      --load_model <path/to/model.pt>

Export

• Export ONNX model.

  $ python3 main.py \
      --input_width 640 \
      --input_height 480 \
      --num_classes 80 \
      --topk 40 \ # Max number of detection
      --mixed_precision True \
      export \
      --load_model <path/to/model.pt> \
      --export_onnx <path/to/model.onnx>

• Build TensorRT engine on Jetson device or PC with TensorRT installed.

  $ export PATH=/usr/src/tensorrt/bin/:$PATH # Add path to TensorRT binary
  $ trtexec --fp16 --onnx=<path/to/model.onnx> --dumpProfile --saveEngine=<path/to/tensorrt.engine>

• Run TensorRT demo.

  $ python3 demo_tensorrt.py \
      --test_image_dir <path/to/image/directory> \
      --test_output_dir <path/to/output/directory> \
      --load_engine <path/to/tensorrt.engine>

RGB-D images

See Instructions for RGB-D images for details.