nvidia retinanet container

+# NVIDIA ODTK change log
+
+## Version 0.2.6 -- 2021-04-04
+
+### Added
+* `--no-apex` option to `odtk train` and `odtk infer`.
+  * This parameter allows you to switch to Pytorch native AMP and DistributedDataParallel.
+* Adding validation stats to TensorBoard.
+
+### Changed
+* Pytorch Docker container 20.11 from 20.06
+* Added training and inference support for PyTorch native AMP, and torch.nn.parallel.DistributedDataParallel (use `--no-apex`).
+* Switched the Pytorch Model and Data Memory Format to Channels Last. (see [Memory Format Tutorial](https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html))
+* Bug fixes:
+  * Workaround for `'No detections!'` during vlidation added. (see [#52663](https://github.com/pytorch/pytorch/issues/52663))
+  * Freeze unused parameters from torchvision models from autograd gradient calculations.
+  * Make tensorboard writer exclusive to the master process to prevent race conditions.
+* Renamed instances of `retinanet` to `odtk` (folder, C++ namepsaces, etc.)
+
+
+## Version 0.2.5 -- 2020-06-27
+
+### Added
+* `--dynamic-batch-opts` option to `odtk export`.
+  * This parameter allows you to provide TensorRT Optimiation Profile batch sizes for engine export (min, opt, max).
+
+### Changed
+* Updated TensorRT plugins to allow for dynamic batch sizes (see https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#work_dynamic_shapes and https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_plugin_v2_dynamic_ext.html).
+
+
+## Version 0.2.4 -- 2020-04-20
+
+### Added
+* `--anchor-ious` option to `odtk train`.
+  * This parameter allows you to adjust the background and foreground anchor IoU threshold. The default values are `[0.4, 0.5].` 
+  * Example `--anchor-ious 0.3 0.5`. This would mean that any anchor with an IoU of less than 0.3 is assigned to background, 
+  and that any anchor with an IoU of greater than 0.5 is assigned to the foreground object, which is atmost one.
+
+## Version 0.2.3 -- 2020-04-14
+
+### Added
+* `MobileNetV2FPN` backbone
+
+## Version 0.2.2 -- 2020-04-01
+
+### Added
+* Rotated bounding box detections models can now be exported to ONNX and TensorRT using `odtk export model.pth model.plan --rotated-bbox`
+* The `--rotated-bbox` flag is automatically applied when running `odtk infer` or `odtk export` _on a model trained with ODTK version 0.2.2 or later_. 
+
+### Changed
+
+* Improvements to the rotated IoU calculations.
+
+### Limitations
+
+* The C++ API cannot currently infer rotated bounding box models.
+
+## Version 0.2.1 -- 2020-03-18
+
+### Added
+* The DALI dataloader (flag `--with-dali`) now supports image augmentation using:
+   * `--augment-brightness` : Randomly adjusts brightness of image
+   * `--augment-contrast` : Randomly adjusts contrast of image
+   * `--augment-hue` : Randomly adjusts hue of image
+   * `--augment-saturation` : Randomly adjusts saturation of image
+
+### Changed
+* The code in `box.py` for generating anchors has been improved.
+
+## Version 0.2.0 -- 2020-03-13
+
+Version 0.2.0 introduces rotated detections.
+
+### Added
+* `train arguments`:
+  * `--rotated-bbox`: Trains a model is predict rotated bounding boxes `[x, y, w, h, theta]` instead of axis aligned boxes `[x, y, w, h]`.
+* `infer arguments`:
+  * `--rotated-bbox`: Infer a rotated model.
+
+### Changed
+The project has reverted to the name **Object Detection Toolkit** (ODTK), to better reflect the multi-network nature of the repo.
+* `retinanet` has been replaced with `odtk`. All subcommands remain the same. 
+
+### Limitations
+* Models trained using the `--rotated-bbox` flag cannot be exported to ONNX or a TensorRT Engine.
+* PyTorch raises two warnings which can be ignored:
+
+Warning 1: NCCL watchdog
+```
+[E ProcessGroupNCCL.cpp:284] NCCL watchdog thread terminated
+```
+
+Warning 2: Save state warning
+```
+/opt/conda/lib/python3.6/site-packages/torch/optim/lr_scheduler.py:201: UserWarning: Please also save or load the state of the optimzer when saving or loading the scheduler.
+  warnings.warn(SAVE_STATE_WARNING, UserWarning)
+```
+
+## Version 0.1.1 -- 2020-03-06
+
+### Added
+ * `train` arguments
+   * `--augment-rotate`: Randomly rotates the training images by 0°, 90°, 180° or 270°.
+   * `--augment-brightness` : Randomly adjusts brightness of image
+   * `--augment-contrast` : Randomly adjusts contrast of image
+   * `--augment-hue` : Randomly adjusts hue of image
+   * `--augment-saturation` : Randomly adjusts saturation of image
+   * `--regularization-l2` : Sets the L2 regularization of the optimizer.

+Reporting problems, asking questions
+------------------------------------
+
+
+We appreciate feedback, questions or bug reports. When you need help with the code, try to follow the process outlined in the Stack Overflow (https://stackoverflow.com/help/mcve) document. 
+
+At a minimum, your issues should describe the following:
+
+* What command you ran 
+* The hardware and container that you are using
+* The version of ODTK you are using
+* What was the result you observed
+* What was the result you expected

+FROM nvcr.io/nvidia/pytorch:20.11-py3
+
+COPY . odtk/
+RUN pip install --no-cache-dir -e odtk/

+FROM nvcr.io/nvidia/pytorch:20.03-py3
+COPY . /workspace/retinanet-examples/
+RUN apt-get update && apt-get install -y libssl1.0.0 libgstreamer1.0-0 gstreamer1.0-tools gstreamer1.0-plugins-good gstreamer1.0-plugins-bad gstreamer1.0-plugins-ugly gstreamer1.0-libav libgstrtspserver-1.0-0 libjansson4 ffmpeg libjson-glib-1.0 libgles2-mesa
+RUN git clone https://github.com/edenhill/librdkafka.git /librdkafka && \
+    cd /librdkafka && ./configure && make -j && make -j install && \
+    mkdir -p /opt/nvidia/deepstream/deepstream-4.0/lib && \
+    cp /usr/local/lib/librdkafka* /opt/nvidia/deepstream/deepstream-4.0/lib && \
+    rm -rf /librdkafka
+WORKDIR /workspace/retinanet-examples/extras/deepstream/DeepStream_Release/deepstream_sdk_v4.0.2_x86_64
+RUN tar -xvf binaries.tbz2 -C / && \
+    ./install.sh
+# config files + sample apps
+RUN chmod u+x ./sources/tools/nvds_logger/setup_nvds_logger.sh
+
+WORKDIR /usr/lib/x86_64-linux-gnu
+RUN ln -sf libnvcuvid.so.1 libnvcuvid.so
+
+WORKDIR /workspace/retinanet-examples
+RUN pip install --no-cache-dir -e .
+RUN mkdir extras/deepstream/deepstream-sample/build && \
+    cd extras/deepstream/deepstream-sample/build && \
+    cmake -DDeepStream_DIR=/workspace/retinanet-examples/extras/deepstream/DeepStream_Release/deepstream_sdk_v4.0.2_x86_64 .. && make -j
+WORKDIR /workspace/retinanet-examples/extras/deepstream

+# Inference
+
+We provide two ways inferring using `odtk`:
+* PyTorch inference using a trained model (FP32 or FP16 precision)
+* Export trained pytorch model to TensorRT for optimized inference (FP32, FP16 or INT8 precision)
+
+`odtk infer` will run distributed inference across all available GPUs. When using PyTorch, the default behavior is to run inference with mixed precision. The precision used when running inference with a TensorRT engine will correspond to the precision chosen when the model was exported to TensorRT (see [TensorRT section](#exporting-trained-pytorch-model-to-tensorrt) below). 
+
+**NOTE**: Availability of HW support for fast FP16 and INT8 precision like [NVIDIA Tensor Cores](https://www.nvidia.com/en-us/data-center/tensorcore/) depends on your GPU architecture: Volta or newer GPUs support both FP16 and INT8, and Pascal GPUs can support either FP16 or INT8. 
+
+## PyTorch Inference
+
+Evaluate trained PyTorch detection model on COCO 2017 (mixed precision):
+
+```bash
+odtk infer model.pth --images=/data/coco/val2017 --annotations=instances_val2017.json --batch 8
+```
+**NOTE**: `--batch N` specifies *global* batch size to be used for inference. The batch size per GPU will be `N // num_gpus`.
+
+Use full precision (FP32) during evaluation:
+
+```bash
+odtk infer model.pth --images=/data/coco/val2017 --annotations=instances_val2017.json --full-precision
+```
+
+Evaluate PyTorch detection model with a small input image size:
+
+```bash
+odtk infer model.pth --images=/data/coco/val2017 --annotations=instances_val2017.json  --resize 400 --max-size 640
+```
+Here, the shorter side of the input images will be resized to `resize` as long as the longer side doesn't get larger than `max-size`, otherwise the longer side of the input image will be resized to `max-size`.
+
+**NOTE**: To get best accuracy, training the model at the preferred export size is encouraged.
+
+Run inference using your own dataset:
+
+```bash
+odtk infer model.pth --images=/data/your_images --output=detections.json
+```
+
+## Exporting trained PyTorch model to TensorRT
+
+`odtk` provides an simple workflow to optimize a trained PyTorch model for inference deployment using TensorRT. The PyTorch model is exported to [ONNX](https://github.com/onnx/onnx), and then the ONNX model is consumed and optimized by TensorRT.
+To learn more about TensorRT optimization, refer here: https://developer.nvidia.com/tensorrt
+
+**NOTE**: When a model is optimized with TensorRT, the output is a TensorRT engine (.plan file) that can be used for deployment. This TensorRT engine has several fixed properties that are specified during the export process.
+* Input image size: TensorRT engines only support a fixed input size.
+* Precision: TensorRT supports FP32, FP16, or INT8 precision.
+* Target GPU: TensorRT optimizations are tied to the type of GPU on the system where optimization is performed. They are not transferable across different types of GPUs. Put another way, if you aim to deploy your TensorRT engine on a Tesla T4 GPU, you must run the optimization on a system with a T4 GPU. 
+
+The workflow for exporting a trained PyTorch detection model to TensorRT is as simple as:
+
+```bash
+odtk export model.pth model_fp16.plan --size 1280
+```
+This will create a TensorRT engine optimized for batch size 1, using an input size of 1280x1280. By default, the engine will be created to run in FP16 precision.
+
+Export your model to use full precision using a non-square input size:
+```bash
+odtk export model.pth model_fp32.plan --full-precision --size 800 1280
+```
+
+In order to use INT8 precision with TensorRT, you need to provide calibration images (images that are representative of what will be seen at runtime) that will be used to rescale the network.
+```bash
+odtk export model.pth model_int8.plan --int8 --calibration-images /data/val/ --calibration-batches 2 --calibration-table model_calibration_table
+```
+
+This will randomly select 16 images from `/data/val/` to calibrate the network for INT8 precision. The results from calibration will be saved to `model_calibration_table` that can be used to create subsequent INT8 engines for this model without needed to recalibrate. 
+
+**NOTE:** Number of images in `/data/val/` must be greater than or equal to the kOPT(middle) optimization profile from `--dynamic-batch-opts`. Here, the default kOPT is 8.
+
+Build an INT8 engine for a previously calibrated model:
+```bash
+odtk export model.pth model_int8.plan --int8 --calibration-table model_calibration_table
+```
+
+## Deployment with TensorRT on NVIDIA Jetson AGX Xavier
+
+We provide a path for deploying trained models with TensorRT onto embedded platforms like [NVIDIA Jetson AGX Xavier](https://developer.nvidia.com/embedded/buy/jetson-agx-xavier-devkit), where PyTorch is not readily available. 
+
+You will need to export your trained PyTorch model to ONNX representation on your host system, and copy the resulting ONNX model to your Jetson AGX Xavier:
+```bash
+odtk export model.pth model.onnx --size 800 1280
+```
+
+Refer to additional documentation on using the example cppapi code to build the TensorRT engine and run inference here: [cppapi example code](extras/cppapi/README.md)
+
+## Rotated detections
+
+*Rotated ODTK* allows users to train and infer rotated bounding boxes in imagery. 
+
+### Inference
+
+An example command:
+```
+odtk infer model.pth --images /data/val --annotations /data/val_rotated.json --output /data/detections.json \ 
+    --resize 768 --rotated-bbox
+```
+
+### Export
+
+Rotated bounding box models can be exported to create TensorRT engines by using the axis aligned command with the addition of `--rotated-bbox`.
\ No newline at end of file

+Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+ * Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+ * Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in the
+   documentation and/or other materials provided with the distribution.
+ * Neither the name of NVIDIA CORPORATION nor the names of its
+   contributors may be used to endorse or promote products derived
+   from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

+# NVIDIA Object Detection Toolkit (ODTK)
+
+**Fast** and **accurate** single stage object detection with end-to-end GPU optimization.
+
+## Description
+
+ODTK is a single shot object detector with various backbones and detection heads. This allows performance/accuracy trade-offs.
+
+It is optimized for end-to-end GPU processing using:
+* The [PyTorch](https://pytorch.org) deep learning framework with [ONNX](https://onnx.ai) support
+* NVIDIA [Apex](https://github.com/NVIDIA/apex) for mixed precision and distributed training
+* NVIDIA [DALI](https://github.com/NVIDIA/DALI) for optimized data pre-processing
+* NVIDIA [TensorRT](https://developer.nvidia.com/tensorrt) for high-performance inference
+* NVIDIA [DeepStream](https://developer.nvidia.com/deepstream-sdk) for optimized real-time video streams support
+
+## Rotated bounding box detections
+
+This repo now supports rotated bounding box detections. See [rotated detections training](TRAINING.md#rotated-detections) and [rotated detections inference](INFERENCE.md#rotated-detections) documents for more information on how to use the `--rotated-bbox` command. 
+
+Bounding box annotations are described by `[x, y, w, h, theta]`. 
+
+## Performance
+
+The detection pipeline allows the user to select a specific backbone depending on the latency-accuracy trade-off preferred.
+
+ODTK **RetinaNet** model accuracy and inference latency & FPS (frames per seconds) for [COCO 2017](http://cocodataset.org/#detection-2017) (train/val) after full training schedule. Inference results include bounding boxes post-processing for a batch size of 1. Inference measured at `--resize 800` using `--with-dali` on a FP16 TensorRT engine.
+
+Backbone |  mAP @[IoU=0.50:0.95] | Training Time on [DGX1v](https://www.nvidia.com/en-us/data-center/dgx-1/) | Inference latency FP16 on [V100](https://www.nvidia.com/en-us/data-center/tesla-v100/) | Inference latency INT8 on [T4](https://www.nvidia.com/en-us/data-center/tesla-t4/) | Inference latency FP16 on [A100](https://www.nvidia.com/en-us/data-center/a100/) | Inference latency INT8 on [A100](https://www.nvidia.com/en-us/data-center/a100/)
+--- | :---: | :---: | :---: | :---: | :---: | :---:
+[ResNet18FPN](https://github.com/NVIDIA/retinanet-examples/releases/download/19.04/retinanet_rn18fpn.zip) | 0.318 | 5 hrs  | 14 ms;</br>71 FPS | 18 ms;</br>56 FPS | 9 ms;</br>110 FPS | 7 ms;</br>141 FPS
+[MobileNetV2FPN](https://github.com/NVIDIA/retinanet-examples/releases/download/v0.2.3/retinanet_mobilenetv2fpn.pth) | 0.333 | | 14 ms;</br>74 FPS | 18 ms;</br>56 FPS | 9 ms;</br>114 FPS | 7 ms;</br>138 FPS
+[ResNet34FPN](https://github.com/NVIDIA/retinanet-examples/releases/download/19.04/retinanet_rn34fpn.zip) | 0.343 | 6 hrs  | 16 ms;</br>64 FPS | 20 ms;</br>50 FPS | 10 ms;</br>103 FPS | 7 ms;</br>142 FPS
+[ResNet50FPN](https://github.com/NVIDIA/retinanet-examples/releases/download/19.04/retinanet_rn50fpn.zip) | 0.358 | 7 hrs  | 18 ms;</br>56 FPS | 22 ms;</br>45 FPS | 11 ms;</br>93 FPS | 8 ms;</br>129 FPS
+[ResNet101FPN](https://github.com/NVIDIA/retinanet-examples/releases/download/19.04/retinanet_rn101fpn.zip) | 0.376 | 10 hrs | 22 ms;</br>46 FPS | 27 ms;</br>37 FPS | 13 ms;</br>78 FPS | 9 ms;</br>117 FPS
+[ResNet152FPN](https://github.com/NVIDIA/retinanet-examples/releases/download/19.04/retinanet_rn152fpn.zip) | 0.393 | 12 hrs | 26 ms;</br>38 FPS | 33 ms;</br>31 FPS | 15 ms;</br>66 FPS | 10 ms;</br>103 FPS
+
+## Installation
+
+For best performance, use the latest [PyTorch NGC docker container](https://ngc.nvidia.com/catalog/containers/nvidia:pytorch). Clone this repository, build and run your own image:
+
+```bash
+git clone https://github.com/nvidia/retinanet-examples
+docker build -t odtk:latest retinanet-examples/
+docker run --gpus all --rm --ipc=host -it odtk:latest
+```
+
+## Usage
+
+Training, inference, evaluation and model export can be done through the `odtk` utility. 
+For more details, including a list of parameters, please refer to the [TRAINING](TRAINING.md) and [INFERENCE](INFERENCE.md) documentation.
+
+### Training
+
+Train a detection model on [COCO 2017](http://cocodataset.org/#download) from pre-trained backbone:
+```bash
+odtk train retinanet_rn50fpn.pth --backbone ResNet50FPN \
+    --images /coco/images/train2017/ --annotations /coco/annotations/instances_train2017.json \
+    --val-images /coco/images/val2017/ --val-annotations /coco/annotations/instances_val2017.json
+```
+
+### Fine Tuning
+
+Fine-tune a pre-trained model on your dataset. In the example below we use [Pascal VOC](http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html) with [JSON annotations](https://storage.googleapis.com/coco-dataset/external/PASCAL_VOC.zip):
+```bash
+odtk train model_mydataset.pth --backbone ResNet50FPN \
+    --fine-tune retinanet_rn50fpn.pth \
+    --classes 20 --iters 10000 --val-iters 1000 --lr 0.0005 \
+    --resize 512 --jitter 480 640 --images /voc/JPEGImages/ \
+    --annotations /voc/pascal_train2012.json --val-annotations /voc/pascal_val2012.json
+```
+
+Note: the shorter side of the input images will be resized to `resize` as long as the longer side doesn't get larger than `max-size`. During training, the images will be randomly randomly resized to a new size within the `jitter` range.
+
+### Inference
+
+Evaluate your detection model on [COCO 2017](http://cocodataset.org/#download):
+```bash
+odtk infer retinanet_rn50fpn.pth --images /coco/images/val2017/ --annotations /coco/annotations/instances_val2017.json
+```
+
+Run inference on [your dataset](#datasets):
+```bash
+odtk infer retinanet_rn50fpn.pth --images /dataset/val --output detections.json
+```
+
+### Optimized Inference with TensorRT
+
+For faster inference, export the detection model to an optimized FP16 TensorRT engine:
+```bash
+odtk export model.pth engine.plan
+```
+
+Evaluate the model with TensorRT backend on [COCO 2017](http://cocodataset.org/#download):
+```bash
+odtk infer engine.plan --images /coco/images/val2017/ --annotations /coco/annotations/instances_val2017.json
+```
+
+### INT8 Inference with TensorRT
+
+For even faster inference, do INT8 calibration to create an optimized INT8 TensorRT engine:
+```bash
+odtk export model.pth engine.plan --int8 --calibration-images /coco/images/val2017/
+```
+This will create an INT8CalibrationTable file that can be used to create INT8 TensorRT engines for the same model later on without needing to do calibration.
+
+Or create an optimized INT8 TensorRT engine using a cached calibration table:
+```bash
+odtk export model.pth engine.plan --int8 --calibration-table /path/to/INT8CalibrationTable
+```
+
+## Datasets
+
+RetinaNet supports annotations in the [COCO JSON format](http://cocodataset.org/#format-data).
+When converting the annotations from your own dataset into JSON, the following entries are required:
+```
+{
+    "images": [{
+        "id" : int,
+        "file_name" : str
+    }],
+    "annotations": [{
+        "id" : int,
+        "image_id" : int, 
+        "category_id" : int,
+        "bbox" : [x, y, w, h]   # all floats
+        "area": float           # w * h. Required for validation scores
+        "iscrowd": 0            # Required for validation scores
+    }],
+    "categories": [{
+        "id" : int
+    ]}
+}
+```
+
+If using the `--rotated-bbox` flag for rotated detections, add an additional float `theta` to the annotations. To get validation scores you also need to fill the `segmentation` section.
+```
+        "bbox" : [x, y, w, h, theta]    # all floats, where theta is measured in radians anti-clockwise from the x-axis.
+        "segmentation" : [[x1, y1, x2, y2, x3, y3, x4, y4]]
+                                        # Required for validation scores.
+```
+
+## Disclaimer
+
+This is a research project, not an official NVIDIA product.
+
+## Jetpack compatibility
+
+This branch uses TensorRT 7. If you are training and inferring models using PyTorch, or are creating TensorRT engines on Tesla GPUs (eg V100, T4), then you should use this branch.
+
+If you wish to deploy your model to a Jetson device (eg - Jetson AGX Xavier) running Jetpack version 4.3, then you should use the `19.10` branch of this repo.
+
+## References
+
+- [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002).
+  Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, Piotr Dollár.
+  ICCV, 2017.
+- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://arxiv.org/abs/1706.02677).
+  Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, Kaiming He.
+  June 2017.
+- [Feature Pyramid Networks for Object Detection](https://arxiv.org/abs/1612.03144).
+  Tsung-Yi Lin, Piotr Dollár, Ross Girshick, Kaiming He, Bharath Hariharan, Serge Belongie.
+  CVPR, 2017.
+- [Deep Residual Learning for Image Recognition](http://arxiv.org/abs/1512.03385).
+  Kaiming He, Xiangyu Zhang, Shaoqing Renm Jian Sun.
+  CVPR, 2016.

+# Training
+
+There are two main ways to train a model with `odtk`:
+* Fine-tuning the detection model using a model already trained on a large dataset (like MS-COCO)
+* Fully training the detection model from random initialization using a pre-trained backbone (usually ImageNet)
+
+## Fine-tuning
+
+Fine-tuning an existing model trained on COCO allows you to use transfer learning to get a accurate model for your own dataset with minimal training.
+When fine-tuning, we re-initialize the last layer of the classification head so the network will re-learn how to map features to classes scores regardless of the number of classes in your own dataset.
+
+You can fine-tune a pre-trained model on your dataset. In the example below we take a model trained on COCO, and then fine-tune using [Pascal VOC](http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html) with [JSON annotations](https://storage.googleapis.com/coco-dataset/external/PASCAL_VOC.zip):
+```bash
+odtk train model_mydataset.pth \
+    --fine-tune retinanet_rn50fpn.pth \
+    --classes 20 --iters 10000 --val-iters 1000 --lr 0.0005 \
+    --resize 512 --jitter 480 640 --images /voc/JPEGImages/ \
+    --annotations /voc/pascal_train2012.json --val-annotations /voc/pascal_val2012.json
+```
+
+Even though the COCO model was trained on 80 classes, we can easily use tranfer learning to fine-tune it on the Pascal VOC model representing only 20 classes.
+
+The shorter side of the input images will be resized to `resize` as long as the longer side doesn't get larger than `max-size`.
+During training the images will be randomly resized to a new size within the `jitter` range.
+
+We usually want to fine-tune the model with a lower learning rate `lr` than during full training and for less iterations `iters`.
+
+## Full Training
+
+If you do not have a pre-trained model, if your dataset is substantially large, or if you have written your own backbone, then you should fully train the detection model.
+
+Full training usually starts from a pre-trained backbone (automatically downloaded with the current backbones we offer) that has been pre-trained on a classification task with a large dataset like [ImageNet](http://www.image-net.org).
+This is especially necessary for backbones using batch normalization as they require large batch sizes during training that cannot be provided when training on the detection task as the input images have to be relatively large.
+
+Train a detection model on [COCO 2017](http://cocodataset.org/#download) from pre-trained backbone:
+```bash
+odtk train retinanet_rn50fpn.pth --backbone ResNet50FPN \
+    --images /coco/images/train2017/ --annotations /coco/annotations/instances_train2017.json \
+    --val-images /coco/images/val2017/ --val-annotations /coco/annotations/instances_val2017.json
+```
+
+## Training arguments
+
+### Positional arguments
+* The only positional argument is the name of the model. This can be a full path, or relative to the current directory.
+```bash
+odtk train model.pth
+```
+
+### Other arguments
+The following arguments are available during training:
+
+* `--annotations` (str): Path to COCO style annotations (required).
+* `--images` (str): Path to a directory of images (required).
+* `--lr` (float): Sets the learning rate. Default: 0.01.
+* `--full-precision`: By default we train using mixed precision. Include this argument to instead train in full precision.
+* `--warmup` (int): The number of initial iterations during which we want to linearly ramp-up the learning rate to avoid early divergence of the loss. Default: 1000
+* `--backbone` (str): Specify one of the supported backbones. Default: `ResNet50FPN`
+* `--classes` (int): The number of classes in your dataset. Default: 80
+* `--batch` (int): The size of each training batch. Default: 2 x number of GPUs.
+* `--max-size` (int): The longest edge of your training image will be resized, so that it is always less than or equal to `max-size`. Default: 1333. 
+* `--jitter` (int int): The shortest edge of your training images will be resized to int1 >= shortest edge >= int2, unless the longest edge exceeds `max-size`, in which case the longest edge will be resized to `max-size` and the shortest length will be sized to keep the aspect ratio constant. Default: 640 1024.
+* `--resize` (int): During validation inference, the shortest edge of your training images will be resized to int, unless the longest edge exceeds `max-size`, in which case the longest edge will be resized to `max-size` and the shortest length will be sized to keep the aspect ratio constant. Default: 800.
+* `--iters` (int): The number of iterations to process. An iteration is the processing (forward and backward pass) of one batch. Number of epochs is (`iters` x `batch`) / `len(data)`. Default: 90000.
+* `--milestones` (int int): The learning rate is multiplied by `--gamma` every time it reaches a milestone. Default: 60000 80000.
+* `--gamma` (float): The learning rate is multiplied by `--gamma` every time it reaches a milestone. Default: 0.1.
+* `--override`: Do not continue training from `model.pth`, instead overwrite it.
+* `--val-annotations` (str): Path to COCO style annotations. If supplied, `pycocotools` will be used to give validation mAP.
+* `--val-images` (str): Path to directory of validation images.
+* `--val-iters` (int): Run inference on the validation set every int iterations.
+* `--fine-tune` (str): Fine tune from a model at path str.
+* `--with-dali`: Load data using DALI.
+* `--augment-rotate`: Randomly rotates the training images by 0°, 90°, 180° or 270°.
+* `--augment-brightness` (float): Randomly adjusts brightness of image. The value sets the standard deviation of a Gaussian distribution. The degree of augmentation is selected from this distribution. Default: 0.002
+* `--augment-contrast` (float): Randomly adjusts contrast of image. The value sets the standard deviation of a Gaussian distribution. The degree of augmentation is selected from this distribution. Default: 0.002
+* `--augment-hue` (float): Randomly adjusts hue of image. The value sets the standard deviation of a Gaussian distribution. The degree of augmentation is selected from this distribution. Default: 0.0002
+* `--augment-saturation` (float): Randomly adjusts saturation of image. The value sets the standard deviation of a Gaussian distribution. The degree of augmentation is selected from this distribution. Default: 0.002
+* `--regularization-l2` (float): Sets the L2 regularization of the optimizer. Default: 0.0001
+
+You can also monitor the loss and learning rate schedule of the training using TensorBoard bu specifying a `logdir` path.
+
+## Rotated detections
+
+*Rotated ODTK* allows users to train and infer rotated bounding boxes in imagery. 
+
+### Dataset
+Annotations need to conform to the COCO standard, with the addition of an angle (radians) in the bounding box (bbox) entry `[xmin, ymin, width, height, **theta**]`. `xmin`, `ymin`, `width` and `height` are in the axis aligned coordinates, ie floats, measured from the top left of the image. `theta` is in radians, measured anti-clockwise from the x-axis. We constrain theta between - \pi/4 and \pi/4.
+
+In order for the validation metrics to calculate, you also need to fill the `segmentation` entry with the coordinates of the corners of your bounding box.
+
+If using the `--rotated-bbox` flag for rotated detections, add an additional float `theta` to the annotations. To get validation scores you also need to fill the `segmentation` section.
+```
+        "bbox" : [x, y, w, h, theta]    # all floats, where theta is measured in radians anti-clockwise from the x-axis.
+        "segmentation" : [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
+                                        # Required for validation scores.
+``` 
+
+### Anchors
+
+As with all single shot detectors, the anchor boxes may need to be adjusted to suit your dataset. You may need to adjust the anchors in `odtk/model.py`
+
+The default anchors are:
+
+```python
+self.ratios = [0.5, 1.0, 2.0]
+self.scales = [4 * 2**(i/3) for i in range(3)]
+self.angles = [-np.pi/6, 0, np.pi/6] 
+```
+
+### Training
+
+We recommend reducing your learning rate, for example using `--lr 0.0005`. 
+
+An example training command for training remote sensing imagery. Note that `--augment-rotate` has been used to randomly rotated the imagery during training.
+```
+odtk train model.pth --images /data/train --annotations /data/train_rotated.json  --backbone ResNet50FPN \ 
+    --lr 0.00005 --fine-tune /data/saved_models/retinanet_rn50fpn.pth \ 
+    --val-images /data/val --val-annotations /data/val_rotated.json --classes 1 \ 
+    --jitter 688 848 --resize 768 \ 
+    --augment-rotate --augment-brightness 0.01 --augment-contrast 0.01 --augment-hue 0.002 \ 
+    --augment-saturation 0.01 --batch 16 --regularization-l2 0.0001  --val-iters 20000 --rotated-bbox
+```
+

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+
+#include <opencv2/opencv.hpp>
+#include <opencv2/opencv.hpp>
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <iterator>
+#include <vector>
+#include <assert.h>
+#include <algorithm>
+#include "NvInfer.h"
+
+using namespace std;
+using namespace cv;
+
+class ImageStream {
+public:
+    ImageStream(int batchSize, Dims inputDims, const vector<string> calibrationImages)
+        : _batchSize(batchSize)
+        , _calibrationImages(calibrationImages)
+        , _currentBatch(0)
+        , _maxBatches(_calibrationImages.size() / _batchSize)
+        , _inputDims(inputDims) {
+        _batch.resize(_batchSize * _inputDims.d[1] * _inputDims.d[2] * _inputDims.d[3]);
+    }
+
+    int getBatchSize() const { return _batchSize;}
+
+    int getMaxBatches() const { return _maxBatches;}
+
+    float* getBatch() { return &_batch[0];}
+
+    Dims getInputDims() { return _inputDims;}
+
+    bool next() {
+        
+        if (_currentBatch == _maxBatches)
+            return false;
+
+        for (int i = 0; i < _batchSize; i++) {
+            auto image = imread(_calibrationImages[_batchSize * _currentBatch + i].c_str(), IMREAD_COLOR);
+            cv::resize(image, image, Size(_inputDims.d[3], _inputDims.d[2]));
+            cv::Mat pixels;
+            image.convertTo(pixels, CV_32FC3, 1.0 / 255, 0);
+
+            vector<float> img;
+
+            if (pixels.isContinuous())
+                img.assign((float*)pixels.datastart, (float*)pixels.dataend);
+            else
+                return false;
+
+            auto hw = _inputDims.d[2] * _inputDims.d[3];
+            auto channels = _inputDims.d[1];
+            auto vol = channels * hw;
+
+            for (int c = 0; c < channels; c++) {
+                for (int j = 0; j < hw; j++) {
+                    _batch[i * vol + c * hw + j] = (img[channels * j + 2 - c] - _mean[c]) / _std[c];
+                }
+            }
+        }
+
+        _currentBatch++;
+        return true;
+    }
+            
+    void reset() {
+        _currentBatch = 0;
+    }
+
+private:
+    int _batchSize;
+    vector<string> _calibrationImages;
+    int _currentBatch;
+    int _maxBatches;
+    Dims _inputDims;
+
+    vector<float> _mean {0.485, 0.456, 0.406};
+    vector<float> _std {0.229, 0.224, 0.225};
+    vector<float> _batch;
+    
+};
+
+class Int8EntropyCalibrator: public IInt8EntropyCalibrator2 {
+public:
+    Int8EntropyCalibrator(ImageStream& stream, const string networkName, const string calibrationCacheName, bool readCache = true)
+        : _stream(stream)
+        , _networkName(networkName)
+        , _calibrationCacheName(calibrationCacheName)
+        , _readCache(readCache) {
+            Dims d = _stream.getInputDims();
+            _inputCount = _stream.getBatchSize() * d.d[1] * d.d[2] * d.d[3];
+            cudaMalloc(&_deviceInput, _inputCount * sizeof(float));
+        }
+
+    int getBatchSize() const override {return _stream.getBatchSize();}
+
+    virtual ~Int8EntropyCalibrator() {cudaFree(_deviceInput);}
+
+    bool getBatch(void* bindings[], const char* names[], int nbBindings) override {
+
+        if (!_stream.next())
+            return false;
+
+        cudaMemcpy(_deviceInput, _stream.getBatch(), _inputCount * sizeof(float), cudaMemcpyHostToDevice);
+        bindings[0] = _deviceInput;
+        return true;
+    }
+
+    const void* readCalibrationCache(size_t& length) { 
+        _calibrationCache.clear();
+        ifstream input(calibrationTableName(), ios::binary);
+        input >> noskipws;
+        if (_readCache && input.good())
+            copy(istream_iterator<char>(input), istream_iterator<char>(), back_inserter(_calibrationCache));
+
+        length = _calibrationCache.size();
+        return length ? &_calibrationCache[0] : nullptr;
+    }
+
+    void writeCalibrationCache(const void* cache, size_t length) {
+        std::ofstream output(calibrationTableName(), std::ios::binary);
+        output.write(reinterpret_cast<const char*>(cache), length);
+    }
+
+private:
+    std::string calibrationTableName() {
+        // Use calibration cache if provided
+        if(_calibrationCacheName.length() > 0)
+            return _calibrationCacheName;
+
+        assert(_networkName.length() > 0);
+        Dims d = _stream.getInputDims();
+        return std::string("Int8CalibrationTable_") + _networkName + to_string(d.d[2]) + "x" + to_string(d.d[3]) + "_" + to_string(_stream.getMaxBatches());
+    }
+
+    ImageStream _stream;
+    const string _networkName;
+    const string _calibrationCacheName;
+    bool _readCache {true};
+    size_t _inputCount;
+    void* _deviceInput {nullptr};
+    vector<char> _calibrationCache;
+
+};

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include "decode.h"
+#include "utils.h"
+
+#include <algorithm>
+#include <cstdint>
+
+#include <thrust/device_ptr.h>
+#include <thrust/sequence.h>
+#include <thrust/execution_policy.h>
+#include <thrust/gather.h>
+#include <thrust/tabulate.h>
+#include <thrust/count.h>
+#include <thrust/find.h>
+#include <cub/device/device_radix_sort.cuh>
+#include <cub/iterator/counting_input_iterator.cuh>
+
+#include <stdio.h>
+
+namespace odtk {
+namespace cuda {
+
+int decode(int batch_size,
+  const void *const *inputs, void *const *outputs,
+  size_t height, size_t width, size_t scale,
+  size_t num_anchors, size_t num_classes,
+  const std::vector<float> &anchors, float score_thresh, int top_n,
+  void *workspace, size_t workspace_size, cudaStream_t stream) {
+
+  int scores_size = num_anchors * num_classes * height * width;
+  
+  if (!workspace || !workspace_size) {
+    // Return required scratch space size cub style
+    workspace_size  = get_size_aligned<float>(anchors.size()); // anchors
+    workspace_size += get_size_aligned<bool>(scores_size);     // flags
+    workspace_size += get_size_aligned<int>(scores_size);      // indices
+    workspace_size += get_size_aligned<int>(scores_size);      // indices_sorted
+    workspace_size += get_size_aligned<float>(scores_size);    // scores
+    workspace_size += get_size_aligned<float>(scores_size);    // scores_sorted
+  
+    size_t temp_size_flag = 0;
+    cub::DeviceSelect::Flagged((void *)nullptr, temp_size_flag,
+      cub::CountingInputIterator<int>(scores_size),
+      (bool *)nullptr, (int *)nullptr, (int *)nullptr, scores_size);
+    size_t temp_size_sort = 0;
+    cub::DeviceRadixSort::SortPairsDescending((void *)nullptr, temp_size_sort,
+      (float *)nullptr, (float *)nullptr, (int *)nullptr, (int *)nullptr, scores_size);
+    workspace_size += std::max(temp_size_flag, temp_size_sort);
+
+    return workspace_size;
+  }
+
+  auto anchors_d = get_next_ptr<float>(anchors.size(), workspace, workspace_size);
+  cudaMemcpyAsync(anchors_d, anchors.data(), anchors.size() * sizeof *anchors_d, cudaMemcpyHostToDevice, stream);
+
+  auto on_stream = thrust::cuda::par.on(stream);
+
+  auto flags = get_next_ptr<bool>(scores_size, workspace, workspace_size);
+  auto indices = get_next_ptr<int>(scores_size, workspace, workspace_size);
+  auto indices_sorted = get_next_ptr<int>(scores_size, workspace, workspace_size);
+  auto scores = get_next_ptr<float>(scores_size, workspace, workspace_size);
+  auto scores_sorted = get_next_ptr<float>(scores_size, workspace, workspace_size);
+
+  
+  for (int batch = 0; batch < batch_size; batch++) {
+  
+    auto in_scores = static_cast<const float *>(inputs[0]) + batch * scores_size;
+    auto in_boxes = static_cast<const float *>(inputs[1]) + batch * (scores_size / num_classes) * 4;
+
+    auto out_scores = static_cast<float *>(outputs[0]) + batch * top_n;
+    auto out_boxes = static_cast<float4 *>(outputs[1]) + batch * top_n;
+    auto out_classes = static_cast<float *>(outputs[2]) + batch * top_n;
+
+    // Discard scores below threshold
+    thrust::transform(on_stream, in_scores, in_scores + scores_size,
+      flags, thrust::placeholders::_1 > score_thresh);
+
+    int *num_selected = reinterpret_cast<int *>(indices_sorted);
+    cub::DeviceSelect::Flagged(workspace, workspace_size,
+      cub::CountingInputIterator<int>(0),
+      flags, indices, num_selected, scores_size, stream);
+    cudaStreamSynchronize(stream);
+    int num_detections = *thrust::device_pointer_cast(num_selected);
+
+    // Only keep top n scores
+    auto indices_filtered = indices;
+    if (num_detections > top_n) {
+      thrust::gather(on_stream, indices, indices + num_detections,
+        in_scores, scores);
+      cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
+        scores, scores_sorted, indices, indices_sorted, num_detections, 0, sizeof(*scores)*8, stream);
+        indices_filtered = indices_sorted;
+        num_detections = top_n;
+    }
+
+    // Gather boxes
+    bool has_anchors = !anchors.empty();
+    thrust::transform(on_stream, indices_filtered, indices_filtered + num_detections,
+      thrust::make_zip_iterator(thrust::make_tuple(out_scores, out_boxes, out_classes)),
+      [=] __device__ (int i) {
+        int x = i % width;
+        int y = (i / width) % height;
+        int a = (i / num_classes / height / width) % num_anchors;
+        int cls = (i / height / width) % num_classes;
+        float4 box = float4{
+          in_boxes[((a * 4 + 0) * height + y) * width + x],
+          in_boxes[((a * 4 + 1) * height + y) * width + x],
+          in_boxes[((a * 4 + 2) * height + y) * width + x],
+          in_boxes[((a * 4 + 3) * height + y) * width + x]
+        };
+
+        if (has_anchors) {
+          // Add anchors offsets to deltas
+          float x = (i % width) * scale;
+          float y = ((i / width)  % height) * scale;
+          float *d = anchors_d + 4*a;
+          
+          float x1 = x + d[0];
+          float y1 = y + d[1];
+          float x2 = x + d[2];
+          float y2 = y + d[3];
+          float w = x2 - x1 + 1.0f;
+          float h = y2 - y1 + 1.0f;
+          float pred_ctr_x = box.x * w + x1 + 0.5f * w;
+          float pred_ctr_y = box.y * h + y1 + 0.5f * h;
+          float pred_w = exp(box.z) * w;
+          float pred_h = exp(box.w) * h;
+
+          box = float4{
+            max(0.0f, pred_ctr_x - 0.5f * pred_w),
+            max(0.0f, pred_ctr_y - 0.5f * pred_h),
+            min(pred_ctr_x + 0.5f * pred_w - 1.0f, width * scale - 1.0f),
+            min(pred_ctr_y + 0.5f * pred_h - 1.0f, height * scale - 1.0f)
+          };
+        }
+        
+        return thrust::make_tuple(in_scores[i], box, cls);
+      });
+
+    // Zero-out unused scores
+    if (num_detections < top_n) {
+      thrust::fill(on_stream, out_scores + num_detections,
+        out_scores + top_n, 0.0f);
+      thrust::fill(on_stream, out_classes + num_detections,
+        out_classes + top_n, 0.0f);
+    }
+  }
+
+  return 0;
+}
+
+}
+}

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+
+#include <vector>
+
+namespace odtk {
+namespace cuda {
+
+int decode(int batch_size,
+    const void *const *inputs, void *const *outputs,
+    size_t height, size_t width, size_t scale,
+    size_t num_anchors, size_t num_classes,
+    const std::vector<float> &anchors, float score_thresh, int top_n,
+    void *workspace, size_t workspace_size, cudaStream_t stream);
+
+
+}
+}
\ No newline at end of file

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include "decode_rotate.h"
+#include "utils.h"
+
+#include <algorithm>
+#include <cstdint>
+
+#include <thrust/device_ptr.h>
+#include <thrust/sequence.h>
+#include <thrust/execution_policy.h>
+#include <thrust/gather.h>
+#include <thrust/tabulate.h>
+#include <thrust/count.h>
+#include <thrust/find.h>
+#include <cub/device/device_radix_sort.cuh>
+#include <cub/iterator/counting_input_iterator.cuh>
+
+namespace odtk {
+namespace cuda {
+
+int decode_rotate(int batch_size,
+  const void *const *inputs, void *const *outputs,
+  size_t height, size_t width, size_t scale,
+  size_t num_anchors, size_t num_classes,
+  const std::vector<float> &anchors, float score_thresh, int top_n,
+  void *workspace, size_t workspace_size, cudaStream_t stream) {
+
+  int scores_size = num_anchors * num_classes * height * width;
+
+  if (!workspace || !workspace_size) {
+    // Return required scratch space size cub style
+    workspace_size  = get_size_aligned<float>(anchors.size()); // anchors
+    workspace_size += get_size_aligned<bool>(scores_size);     // flags
+    workspace_size += get_size_aligned<int>(scores_size);      // indices
+    workspace_size += get_size_aligned<int>(scores_size);      // indices_sorted
+    workspace_size += get_size_aligned<float>(scores_size);    // scores
+    workspace_size += get_size_aligned<float>(scores_size);    // scores_sorted
+
+    size_t temp_size_flag = 0;
+    cub::DeviceSelect::Flagged((void *)nullptr, temp_size_flag,
+      cub::CountingInputIterator<int>(scores_size),
+      (bool *)nullptr, (int *)nullptr, (int *)nullptr, scores_size);
+    size_t temp_size_sort = 0;
+    cub::DeviceRadixSort::SortPairsDescending((void *)nullptr, temp_size_sort,
+      (float *)nullptr, (float *)nullptr, (int *)nullptr, (int *)nullptr, scores_size);
+    workspace_size += std::max(temp_size_flag, temp_size_sort);
+
+    return workspace_size;
+  }
+
+  auto anchors_d = get_next_ptr<float>(anchors.size(), workspace, workspace_size);
+  cudaMemcpyAsync(anchors_d, anchors.data(), anchors.size() * sizeof *anchors_d, cudaMemcpyHostToDevice, stream);
+
+  auto on_stream = thrust::cuda::par.on(stream);
+
+  auto flags = get_next_ptr<bool>(scores_size, workspace, workspace_size);
+  auto indices = get_next_ptr<int>(scores_size, workspace, workspace_size);
+  auto indices_sorted = get_next_ptr<int>(scores_size, workspace, workspace_size);
+  auto scores = get_next_ptr<float>(scores_size, workspace, workspace_size);
+  auto scores_sorted = get_next_ptr<float>(scores_size, workspace, workspace_size);
+
+  for (int batch = 0; batch < batch_size; batch++) {
+    auto in_scores = static_cast<const float *>(inputs[0]) + batch * scores_size;
+    auto in_boxes = static_cast<const float *>(inputs[1]) + batch * (scores_size / num_classes) * 6; //From 4
+
+    auto out_scores = static_cast<float *>(outputs[0]) + batch * top_n;
+    auto out_boxes = static_cast<float6 *>(outputs[1]) + batch * top_n; // From float4
+    auto out_classes = static_cast<float *>(outputs[2]) + batch * top_n;
+
+    // Discard scores below threshold
+    thrust::transform(on_stream, in_scores, in_scores + scores_size,
+      flags, thrust::placeholders::_1 > score_thresh);
+
+    int *num_selected = reinterpret_cast<int *>(indices_sorted);
+    cub::DeviceSelect::Flagged(workspace, workspace_size, cub::CountingInputIterator<int>(0),
+      flags, indices, num_selected, scores_size, stream);
+    cudaStreamSynchronize(stream);
+    int num_detections = *thrust::device_pointer_cast(num_selected);
+
+    // Only keep top n scores
+    auto indices_filtered = indices;
+    if (num_detections > top_n) {
+      thrust::gather(on_stream, indices, indices + num_detections,
+        in_scores, scores);
+      cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
+        scores, scores_sorted, indices, indices_sorted, num_detections, 0, sizeof(*scores)*8, stream);
+        indices_filtered = indices_sorted;
+        num_detections = top_n;
+    }
+
+    // Gather boxes
+    bool has_anchors = !anchors.empty();
+    thrust::transform(on_stream, indices_filtered, indices_filtered + num_detections,
+      thrust::make_zip_iterator(thrust::make_tuple(out_scores, out_boxes, out_classes)),
+      [=] __device__ (int i) {
+        int x = i % width;
+        int y = (i / width) % height;
+        int a = (i / num_classes / height / width) % num_anchors;
+        int cls = (i / height / width) % num_classes;
+
+        float6 box = make_float6(
+          make_float4(
+            in_boxes[((a * 6 + 0) * height + y) * width + x],
+            in_boxes[((a * 6 + 1) * height + y) * width + x],
+            in_boxes[((a * 6 + 2) * height + y) * width + x],
+            in_boxes[((a * 6 + 3) * height + y) * width + x]
+          ),
+          make_float2(
+            in_boxes[((a * 6 + 4) * height + y) * width + x],
+            in_boxes[((a * 6 + 5) * height + y) * width + x]
+          )
+        );
+
+        if (has_anchors) {
+          // Add anchors offsets to deltas
+          float x = (i % width) * scale;
+          float y = ((i / width)  % height) * scale;
+          float *d = anchors_d + 4*a;
+
+          float x1 = x + d[0];
+          float y1 = y + d[1];
+          float x2 = x + d[2];
+          float y2 = y + d[3];
+
+          float w = x2 - x1 + 1.0f;
+          float h = y2 - y1 + 1.0f;
+          float pred_ctr_x = box.x1 * w + x1 + 0.5f * w;
+          float pred_ctr_y = box.y1 * h + y1 + 0.5f * h;
+          float pred_w = exp(box.x2) * w;
+          float pred_h = exp(box.y2) * h;
+          float pred_sin = box.s;
+          float pred_cos = box.c;
+
+          box = make_float6(
+            make_float4(
+              max(0.0f, pred_ctr_x - 0.5f * pred_w),
+              max(0.0f, pred_ctr_y - 0.5f * pred_h),
+              min(pred_ctr_x + 0.5f * pred_w - 1.0f, width * scale - 1.0f),
+              min(pred_ctr_y + 0.5f * pred_h - 1.0f, height * scale - 1.0f)
+            ),
+            make_float2(pred_sin, pred_cos)
+          );
+        }
+
+        return thrust::make_tuple(in_scores[i], box, cls);
+      });
+
+    // Zero-out unused scores
+    if (num_detections < top_n) {
+      thrust::fill(on_stream, out_scores + num_detections,
+        out_scores + top_n, 0.0f);
+      thrust::fill(on_stream, out_classes + num_detections,
+        out_classes + top_n, 0.0f);
+    }
+  }
+
+  return 0;
+}
+
+}
+}

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+
+#include <vector>
+
+namespace odtk {
+namespace cuda {
+
+int decode_rotate(int batchSize,
+    const void *const *inputs, void *const *outputs,
+    size_t height, size_t width, size_t scale,
+    size_t num_anchors, size_t num_classes,
+    const std::vector<float> &anchors, float score_thresh, int top_n,
+    void *workspace, size_t workspace_size, cudaStream_t stream);
+
+}
+}
\ No newline at end of file

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include "nms.h"
+#include "utils.h"
+
+#include <algorithm>
+#include <iostream>
+#include <stdexcept>
+#include <cstdint>
+#include <vector>
+#include <cmath>
+
+#include <cuda.h>
+#include <thrust/device_ptr.h>
+#include <thrust/sequence.h>
+#include <thrust/execution_policy.h>
+#include <thrust/gather.h>
+#include <cub/device/device_radix_sort.cuh>
+#include <cub/iterator/counting_input_iterator.cuh>
+
+namespace odtk {
+namespace cuda {
+
+__global__ void nms_kernel(
+      const int num_per_thread, const float threshold, const int num_detections,
+      const int *indices, float *scores, const float *classes, const float4 *boxes) {
+
+  // Go through detections by descending score
+  for (int m = 0; m < num_detections; m++) {
+    for (int n = 0; n < num_per_thread; n++) {
+      int i = threadIdx.x * num_per_thread + n;
+      if (i < num_detections && m < i && scores[m] > 0.0f) {
+        int idx = indices[i];
+        int max_idx = indices[m];
+        int icls = classes[idx];
+        int mcls = classes[max_idx];
+        if (mcls == icls) {
+          float4 ibox = boxes[idx];
+          float4 mbox = boxes[max_idx];
+          float x1 = max(ibox.x, mbox.x);
+          float y1 = max(ibox.y, mbox.y);
+          float x2 = min(ibox.z, mbox.z);
+          float y2 = min(ibox.w, mbox.w);
+          float w = max(0.0f, x2 - x1 + 1);
+          float h = max(0.0f, y2 - y1 + 1);
+          float iarea = (ibox.z - ibox.x + 1) * (ibox.w - ibox.y + 1);
+          float marea = (mbox.z - mbox.x + 1) * (mbox.w - mbox.y + 1);
+          float inter = w * h;
+          float overlap = inter / (iarea + marea - inter);
+          if (overlap > threshold) {
+            scores[i] = 0.0f;
+          }
+        }
+      }
+    }
+
+    // Sync discarded detections
+    __syncthreads();
+  }
+}
+
+int nms(int batch_size,
+  const void *const *inputs, void *const *outputs,
+  size_t count, int detections_per_im, float nms_thresh,
+  void *workspace, size_t workspace_size, cudaStream_t stream) {
+
+  if (!workspace || !workspace_size) {
+    // Return required scratch space size cub style
+    workspace_size  = get_size_aligned<bool>(count);  // flags
+    workspace_size += get_size_aligned<int>(count);   // indices
+    workspace_size += get_size_aligned<int>(count);   // indices_sorted
+    workspace_size += get_size_aligned<float>(count); // scores
+    workspace_size += get_size_aligned<float>(count); // scores_sorted
+  
+    size_t temp_size_flag = 0;
+    cub::DeviceSelect::Flagged((void *)nullptr, temp_size_flag,
+      cub::CountingInputIterator<int>(count),
+      (bool *)nullptr, (int *)nullptr, (int *)nullptr, count);
+    size_t temp_size_sort = 0;
+    cub::DeviceRadixSort::SortPairsDescending((void *)nullptr, temp_size_sort,
+      (float *)nullptr, (float *)nullptr, (int *)nullptr, (int *)nullptr, count);
+    workspace_size += std::max(temp_size_flag, temp_size_sort);
+
+    return workspace_size;
+  }
+
+  auto on_stream = thrust::cuda::par.on(stream);
+
+  auto flags = get_next_ptr<bool>(count, workspace, workspace_size);
+  auto indices = get_next_ptr<int>(count, workspace, workspace_size);
+  auto indices_sorted = get_next_ptr<int>(count, workspace, workspace_size);
+  auto scores = get_next_ptr<float>(count, workspace, workspace_size);
+  auto scores_sorted = get_next_ptr<float>(count, workspace, workspace_size);
+
+  for (int batch = 0; batch < batch_size; batch++) {
+    auto in_scores = static_cast<const float *>(inputs[0]) + batch * count;
+    auto in_boxes = static_cast<const float4 *>(inputs[1]) + batch * count;
+    auto in_classes = static_cast<const float *>(inputs[2]) + batch * count;
+    
+    auto out_scores = static_cast<float *>(outputs[0]) + batch * detections_per_im;
+    auto out_boxes = static_cast<float4 *>(outputs[1]) + batch * detections_per_im;
+    auto out_classes = static_cast<float *>(outputs[2]) + batch * detections_per_im;
+    
+    // Discard null scores
+    thrust::transform(on_stream, in_scores, in_scores + count,
+      flags, thrust::placeholders::_1 > 0.0f);
+
+    int *num_selected = reinterpret_cast<int *>(indices_sorted);
+    cub::DeviceSelect::Flagged(workspace, workspace_size, cub::CountingInputIterator<int>(0),
+      flags, indices, num_selected, count, stream);
+    cudaStreamSynchronize(stream);
+    int num_detections = *thrust::device_pointer_cast(num_selected);
+
+    // Sort scores and corresponding indices
+    thrust::gather(on_stream, indices, indices + num_detections, in_scores, scores);
+    cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
+      scores, scores_sorted, indices, indices_sorted, num_detections, 0, sizeof(*scores)*8, stream);
+
+    // Launch actual NMS kernel - 1 block with each thread handling n detections
+    const int max_threads = 1024;
+    int num_per_thread = ceil((float)num_detections / max_threads);
+    nms_kernel<<<1, max_threads, 0, stream>>>(num_per_thread, nms_thresh, num_detections,
+      indices_sorted, scores_sorted, in_classes, in_boxes);
+
+    // Re-sort with updated scores
+    cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
+      scores_sorted, scores, indices_sorted, indices, num_detections, 0, sizeof(*scores)*8, stream);
+
+    // Gather filtered scores, boxes, classes
+    num_detections = min(detections_per_im, num_detections);
+    cudaMemcpyAsync(out_scores, scores, num_detections * sizeof *scores, cudaMemcpyDeviceToDevice, stream);
+    if (num_detections < detections_per_im) {
+      thrust::fill_n(on_stream, out_scores + num_detections, detections_per_im - num_detections, 0);
+    }
+    thrust::gather(on_stream, indices, indices + num_detections, in_boxes, out_boxes);
+    thrust::gather(on_stream, indices, indices + num_detections, in_classes, out_classes);
+  }
+  
+  return 0;
+}
+
+}
+}

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+
+namespace odtk {
+namespace cuda {
+
+int nms(int batchSize,
+    const void *const *inputs, void *const *outputs,
+    size_t count, int detections_per_im, float nms_thresh,
+    void *workspace, size_t workspace_size, cudaStream_t stream);
+
+}
+}
\ No newline at end of file

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+
+namespace odtk {
+namespace cuda {
+
+int nms_rotate(int batchSize,
+    const void *const *inputs, void *const *outputs,
+    size_t count, int detections_per_im, float nms_thresh,
+    void *workspace, size_t workspace_size, cudaStream_t stream);
+
+int iou(
+    const void *const *inputs, void *const *outputs,
+    int num_boxes, int num_anchors, cudaStream_t stream);
+}
+}
\ No newline at end of file

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+#include <stdexcept>
+#include <cstdint>
+#include <thrust/functional.h>
+
+#define CUDA_ALIGN 256
+
+struct float6
+{
+  float x1, y1, x2, y2, s, c; 
+};
+
+inline __host__ __device__ float6 make_float6(float4 f, float2 t)
+{
+  float6 fs;
+  fs.x1 = f.x; fs.y1 = f.y; fs.x2 = f.z; fs.y2 = f.w; fs.s = t.x; fs.c = t.y;
+  return fs;
+}
+
+template <typename T>
+inline size_t get_size_aligned(size_t num_elem) {
+    size_t size = num_elem * sizeof(T);
+    size_t extra_align = 0;
+    if (size % CUDA_ALIGN != 0) {
+        extra_align = CUDA_ALIGN - size % CUDA_ALIGN;
+    }
+    return size + extra_align;
+}
+
+template <typename T>
+inline T *get_next_ptr(size_t num_elem, void *&workspace, size_t &workspace_size) {
+  size_t size = get_size_aligned<T>(num_elem);
+  if (size > workspace_size) {
+    throw std::runtime_error("Workspace is too small!");
+  }
+  workspace_size -= size;
+  T *ptr = reinterpret_cast<T *>(workspace);
+  workspace = reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(workspace) + size);
+  return ptr;
+}

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include "engine.h"
+
+#include <iostream>
+#include <fstream>
+
+#include <NvOnnxConfig.h>
+#include <NvOnnxParser.h>
+
+#include "plugins/DecodePlugin.h"
+#include "plugins/NMSPlugin.h"
+#include "plugins/DecodeRotatePlugin.h"
+#include "plugins/NMSRotatePlugin.h"
+#include "calibrator.h"
+
+#include <stdio.h>
+#include <string>
+
+using namespace nvinfer1;
+using namespace nvonnxparser;
+
+namespace odtk {
+
+class Logger : public ILogger {
+public:
+    Logger(bool verbose)
+        : _verbose(verbose) {
+    }
+
+    void log(Severity severity, const char *msg) override {
+        if (_verbose || ((severity != Severity::kINFO) && (severity != Severity::kVERBOSE)))
+            cout << msg << endl;
+    }
+
+private:
+   bool _verbose{false};
+};
+
+void Engine::_load(const string &path) {
+    ifstream file(path, ios::in | ios::binary);
+    file.seekg (0, file.end);
+    size_t size = file.tellg();
+    file.seekg (0, file.beg);
+
+    char *buffer = new char[size];
+    file.read(buffer, size);
+    file.close();
+
+    _engine = _runtime->deserializeCudaEngine(buffer, size, nullptr);
+
+    delete[] buffer;
+}
+
+void Engine::_prepare() {
+    _context = _engine->createExecutionContext();
+    _context->setOptimizationProfileAsync(0, _stream);
+    cudaStreamCreate(&_stream);
+}
+
+Engine::Engine(const string &engine_path, bool verbose) {
+    Logger logger(verbose);
+    _runtime = createInferRuntime(logger);
+    _load(engine_path);
+    _prepare();
+}
+
+Engine::~Engine() {
+    if (_stream) cudaStreamDestroy(_stream);
+    if (_context) _context->destroy();
+    if (_engine) _engine->destroy();
+    if (_runtime) _runtime->destroy();
+}
+
+Engine::Engine(const char *onnx_model, size_t onnx_size, const vector<int>& dynamic_batch_opts,
+    string precision, float score_thresh, int top_n, const vector<vector<float>>& anchors, 
+    bool rotated, float nms_thresh, int detections_per_im, const vector<string>& calibration_images,
+    string model_name, string calibration_table, bool verbose, size_t workspace_size) {
+
+    Logger logger(verbose);
+    _runtime = createInferRuntime(logger);
+
+    bool fp16 = precision.compare("FP16") == 0;
+    bool int8 = precision.compare("INT8") == 0;
+
+    // Create builder
+    auto builder = createInferBuilder(logger);
+    const auto builderConfig = builder->createBuilderConfig();
+    // Allow use of FP16 layers when running in INT8
+    if(fp16 || int8) builderConfig->setFlag(BuilderFlag::kFP16);
+    builderConfig->setMaxWorkspaceSize(workspace_size);
+    
+    // Parse ONNX FCN
+    cout << "Building " << precision << " core model..." << endl;
+    const auto flags = 1U << static_cast<int>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
+    auto network = builder->createNetworkV2(flags);
+    auto parser = createParser(*network, logger);
+    parser->parse(onnx_model, onnx_size);
+    
+    auto input = network->getInput(0);
+    auto inputDims = input->getDimensions();
+    auto profile = builder->createOptimizationProfile();
+    auto inputName = input->getName();
+    auto profileDimsmin = Dims4{dynamic_batch_opts[0], inputDims.d[1], inputDims.d[2], inputDims.d[3]};
+    auto profileDimsopt = Dims4{dynamic_batch_opts[1], inputDims.d[1], inputDims.d[2], inputDims.d[3]};
+    auto profileDimsmax = Dims4{dynamic_batch_opts[2], inputDims.d[1], inputDims.d[2], inputDims.d[3]};
+
+    profile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMIN, profileDimsmin);
+    profile->setDimensions(inputName, nvinfer1::OptProfileSelector::kOPT, profileDimsopt);
+    profile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMAX, profileDimsmax);
+    
+    if(profile->isValid())
+        builderConfig->addOptimizationProfile(profile);
+
+    std::unique_ptr<Int8EntropyCalibrator> calib;
+    if (int8) {
+        builderConfig->setFlag(BuilderFlag::kINT8);
+        // Calibration is performed using kOPT values of the profile.
+        // Calibration batch size must match this profile.
+        builderConfig->setCalibrationProfile(profile);
+        ImageStream stream(dynamic_batch_opts[1], inputDims, calibration_images);
+        calib = std::unique_ptr<Int8EntropyCalibrator>(new Int8EntropyCalibrator(stream, model_name, calibration_table));
+        builderConfig->setInt8Calibrator(calib.get());
+    }
+
+    // Add decode plugins
+    cout << "Building accelerated plugins..." << endl;
+    vector<DecodePlugin> decodePlugins;
+    vector<DecodeRotatePlugin> decodeRotatePlugins;
+    vector<ITensor *> scores, boxes, classes;
+    auto nbOutputs = network->getNbOutputs();
+    
+    for (int i = 0; i < nbOutputs / 2; i++) {
+        auto classOutput = network->getOutput(i);
+        auto boxOutput = network->getOutput(nbOutputs / 2 + i);
+        auto outputDims = classOutput->getDimensions();
+        int scale = inputDims.d[2] / outputDims.d[2];
+        auto decodePlugin = DecodePlugin(score_thresh, top_n, anchors[i], scale);
+        auto decodeRotatePlugin = DecodeRotatePlugin(score_thresh, top_n, anchors[i], scale);
+        decodePlugins.push_back(decodePlugin); 
+        decodeRotatePlugins.push_back(decodeRotatePlugin);
+        vector<ITensor *> inputs = {classOutput, boxOutput};
+        auto layer = (!rotated) ? network->addPluginV2(inputs.data(), inputs.size(), decodePlugin) \
+                    : network->addPluginV2(inputs.data(), inputs.size(), decodeRotatePlugin);
+        scores.push_back(layer->getOutput(0));
+        boxes.push_back(layer->getOutput(1));
+        classes.push_back(layer->getOutput(2));
+    }
+
+    // Cleanup outputs
+    for (int i = 0; i < nbOutputs; i++) {
+        auto output = network->getOutput(0);
+        network->unmarkOutput(*output);
+    }
+
+    // Concat tensors from each feature map
+    vector<ITensor *> concat;
+    for (auto tensors : {scores, boxes, classes}) {
+        auto layer = network->addConcatenation(tensors.data(), tensors.size());
+        concat.push_back(layer->getOutput(0));
+    }
+    
+    // Add NMS plugin
+    auto nmsPlugin = NMSPlugin(nms_thresh, detections_per_im);
+    auto nmsRotatePlugin = NMSRotatePlugin(nms_thresh, detections_per_im);
+    auto layer = (!rotated) ? network->addPluginV2(concat.data(), concat.size(), nmsPlugin) \
+                : network->addPluginV2(concat.data(), concat.size(), nmsRotatePlugin);
+    vector<string> names = {"scores", "boxes", "classes"};
+    for (int i = 0; i < layer->getNbOutputs(); i++) {
+        auto output = layer->getOutput(i);
+        network->markOutput(*output);
+        output->setName(names[i].c_str());
+    }
+    
+    // Build engine
+    cout << "Applying optimizations and building TRT CUDA engine..." << endl;
+    _engine = builder->buildEngineWithConfig(*network, *builderConfig);
+
+    // Housekeeping
+    parser->destroy();
+    network->destroy();
+    builderConfig->destroy();
+    builder->destroy();
+
+    _prepare();
+}
+
+void Engine::save(const string &path) {
+    cout << "Writing to " << path << "..." << endl;
+    auto serialized = _engine->serialize();
+    ofstream file(path, ios::out | ios::binary);
+    file.write(reinterpret_cast<const char*>(serialized->data()), serialized->size());
+
+    serialized->destroy();    
+}
+
+void Engine::infer(vector<void *> &buffers, int batch){
+    auto dims = _engine->getBindingDimensions(0);
+    _context->setBindingDimensions(0, Dims4(batch, dims.d[1], dims.d[2], dims.d[3]));
+    _context->enqueueV2(buffers.data(), _stream, nullptr);
+    cudaStreamSynchronize(_stream);
+}
+
+vector<int> Engine::getInputSize() {
+    auto dims = _engine->getBindingDimensions(0);
+    return {dims.d[2], dims.d[3]};
+}
+
+int Engine::getMaxBatchSize() {
+    return _engine->getMaxBatchSize();
+}
+
+int Engine::getMaxDetections() {
+    return _engine->getBindingDimensions(1).d[1];
+}
+
+int Engine::getStride() {
+    return 1;
+}
+
+}

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include <NvInfer.h>
+
+#include <cuda_runtime.h>
+
+using namespace std;
+using namespace nvinfer1;
+
+namespace odtk {
+
+// RetinaNet wrapper around TensorRT CUDA engine
+class Engine {
+public:
+    // Create engine from engine path
+    Engine(const string &engine_path, bool verbose=false);
+
+    // Create engine from serialized onnx model
+    
+    Engine(const char *onnx_model, size_t onnx_size, const vector<int>& dynamic_batch_opts,
+        string precision, float score_thresh, int top_n, const vector<vector<float>>& anchors,
+        bool rotated, float nms_thresh, int detections_per_im, const vector<string>& calibration_images, 
+        string model_name, string calibration_table, bool verbose, size_t workspace_size=(1ULL << 30));
+    
+    ~Engine();
+
+    // Save model to path
+    void save(const string &path);
+
+    // Infer using pre-allocated GPU buffers {data, scores, boxes, classes}
+    void infer(vector<void *> &buffers, int batch);
+
+    // Get (h, w) size of the fixed input
+    vector<int> getInputSize();
+
+    // Get max allowed batch size
+    int getMaxBatchSize();
+
+    // Get max number of detections
+    int getMaxDetections();
+
+    // Get stride
+    int getStride();
+    
+private:
+    IRuntime *_runtime = nullptr;
+    ICudaEngine *_engine = nullptr;
+    IExecutionContext *_context = nullptr;
+    cudaStream_t _stream = nullptr;
+
+    void _load(const string &path);
+    void _prepare();
+
+};
+
+}

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#include <algorithm>
+#include <iostream>
+#include <stdexcept>
+#include <cstdint>
+#include <cmath>
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <vector>
+#include <optional>
+
+#include "engine.h"
+#include "cuda/decode.h"
+#include "cuda/decode_rotate.h"
+#include "cuda/nms.h"
+#include "cuda/nms_iou.h"
+#include <stdio.h>
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+
+vector<at::Tensor> iou(at::Tensor boxes, at::Tensor anchors) {
+
+    CHECK_INPUT(boxes);
+    CHECK_INPUT(anchors);
+
+    int num_boxes = boxes.numel() / 8;
+    int num_anchors = anchors.numel() / 8;
+    auto options = boxes.options();
+
+    auto iou_vals = at::zeros({num_boxes*num_anchors}, options);
+
+    // Calculate Polygon IOU
+    vector<void *> inputs = {boxes.data_ptr(), anchors.data_ptr()};
+    vector<void *> outputs = {iou_vals.data_ptr()};
+
+    odtk::cuda::iou(inputs.data(), outputs.data(), num_boxes, num_anchors, at::cuda::getCurrentCUDAStream());
+
+    auto shape = std::vector<int64_t>{num_anchors, num_boxes};
+
+    return {iou_vals.reshape(shape)};
+}
+
+vector<at::Tensor> decode(at::Tensor cls_head, at::Tensor box_head,
+        vector<float> &anchors, int scale, float score_thresh, int top_n, bool rotated=false) {
+
+    CHECK_INPUT(cls_head);
+    CHECK_INPUT(box_head);
+
+    int num_boxes = (!rotated) ? 4 : 6;
+    int batch = cls_head.size(0);
+    int num_anchors = anchors.size() / 4;
+    int num_classes = cls_head.size(1) / num_anchors;
+    int height = cls_head.size(2);
+    int width = cls_head.size(3);
+    auto options = cls_head.options();
+
+    auto scores = at::zeros({batch, top_n}, options);
+    auto boxes = at::zeros({batch, top_n, num_boxes}, options);
+    auto classes = at::zeros({batch, top_n}, options);
+
+    vector<void *> inputs = {cls_head.data_ptr(), box_head.data_ptr()};
+    vector<void *> outputs = {scores.data_ptr(), boxes.data_ptr(), classes.data_ptr()};
+
+    if(!rotated) {
+        // Create scratch buffer
+        int size = odtk::cuda::decode(batch, nullptr, nullptr, height, width, scale,
+            num_anchors, num_classes, anchors, score_thresh, top_n, nullptr, 0, nullptr);
+        auto scratch = at::zeros({size}, options.dtype(torch::kUInt8));
+
+        // Decode boxes
+        odtk::cuda::decode(batch, inputs.data(), outputs.data(), height, width, scale,
+            num_anchors, num_classes, anchors, score_thresh, top_n,
+            scratch.data_ptr(), size, at::cuda::getCurrentCUDAStream());
+    
+    }
+    else {
+        // Create scratch buffer
+        int size = odtk::cuda::decode_rotate(batch, nullptr, nullptr, height, width, scale,
+            num_anchors, num_classes, anchors, score_thresh, top_n, nullptr, 0, nullptr);
+        auto scratch = at::zeros({size}, options.dtype(torch::kUInt8));
+
+        // Decode boxes
+        odtk::cuda::decode_rotate(batch, inputs.data(), outputs.data(), height, width, scale,
+            num_anchors, num_classes, anchors, score_thresh, top_n,
+            scratch.data_ptr(), size, at::cuda::getCurrentCUDAStream());
+    }
+    
+    return {scores, boxes, classes};
+}
+
+vector<at::Tensor> nms(at::Tensor scores, at::Tensor boxes, at::Tensor classes,
+        float nms_thresh, int detections_per_im, bool rotated=false) {
+
+    CHECK_INPUT(scores);
+    CHECK_INPUT(boxes);
+    CHECK_INPUT(classes);
+
+    int num_boxes = (!rotated) ? 4 : 6;
+    int batch = scores.size(0);
+    int count = scores.size(1);
+    auto options = scores.options();
+    auto nms_scores = at::zeros({batch, detections_per_im}, scores.options());
+    auto nms_boxes = at::zeros({batch, detections_per_im, num_boxes}, boxes.options());
+    auto nms_classes = at::zeros({batch, detections_per_im}, classes.options());
+
+    vector<void *> inputs = {scores.data_ptr(), boxes.data_ptr(), classes.data_ptr()};
+    vector<void *> outputs = {nms_scores.data_ptr(), nms_boxes.data_ptr(), nms_classes.data_ptr()};
+
+    if(!rotated) {
+        // Create scratch buffer
+        int size = odtk::cuda::nms(batch, nullptr, nullptr, count,
+            detections_per_im, nms_thresh, nullptr, 0, nullptr);
+        auto scratch = at::zeros({size}, options.dtype(torch::kUInt8));
+
+        // Perform NMS
+        odtk::cuda::nms(batch, inputs.data(), outputs.data(), count, detections_per_im, 
+            nms_thresh, scratch.data_ptr(), size, at::cuda::getCurrentCUDAStream());
+    }
+    else {
+        // Create scratch buffer
+        int size = odtk::cuda::nms_rotate(batch, nullptr, nullptr, count,
+            detections_per_im, nms_thresh, nullptr, 0, nullptr);
+        auto scratch = at::zeros({size}, options.dtype(torch::kUInt8));
+
+        // Perform NMS
+        odtk::cuda::nms_rotate(batch, inputs.data(), outputs.data(), count,
+            detections_per_im, nms_thresh, scratch.data_ptr(), size, at::cuda::getCurrentCUDAStream());
+    }
+    
+
+    return {nms_scores, nms_boxes, nms_classes};
+}
+
+vector<at::Tensor> infer(odtk::Engine &engine, at::Tensor data, bool rotated=false) {
+    CHECK_INPUT(data);
+
+    int num_boxes = (!rotated) ? 4 : 6;
+    int batch = data.size(0);
+    auto input_size = engine.getInputSize();
+    data = at::constant_pad_nd(data, {0, input_size[1] - data.size(3), 0, input_size[0] - data.size(2)});
+
+    int num_detections = engine.getMaxDetections();
+    auto scores = at::zeros({batch, num_detections}, data.options());
+    auto boxes = at::zeros({batch, num_detections, num_boxes}, data.options());
+    auto classes = at::zeros({batch, num_detections}, data.options());
+
+    vector<void *> buffers;
+    for (auto buffer : {data, scores, boxes, classes}) {
+        buffers.push_back(buffer.data<float>());
+    }
+
+    engine.infer(buffers, batch);
+
+    return {scores, boxes, classes};
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    pybind11::class_<odtk::Engine>(m, "Engine")
+        .def(pybind11::init<const char *, size_t, const vector<int>&, string, float, int,
+            const vector<vector<float>>&, bool, float, int, const vector<string>&, string, string, bool>())
+        .def("save", &odtk::Engine::save)
+        .def("infer", &odtk::Engine::infer)
+        .def_property_readonly("stride", &odtk::Engine::getStride)
+        .def_property_readonly("input_size", &odtk::Engine::getInputSize)
+        .def_static("load", [](const string &path) {
+            return new odtk::Engine(path);
+        })
+        .def("__call__", [](odtk::Engine &engine, at::Tensor data, bool rotated=false) {
+            return infer(engine, data, rotated);
+        });
+    m.def("decode", &decode);
+    m.def("nms", &nms);
+    m.def("iou", &iou);
+}
\ No newline at end of file

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once 
+
+#include <NvInfer.h>
+
+#include <cassert>
+#include <vector>
+
+#include "../cuda/decode.h"
+
+using namespace nvinfer1;
+
+#define RETINANET_PLUGIN_NAME "RetinaNetDecode"
+#define RETINANET_PLUGIN_VERSION "1"
+#define RETINANET_PLUGIN_NAMESPACE ""
+
+namespace odtk {
+
+class DecodePlugin : public IPluginV2DynamicExt {
+  float _score_thresh;
+  int _top_n;
+  std::vector<float> _anchors;
+  float _scale;
+
+  size_t _height;
+  size_t _width;
+  size_t _num_anchors;
+  size_t _num_classes;
+  mutable int size = -1;
+
+protected:
+  void deserialize(void const* data, size_t length) {
+    const char* d = static_cast<const char*>(data);
+    read(d, _score_thresh);
+    read(d, _top_n);
+    size_t anchors_size;
+    read(d, anchors_size);
+    while( anchors_size-- ) {
+      float val;
+      read(d, val);
+      _anchors.push_back(val);
+    }
+    read(d, _scale);
+    read(d, _height);
+    read(d, _width);
+    read(d, _num_anchors);
+    read(d, _num_classes);
+  }
+
+  size_t getSerializationSize() const override {
+    return sizeof(_score_thresh) + sizeof(_top_n)
+      + sizeof(size_t) + sizeof(float) * _anchors.size() + sizeof(_scale)
+      + sizeof(_height) + sizeof(_width) + sizeof(_num_anchors) + sizeof(_num_classes);
+  }
+
+  void serialize(void *buffer) const override {
+    char* d = static_cast<char*>(buffer);
+    write(d, _score_thresh);
+    write(d, _top_n);
+    write(d, _anchors.size());
+    for( auto &val : _anchors ) {
+      write(d, val);
+    }
+    write(d, _scale);
+    write(d, _height);
+    write(d, _width);
+    write(d, _num_anchors);
+    write(d, _num_classes);
+  }
+
+public:
+  DecodePlugin(float score_thresh, int top_n, std::vector<float> const& anchors, int scale)
+    : _score_thresh(score_thresh), _top_n(top_n), _anchors(anchors), _scale(scale) {}
+
+  DecodePlugin(float score_thresh, int top_n, std::vector<float> const& anchors, int scale,
+    size_t height, size_t width, size_t num_anchors, size_t num_classes)
+    : _score_thresh(score_thresh), _top_n(top_n), _anchors(anchors), _scale(scale),
+      _height(height), _width(width), _num_anchors(num_anchors), _num_classes(num_classes) {}
+
+  DecodePlugin(void const* data, size_t length) {
+      this->deserialize(data, length);
+  }
+
+  const char *getPluginType() const override {
+    return RETINANET_PLUGIN_NAME;
+  }
+ 
+  const char *getPluginVersion() const override {
+    return RETINANET_PLUGIN_VERSION;
+  }
+  
+  int getNbOutputs() const override {
+    return 3;
+  }
+
+  DimsExprs getOutputDimensions(int outputIndex, const DimsExprs *inputs,
+    int nbInputs, IExprBuilder &exprBuilder) override 
+  {
+    DimsExprs output(inputs[0]);
+    output.d[1] = exprBuilder.constant(_top_n * (outputIndex == 1 ? 4 : 1));
+    output.d[2] = exprBuilder.constant(1);
+    output.d[3] = exprBuilder.constant(1);
+
+    return output;
+  }
+
+  bool supportsFormatCombination(int pos, const PluginTensorDesc *inOut, 
+    int nbInputs, int nbOutputs) override
+  {
+    assert(nbInputs == 2);
+    assert(nbOutputs == 3);
+    assert(pos < 5);
+    return inOut[pos].type == DataType::kFLOAT && inOut[pos].format == nvinfer1::PluginFormat::kLINEAR;
+  }
+
+  int initialize() override { return 0; }
+
+  void terminate() override {}
+
+  size_t getWorkspaceSize(const PluginTensorDesc *inputs, 
+    int nbInputs, const PluginTensorDesc *outputs, int nbOutputs) const override 
+  {
+    if (size < 0) {
+      size = cuda::decode(inputs->dims.d[0], nullptr, nullptr, _height, _width, _scale,
+        _num_anchors, _num_classes, _anchors, _score_thresh, _top_n, 
+        nullptr, 0, nullptr);
+    }
+    return size;
+  }
+
+  int enqueue(const PluginTensorDesc *inputDesc, 
+    const PluginTensorDesc *outputDesc, const void *const *inputs, 
+    void *const *outputs, void *workspace, cudaStream_t stream)  
+  {
+    
+    return cuda::decode(inputDesc->dims.d[0], inputs, outputs, _height, _width, _scale,
+      _num_anchors, _num_classes, _anchors, _score_thresh, _top_n,
+      workspace, getWorkspaceSize(inputDesc, 2, outputDesc, 3), stream);
+    
+  }
+
+  void destroy() override {
+    delete this;
+  };
+
+  const char *getPluginNamespace() const override {
+    return RETINANET_PLUGIN_NAMESPACE;
+  }
+  
+  void setPluginNamespace(const char *N) override {}
+
+  DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const
+  {
+    assert(index < 3);
+    return DataType::kFLOAT;
+  }
+
+  void configurePlugin(const DynamicPluginTensorDesc *in, int nbInputs, 
+    const DynamicPluginTensorDesc *out, int nbOutputs)
+  {
+    assert(nbInputs == 2);
+    assert(nbOutputs == 3);
+    auto const& scores_dims = in[0].desc.dims;
+    auto const& boxes_dims = in[1].desc.dims;
+    assert(scores_dims.d[2] == boxes_dims.d[2]);
+    assert(scores_dims.d[3] == boxes_dims.d[3]);
+    _height = scores_dims.d[2];
+    _width = scores_dims.d[3];
+    _num_anchors = boxes_dims.d[1] / 4; 
+    _num_classes = scores_dims.d[1] / _num_anchors;
+  }
+
+  IPluginV2DynamicExt *clone() const  {
+    return new DecodePlugin(_score_thresh, _top_n, _anchors, _scale, _height, _width, 
+      _num_anchors, _num_classes);
+  }
+  
+private:
+  template<typename T> void write(char*& buffer, const T& val) const {
+    *reinterpret_cast<T*>(buffer) = val;
+    buffer += sizeof(T);
+  }
+
+  template<typename T> void read(const char*& buffer, T& val) {
+    val = *reinterpret_cast<const T*>(buffer);
+    buffer += sizeof(T);
+  }
+};
+
+class DecodePluginCreator : public IPluginCreator {
+public:
+  DecodePluginCreator() {}
+
+  const char *getPluginName () const override {
+    return RETINANET_PLUGIN_NAME;
+  }
+
+  const char *getPluginVersion () const override {
+    return RETINANET_PLUGIN_VERSION;
+  }
+ 
+  const char *getPluginNamespace() const override {
+    return RETINANET_PLUGIN_NAMESPACE;
+  }
+
+  
+  IPluginV2DynamicExt *deserializePlugin (const char *name, const void *serialData, size_t serialLength) override {
+    return new DecodePlugin(serialData, serialLength);
+  }
+
+  void setPluginNamespace(const char *N) override {}
+  const PluginFieldCollection *getFieldNames() override { return nullptr; }
+  IPluginV2DynamicExt *createPlugin (const char *name, const PluginFieldCollection *fc) override { return nullptr; }
+};
+
+REGISTER_TENSORRT_PLUGIN(DecodePluginCreator);
+
+}
+
+#undef RETINANET_PLUGIN_NAME
+#undef RETINANET_PLUGIN_VERSION
+#undef RETINANET_PLUGIN_NAMESPACE
\ No newline at end of file

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once 
+
+#include <NvInfer.h>
+
+#include <cassert>
+#include <vector>
+
+#include "../cuda/decode_rotate.h"
+
+using namespace nvinfer1;
+
+#define RETINANET_PLUGIN_NAME "RetinaNetDecodeRotate"
+#define RETINANET_PLUGIN_VERSION "1"
+#define RETINANET_PLUGIN_NAMESPACE ""
+
+namespace odtk {
+
+class DecodeRotatePlugin : public IPluginV2DynamicExt {
+  float _score_thresh;
+  int _top_n;
+  std::vector<float> _anchors;
+  float _scale;
+
+  size_t _height;
+  size_t _width;
+  size_t _num_anchors;
+  size_t _num_classes;
+  mutable int size = -1;
+
+protected:
+  void deserialize(void const* data, size_t length) {
+    const char* d = static_cast<const char*>(data);
+    read(d, _score_thresh);
+    read(d, _top_n);
+    size_t anchors_size;
+    read(d, anchors_size);
+    while( anchors_size-- ) {
+      float val;
+      read(d, val);
+      _anchors.push_back(val);
+    }
+    read(d, _scale);
+    read(d, _height);
+    read(d, _width);
+    read(d, _num_anchors);
+    read(d, _num_classes);
+  }
+
+  size_t getSerializationSize() const override {
+    return sizeof(_score_thresh) + sizeof(_top_n)
+      + sizeof(size_t) + sizeof(float) * _anchors.size() + sizeof(_scale)
+      + sizeof(_height) + sizeof(_width) + sizeof(_num_anchors) + sizeof(_num_classes);
+  }
+
+  void serialize(void *buffer) const override {
+    char* d = static_cast<char*>(buffer);
+    write(d, _score_thresh);
+    write(d, _top_n);
+    write(d, _anchors.size());
+    for( auto &val : _anchors ) {
+      write(d, val);
+    }
+    write(d, _scale);
+    write(d, _height);
+    write(d, _width);
+    write(d, _num_anchors);
+    write(d, _num_classes);
+  }
+
+public:
+  DecodeRotatePlugin(float score_thresh, int top_n, std::vector<float> const& anchors, int scale)
+    : _score_thresh(score_thresh), _top_n(top_n), _anchors(anchors), _scale(scale) {}
+
+  DecodeRotatePlugin(float score_thresh, int top_n, std::vector<float> const& anchors, int scale,
+    size_t height, size_t width, size_t num_anchors, size_t num_classes)
+    : _score_thresh(score_thresh), _top_n(top_n), _anchors(anchors), _scale(scale),
+      _height(height), _width(width), _num_anchors(num_anchors), _num_classes(num_classes) {}
+
+  DecodeRotatePlugin(void const* data, size_t length) {
+      this->deserialize(data, length);
+  }
+
+  const char *getPluginType() const override {
+    return RETINANET_PLUGIN_NAME;
+  }
+
+  const char *getPluginVersion() const override {
+    return RETINANET_PLUGIN_VERSION;
+  }
+
+  int getNbOutputs() const override {
+    return 3;
+  }
+
+  DimsExprs getOutputDimensions(int outputIndex, const DimsExprs *inputs,
+    int nbInputs, IExprBuilder &exprBuilder) override 
+  {
+    DimsExprs output(inputs[0]);
+    output.d[1] = exprBuilder.constant(_top_n * (outputIndex == 1 ? 6 : 1));
+    output.d[2] = exprBuilder.constant(1);
+    output.d[3] = exprBuilder.constant(1);
+
+    return output;
+  }
+
+
+  bool supportsFormatCombination(int pos, const PluginTensorDesc *inOut, 
+    int nbInputs, int nbOutputs) override
+  {
+    assert(nbInputs == 2);
+    assert(nbOutputs == 3);
+    assert(pos < 5);
+    return inOut[pos].type == DataType::kFLOAT && inOut[pos].format == nvinfer1::PluginFormat::kLINEAR;
+  }
+
+
+  int initialize() override { return 0; }
+
+  void terminate() override {}
+
+  size_t getWorkspaceSize(const PluginTensorDesc *inputs, 
+    int nbInputs, const PluginTensorDesc *outputs, int nbOutputs) const override 
+  {
+    if (size < 0) {
+      size = cuda::decode_rotate(inputs->dims.d[0], nullptr, nullptr, _height, _width, _scale,
+        _num_anchors, _num_classes, _anchors, _score_thresh, _top_n, 
+        nullptr, 0, nullptr);
+    }
+    return size;
+  }
+
+  int enqueue(const PluginTensorDesc *inputDesc, 
+    const PluginTensorDesc *outputDesc, const void *const *inputs, 
+    void *const *outputs, void *workspace, cudaStream_t stream) override 
+  {
+    return cuda::decode_rotate(inputDesc->dims.d[0], inputs, outputs, _height, _width, _scale,
+      _num_anchors, _num_classes, _anchors, _score_thresh, _top_n,
+      workspace, getWorkspaceSize(inputDesc, 2, outputDesc, 3), stream);
+  }
+
+  void destroy() override {
+    delete this;
+  };
+
+  const char *getPluginNamespace() const override {
+    return RETINANET_PLUGIN_NAMESPACE;
+  }
+
+  void setPluginNamespace(const char *N) override {}
+
+  DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const
+  {
+    assert(index < 3);
+    return DataType::kFLOAT;
+  }
+
+  void configurePlugin(const DynamicPluginTensorDesc *in, int nbInputs, 
+    const DynamicPluginTensorDesc *out, int nbOutputs)
+  {
+    assert(nbInputs == 2);
+    assert(nbOutputs == 3);
+    auto const& scores_dims = in[0].desc.dims;
+    auto const& boxes_dims = in[1].desc.dims;
+    assert(scores_dims.d[2] == boxes_dims.d[2]);
+    assert(scores_dims.d[3] == boxes_dims.d[3]);
+    _height = scores_dims.d[2];
+    _width = scores_dims.d[3];
+    _num_anchors = boxes_dims.d[1] / 6; 
+    _num_classes = scores_dims.d[1] / _num_anchors;
+  }
+
+  IPluginV2DynamicExt *clone() const override {
+    return new DecodeRotatePlugin(_score_thresh, _top_n, _anchors, _scale, _height, _width, 
+      _num_anchors, _num_classes);
+  }
+
+private:
+  template<typename T> void write(char*& buffer, const T& val) const {
+    *reinterpret_cast<T*>(buffer) = val;
+    buffer += sizeof(T);
+  }
+
+  template<typename T> void read(const char*& buffer, T& val) {
+    val = *reinterpret_cast<const T*>(buffer);
+    buffer += sizeof(T);
+  }
+};
+
+class DecodeRotatePluginCreator : public IPluginCreator {
+public:
+  DecodeRotatePluginCreator() {}
+
+  const char *getPluginName () const override {
+    return RETINANET_PLUGIN_NAME;
+  }
+
+  const char *getPluginVersion () const override {
+    return RETINANET_PLUGIN_VERSION;
+  }
+ 
+  const char *getPluginNamespace() const override {
+    return RETINANET_PLUGIN_NAMESPACE;
+  }
+
+  IPluginV2DynamicExt *deserializePlugin (const char *name, const void *serialData, size_t serialLength) override {
+    return new DecodeRotatePlugin(serialData, serialLength);
+  }
+
+  void setPluginNamespace(const char *N) override {}
+  const PluginFieldCollection *getFieldNames() override { return nullptr; }
+  IPluginV2DynamicExt *createPlugin (const char *name, const PluginFieldCollection *fc) override { return nullptr; }
+};
+
+REGISTER_TENSORRT_PLUGIN(DecodeRotatePluginCreator);
+
+}
+
+#undef RETINANET_PLUGIN_NAME
+#undef RETINANET_PLUGIN_VERSION
+#undef RETINANET_PLUGIN_NAMESPACE

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+
+#include <NvInfer.h>
+
+#include <vector>
+#include <cassert>
+
+#include "../cuda/nms.h"
+
+using namespace nvinfer1;
+
+#define RETINANET_PLUGIN_NAME "RetinaNetNMS"
+#define RETINANET_PLUGIN_VERSION "1"
+#define RETINANET_PLUGIN_NAMESPACE ""
+
+namespace odtk {
+
+class NMSPlugin : public IPluginV2DynamicExt {
+  float _nms_thresh;
+  int _detections_per_im;
+
+  size_t _count;
+  mutable int size = -1;
+
+protected:
+  void deserialize(void const* data, size_t length) {
+    const char* d = static_cast<const char*>(data);
+    read(d, _nms_thresh);
+    read(d, _detections_per_im);
+    read(d, _count);
+  }
+
+  size_t getSerializationSize() const override {
+    return sizeof(_nms_thresh) + sizeof(_detections_per_im)
+      + sizeof(_count);
+  }
+
+  void serialize(void *buffer) const override {
+    char* d = static_cast<char*>(buffer);
+    write(d, _nms_thresh);
+    write(d, _detections_per_im);
+    write(d, _count);
+  }
+
+public:
+  NMSPlugin(float nms_thresh, int detections_per_im)
+    : _nms_thresh(nms_thresh), _detections_per_im(detections_per_im) {
+    assert(nms_thresh > 0);
+    assert(detections_per_im > 0);
+  }
+
+  NMSPlugin(float nms_thresh, int detections_per_im, size_t count)
+    : _nms_thresh(nms_thresh), _detections_per_im(detections_per_im), _count(count) {
+    assert(nms_thresh > 0);
+    assert(detections_per_im > 0);
+    assert(count > 0);
+  }
+
+  NMSPlugin(void const* data, size_t length) {
+    this->deserialize(data, length);
+  }
+
+  const char *getPluginType() const override {
+    return RETINANET_PLUGIN_NAME;
+  }
+ 
+  const char *getPluginVersion() const override {
+    return RETINANET_PLUGIN_VERSION;
+  }
+  
+  int getNbOutputs() const override {
+    return 3;
+  }
+
+  DimsExprs getOutputDimensions(int outputIndex, const DimsExprs *inputs,
+    int nbInputs, IExprBuilder &exprBuilder) override 
+  {
+    DimsExprs output(inputs[0]);
+    output.d[1] = exprBuilder.constant(_detections_per_im * (outputIndex == 1 ? 4 : 1));
+    output.d[2] = exprBuilder.constant(1);
+    output.d[3] = exprBuilder.constant(1);
+    return output;
+  }
+
+  bool supportsFormatCombination(int pos, const PluginTensorDesc *inOut, 
+    int nbInputs, int nbOutputs) override
+  {
+    assert(nbInputs == 3);
+    assert(nbOutputs == 3);
+    assert(pos < 6);
+    return inOut[pos].type == DataType::kFLOAT && inOut[pos].format == nvinfer1::PluginFormat::kLINEAR;
+  }
+
+  int initialize() override { return 0; }
+
+  void terminate() override {}
+
+  size_t getWorkspaceSize(const PluginTensorDesc *inputs, 
+    int nbInputs, const PluginTensorDesc *outputs, int nbOutputs) const override 
+  {
+    if (size < 0) {
+      size = cuda::nms(inputs->dims.d[0], nullptr, nullptr, _count, 
+        _detections_per_im, _nms_thresh, 
+        nullptr, 0, nullptr);
+    }
+    return size;
+  }
+
+  int enqueue(const PluginTensorDesc *inputDesc, 
+    const PluginTensorDesc *outputDesc, const void *const *inputs, 
+    void *const *outputs, void *workspace, cudaStream_t stream) 
+  {
+    return cuda::nms(inputDesc->dims.d[0], inputs, outputs, _count, 
+      _detections_per_im, _nms_thresh,
+      workspace, getWorkspaceSize(inputDesc, 3, outputDesc, 3), stream);
+  }
+
+  void destroy() override {
+    delete this;
+  }
+
+  const char *getPluginNamespace() const override {
+    return RETINANET_PLUGIN_NAMESPACE;
+  }
+  
+  void setPluginNamespace(const char *N) override {}
+
+  DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const
+  {
+    assert(index < 3);
+    return DataType::kFLOAT;
+  }  
+  
+  void configurePlugin(const DynamicPluginTensorDesc *in, int nbInputs, 
+    const DynamicPluginTensorDesc *out, int nbOutputs)
+  {
+    assert(nbInputs == 3);
+    assert(in[0].desc.dims.d[1] == in[2].desc.dims.d[1]);
+    assert(in[1].desc.dims.d[1] == in[2].desc.dims.d[1] * 4);
+    _count = in[0].desc.dims.d[1];
+  }
+
+  IPluginV2DynamicExt *clone() const {
+    return new NMSPlugin(_nms_thresh, _detections_per_im, _count);
+  }
+
+
+private:
+  template<typename T> void write(char*& buffer, const T& val) const {
+    *reinterpret_cast<T*>(buffer) = val;
+    buffer += sizeof(T);
+  }
+
+  template<typename T> void read(const char*& buffer, T& val) {
+    val = *reinterpret_cast<const T*>(buffer);
+    buffer += sizeof(T);
+  }
+};
+
+class NMSPluginCreator : public IPluginCreator {
+public:
+  NMSPluginCreator() {}
+  
+  const char *getPluginNamespace() const override {
+    return RETINANET_PLUGIN_NAMESPACE;
+  }
+  const char *getPluginName () const override {
+    return RETINANET_PLUGIN_NAME;
+  }
+
+  const char *getPluginVersion () const override {
+    return RETINANET_PLUGIN_VERSION;
+  }
+ 
+  //Was IPluginV2
+  IPluginV2DynamicExt *deserializePlugin (const char *name, const void *serialData, size_t serialLength) override {
+    return new NMSPlugin(serialData, serialLength);
+  }
+  
+  //Was IPluginV2
+  void setPluginNamespace(const char *N) override {}
+  const PluginFieldCollection *getFieldNames() override { return nullptr; }
+  IPluginV2DynamicExt *createPlugin (const char *name, const PluginFieldCollection *fc) override { return nullptr; }
+};
+
+REGISTER_TENSORRT_PLUGIN(NMSPluginCreator);
+
+}
+
+#undef RETINANET_PLUGIN_NAME
+#undef RETINANET_PLUGIN_VERSION
+#undef RETINANET_PLUGIN_NAMESPACE
\ No newline at end of file

+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#pragma once
+
+#include <NvInfer.h>
+
+#include <vector>
+#include <cassert>
+
+#include "../cuda/nms_iou.h"
+
+using namespace nvinfer1;
+
+#define RETINANET_PLUGIN_NAME "RetinaNetNMSRotate"
+#define RETINANET_PLUGIN_VERSION "1"
+#define RETINANET_PLUGIN_NAMESPACE ""
+
+namespace odtk {
+
+class NMSRotatePlugin : public IPluginV2DynamicExt {
+  float _nms_thresh;
+  int _detections_per_im;
+
+  size_t _count;
+  mutable int size = -1;
+
+protected:
+  void deserialize(void const* data, size_t length) {
+    const char* d = static_cast<const char*>(data);
+    read(d, _nms_thresh);
+    read(d, _detections_per_im);
+    read(d, _count);
+  }
+
+  size_t getSerializationSize() const override {
+    return sizeof(_nms_thresh) + sizeof(_detections_per_im)
+      + sizeof(_count);
+  }
+
+  void serialize(void *buffer) const override {
+    char* d = static_cast<char*>(buffer);
+    write(d, _nms_thresh);
+    write(d, _detections_per_im);
+    write(d, _count);
+  }
+
+public:
+  NMSRotatePlugin(float nms_thresh, int detections_per_im)
+    : _nms_thresh(nms_thresh), _detections_per_im(detections_per_im) {
+    assert(nms_thresh > 0);
+    assert(detections_per_im > 0);
+  }
+
+  NMSRotatePlugin(float nms_thresh, int detections_per_im, size_t count)
+    : _nms_thresh(nms_thresh), _detections_per_im(detections_per_im), _count(count) {
+    assert(nms_thresh > 0);
+    assert(detections_per_im > 0);
+    assert(count > 0);
+  }
+
+  NMSRotatePlugin(void const* data, size_t length) {
+    this->deserialize(data, length);
+  }
+
+  const char *getPluginType() const override {
+    return RETINANET_PLUGIN_NAME;
+  }
+ 
+  const char *getPluginVersion() const override {
+    return RETINANET_PLUGIN_VERSION;
+  }
+  
+  int getNbOutputs() const override {
+    return 3;
+  }
+
+  DimsExprs getOutputDimensions(int outputIndex, const DimsExprs *inputs,
+    int nbInputs, IExprBuilder &exprBuilder) override 
+  {
+    DimsExprs output(inputs[0]);
+    output.d[1] = exprBuilder.constant(_detections_per_im * (outputIndex == 1 ? 6 : 1));
+    output.d[2] = exprBuilder.constant(1);
+    output.d[3] = exprBuilder.constant(1);
+    return output;
+  }
+
+  bool supportsFormatCombination(int pos, const PluginTensorDesc *inOut, 
+    int nbInputs, int nbOutputs) override
+  {
+    assert(nbInputs == 3);
+    assert(nbOutputs == 3);
+    assert(pos < 6);
+    return inOut[pos].type == DataType::kFLOAT && inOut[pos].format == nvinfer1::PluginFormat::kLINEAR;
+  }
+
+  int initialize() override { return 0; }
+
+  void terminate() override {}
+
+  size_t getWorkspaceSize(const PluginTensorDesc *inputs, 
+    int nbInputs, const PluginTensorDesc *outputs, int nbOutputs) const override 
+  {
+    if (size < 0) {
+      size = cuda::nms_rotate(inputs->dims.d[0], nullptr, nullptr, _count, 
+        _detections_per_im, _nms_thresh, 
+        nullptr, 0, nullptr);
+    }
+    return size;
+  }
+
+  int enqueue(const PluginTensorDesc *inputDesc, 
+    const PluginTensorDesc *outputDesc, const void *const *inputs, 
+    void *const *outputs, void *workspace, cudaStream_t stream) override 
+  {
+    return cuda::nms_rotate(inputDesc->dims.d[0], inputs, outputs, _count, 
+      _detections_per_im, _nms_thresh,
+      workspace, getWorkspaceSize(inputDesc, 3, outputDesc, 3), stream);
+  }
+
+  void destroy() override {
+    delete this;
+  }
+
+  const char *getPluginNamespace() const override {
+    return RETINANET_PLUGIN_NAMESPACE;
+  }
+  
+  void setPluginNamespace(const char *N) override {}
+
+  DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const
+  {
+    assert(index < 3);
+    return DataType::kFLOAT;
+  }
+
+
+  void configurePlugin(const DynamicPluginTensorDesc *in, int nbInputs, 
+    const DynamicPluginTensorDesc *out, int nbOutputs)
+  {
+    assert(nbInputs == 3);
+    assert(in[0].desc.dims.d[1] == in[2].desc.dims.d[1]);
+    assert(in[1].desc.dims.d[1] == in[2].desc.dims.d[1] * 6);
+    _count = in[0].desc.dims.d[1];
+  }
+
+  IPluginV2DynamicExt *clone() const {
+    return new NMSRotatePlugin(_nms_thresh, _detections_per_im, _count);
+  }
+
+private:
+  template<typename T> void write(char*& buffer, const T& val) const {
+    *reinterpret_cast<T*>(buffer) = val;
+    buffer += sizeof(T);
+  }
+
+  template<typename T> void read(const char*& buffer, T& val) {
+    val = *reinterpret_cast<const T*>(buffer);
+    buffer += sizeof(T);
+  }
+};
+
+class NMSRotatePluginCreator : public IPluginCreator {
+public:
+  NMSRotatePluginCreator() {}
+  
+  const char *getPluginNamespace() const override {
+    return RETINANET_PLUGIN_NAMESPACE;
+  }
+  const char *getPluginName () const override {
+    return RETINANET_PLUGIN_NAME;
+  }
+
+  const char *getPluginVersion () const override {
+    return RETINANET_PLUGIN_VERSION;
+  }
+ 
+  IPluginV2DynamicExt *deserializePlugin (const char *name, const void *serialData, size_t serialLength) override {
+    return new NMSRotatePlugin(serialData, serialLength);
+  }
+
+  void setPluginNamespace(const char *N) override {}
+  const PluginFieldCollection *getFieldNames() override { return nullptr; }
+  IPluginV2DynamicExt *createPlugin (const char *name, const PluginFieldCollection *fc) override { return nullptr; }
+};
+
+REGISTER_TENSORRT_PLUGIN(NMSRotatePluginCreator);
+
+}
+
+#undef RETINANET_PLUGIN_NAME
+#undef RETINANET_PLUGIN_VERSION
+#undef RETINANET_PLUGIN_NAMESPACE

+cmake_minimum_required(VERSION 3.9 FATAL_ERROR)
+
+project(odtk_infer LANGUAGES CXX)
+set(CMAKE_CXX_STANDARD 14)
+find_package(CUDA REQUIRED)
+enable_language(CUDA)
+find_package(OpenCV REQUIRED)
+
+if(DEFINED TensorRT_DIR)
+  include_directories("${TensorRT_DIR}/include")
+  link_directories("${TensorRT_DIR}/lib")
+endif(DEFINED TensorRT_DIR)
+include_directories(${CUDA_INCLUDE_DIRS})
+
+add_library(odtk SHARED
+  ../../csrc/cuda/decode.h
+  ../../csrc/cuda/decode.cu
+  ../../csrc/cuda/nms.h
+  ../../csrc/cuda/nms.cu
+  ../../csrc/cuda/decode_rotate.h
+  ../../csrc/cuda/decode_rotate.cu
+  ../../csrc/cuda/nms_iou.h
+  ../../csrc/cuda/nms_iou.cu
+  ../../csrc/cuda/utils.h
+  ../../csrc/engine.h
+  ../../csrc/engine.cpp
+  ../../csrc/calibrator.h
+)
+set_target_properties(odtk PROPERTIES
+  CUDA_RESOLVE_DEVICE_SYMBOLS ON
+  CUDA_ARCHITECTURES 60 61 70 72 75 80 86
+)
+include_directories(${OpenCV_INCLUDE_DIRS})
+target_link_libraries(odtk PUBLIC nvinfer nvonnxparser ${OpenCV_LIBS})
+
+add_executable(export export.cpp)
+include_directories(${OpenCV_INCLUDE_DIRS})
+target_link_libraries(export PRIVATE odtk ${OpenCV_LIBS})
+
+add_executable(infer infer.cpp)
+include_directories(${OpenCV_INCLUDE_DIRS})
+target_link_libraries(infer PRIVATE odtk ${OpenCV_LIBS} cuda ${CUDA_LIBRARIES})
+
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
+  add_executable(infervideo infervideo.cpp)
+  include_directories(${OpenCV_INCLUDE_DIRS})
+  target_link_libraries(infervideo PRIVATE odtk ${OpenCV_LIBS} cuda ${CUDA_LIBRARIES})
+endif()

+# RetinaNet C++ Inference API - Sample Code
+
+The C++ API allows you to build a TensorRT engine for inference using the ONNX export of a core model.
+
+The following shows how to build and run code samples for exporting an ONNX core model (from RetinaNet or other toolkit supporting the same sort of core model structure) to a TensorRT engine and doing inference on images.
+
+## Building
+
+Building the example requires the following toolkits and libraries to be set up properly on your system:
+* A proper C++ toolchain (MSVS on Windows)
+* [CMake](https://cmake.org/download/) version 3.9 or later
+* NVIDIA [CUDA](https://developer.nvidia.com/cuda-toolkit)
+* NVIDIA [CuDNN](https://developer.nvidia.com/cudnn)
+* NVIDIA [TensorRT](https://developer.nvidia.com/tensorrt)
+* [OpenCV](https://opencv.org/releases.html)
+
+### Linux
+```bash
+mkdir build && cd build
+cmake -DCMAKE_CUDA_FLAGS="--expt-extended-lambda -std=c++14" ..
+make
+```
+
+### Windows
+```bash
+mkdir build && cd build
+cmake -G "Visual Studio 15 2017" -A x64 -T host=x64,cuda=10.0 -DTensorRT_DIR="C:\path\to\tensorrt" -DOpenCV_DIR="C:\path\to\opencv\build" ..
+msbuild odtk_infer.sln
+```
+
+## Running
+
+If you don't have an ONNX core model, generate one from your RetinaNet model:
+```bash
+odtk export model.pth model.onnx
+```
+
+Load the ONNX core model and export it to a RetinaNet TensorRT engine (using FP16 precision):
+```bash
+export{.exe} model.onnx engine.plan
+```
+
+You can also export the ONNX core model to an INT8 TensorRT engine if you have already done INT8 calibration:
+```bash
+export{.exe} model.onnx engine.plan INT8CalibrationTable
+```
+
+Run a test inference (default output if none provided: "detections.png"):
+```bash
+infer{.exe} engine.plan image.jpg [<OUTPUT>.png]
+```
+
+Note: make sure the TensorRT, CuDNN and OpenCV libraries are available in your environment and path.
+
+We have verified these steps with the following configurations:
+* DGX-1V using the provided Docker container (CUDA 10, cuDNN 7.4.2, TensorRT 5.0.2, OpenCV 3.4.3)
+* Jetson AGX Xavier with JetPack 4.1.1 Developer Preview (CUDA 10, cuDNN 7.3.1, TensorRT 5.0.3, OpenCV 3.3.1)
+
+
+
+

+import numpy as np
+from odtk.box import generate_anchors, generate_anchors_rotated
+
+# Generates anchors for export.cpp
+
+# ratios = [1.0, 2.0, 0.5]
+# scales = [4 * 2 ** (i / 3) for i in range(3)]
+ratios = [0.25, 0.5, 1.0, 2.0, 4.0]
+scales = [2 * 2**(2 * i/3) for i in range(3)]
+angles = [-np.pi / 6, 0, np.pi / 6]
+strides = [2**i for i in range(3,8)]
+
+axis = str(np.round([generate_anchors(stride, ratios, scales, 
+            angles).view(-1).tolist() for stride in strides], decimals=2).tolist()
+        ).replace('[', '{').replace(']', '}').replace('}, ', '},\n')
+
+rot = str(np.round([generate_anchors_rotated(stride, ratios, scales, 
+            angles)[0].view(-1).tolist() for stride in strides], decimals=2).tolist()
+        ).replace('[', '{').replace(']', '}').replace('}, ', '},\n')
+
+print("Axis-aligned:\n"+axis+'\n')
+print("Rotated:\n"+rot)

+#include <iostream>
+#include <stdexcept>
+#include <fstream>
+#include <vector>
+#include <chrono>
+
+#include <opencv2/opencv.hpp>
+#include <opencv2/opencv.hpp>
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+
+#include <cuda_runtime.h>
+
+#include "../../csrc/engine.h"
+
+using namespace std;
+using namespace cv;
+
+int main(int argc, char *argv[]) {
+	if (argc<3 || argc>4) {
+		cerr << "Usage: " << argv[0] << " engine.plan image.jpg [<OUTPUT>.png]" << endl;
+		return 1;
+	}
+
+	cout << "Loading engine..." << endl;
+	auto engine = odtk::Engine(argv[1]);
+
+	cout << "Preparing data..." << endl;
+	auto image = imread(argv[2], IMREAD_COLOR);
+	auto inputSize = engine.getInputSize();
+	cv::resize(image, image, Size(inputSize[1], inputSize[0]));
+	cv::Mat pixels;
+	image.convertTo(pixels, CV_32FC3, 1.0 / 255, 0);
+
+	int channels = 3;
+	vector<float> img;
+	vector<float> data (channels * inputSize[0] * inputSize[1]);
+
+	if (pixels.isContinuous())
+		img.assign((float*)pixels.datastart, (float*)pixels.dataend);
+	else {
+		cerr << "Error reading image " << argv[2] << endl;
+		return -1;
+	}
+
+	vector<float> mean {0.485, 0.456, 0.406};
+	vector<float> std {0.229, 0.224, 0.225};
+
+	for (int c = 0; c < channels; c++) {
+		for (int j = 0, hw = inputSize[0] * inputSize[1]; j < hw; j++) {
+			data[c * hw + j] = (img[channels * j + 2 - c] - mean[c]) / std[c];
+		}
+	}        
+
+	// Create device buffers
+	void *data_d, *scores_d, *boxes_d, *classes_d;
+	auto num_det = engine.getMaxDetections();
+	cudaMalloc(&data_d, 3 * inputSize[0] * inputSize[1] * sizeof(float));
+	cudaMalloc(&scores_d, num_det * sizeof(float));
+	cudaMalloc(&boxes_d, num_det * 4 * sizeof(float));
+	cudaMalloc(&classes_d, num_det * sizeof(float));
+
+	// Copy image to device
+	size_t dataSize = data.size() * sizeof(float);
+	cudaMemcpy(data_d, data.data(), dataSize, cudaMemcpyHostToDevice);
+
+	// Run inference n times
+	cout << "Running inference..." << endl;
+	const int count = 100;
+	auto start = chrono::steady_clock::now();
+ 	vector<void *> buffers = { data_d, scores_d, boxes_d, classes_d };
+	for (int i = 0; i < count; i++) {
+		engine.infer(buffers, 1);
+	}
+	auto stop = chrono::steady_clock::now();
+	auto timing = chrono::duration_cast<chrono::duration<double>>(stop - start);
+	cout << "Took " << timing.count() / count << " seconds per inference." << endl;
+
+	cudaFree(data_d);
+
+	// Get back the bounding boxes
+	unique_ptr<float[]> scores(new float[num_det]);
+	unique_ptr<float[]> boxes(new float[num_det * 4]);
+	unique_ptr<float[]> classes(new float[num_det]);
+	cudaMemcpy(scores.get(), scores_d, sizeof(float) * num_det, cudaMemcpyDeviceToHost);
+	cudaMemcpy(boxes.get(), boxes_d, sizeof(float) * num_det * 4, cudaMemcpyDeviceToHost);
+	cudaMemcpy(classes.get(), classes_d, sizeof(float) * num_det, cudaMemcpyDeviceToHost);
+
+	cudaFree(scores_d);
+	cudaFree(boxes_d);
+	cudaFree(classes_d);
+
+	for (int i = 0; i < num_det; i++) {
+		// Show results over confidence threshold
+		if (scores[i] >= 0.3f) {
+			float x1 = boxes[i*4+0];
+			float y1 = boxes[i*4+1];
+			float x2 = boxes[i*4+2];
+			float y2 = boxes[i*4+3];
+			cout << "Found box {" << x1 << ", " << y1 << ", " << x2 << ", " << y2
+				<< "} with score " << scores[i] << " and class " << classes[i] << endl;
+
+			// Draw bounding box on image
+			cv::rectangle(image, Point(x1, y1), Point(x2, y2), cv::Scalar(0, 255, 0));
+		}
+	}
+
+	// Write image
+	string out_file = argc == 4 ? string(argv[3]) : "detections.png";
+	cout << "Saving result to " << out_file << endl;
+	imwrite(out_file, image);
+	
+	return 0;
+}

+#include <iostream>
+#include <stdexcept>
+#include <fstream>
+#include <vector>
+#include <chrono>
+
+#include <opencv2/opencv.hpp>
+#include <opencv2/opencv.hpp>
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+
+#include <cuda_runtime.h>
+
+#include "../../csrc/engine.h"
+
+using namespace std;
+using namespace cv;
+
+int main(int argc, char *argv[]) {
+	if (argc != 4) {
+		cerr << "Usage: " << argv[0] << " engine.plan input.mov output.mp4" << endl;
+		return 1;
+	}
+
+	cout << "Loading engine..." << endl;
+	auto engine = odtk::Engine(argv[1]);
+	VideoCapture src(argv[2]);
+
+	if (!src.isOpened()){
+		cerr << "Could not read " << argv[2] << endl;
+		return 1;
+	}
+	
+	auto fh=src.get(CAP_PROP_FRAME_HEIGHT);
+	auto fw=src.get(CAP_PROP_FRAME_WIDTH);
+	auto fps=src.get(CAP_PROP_FPS);
+	auto nframes=src.get(CAP_PROP_FRAME_COUNT);
+
+	VideoWriter sink;
+	sink.open(argv[3], 0x31637661, fps, Size(fw, fh));
+	Mat frame;
+	Mat resized_frame;
+	Mat inferred_frame;
+	int count=1;
+
+	auto inputSize = engine.getInputSize();
+	// Create device buffers
+	void *data_d, *scores_d, *boxes_d, *classes_d;
+	auto num_det = engine.getMaxDetections();
+	cudaMalloc(&data_d, 3 * inputSize[0] * inputSize[1] * sizeof(float));
+	cudaMalloc(&scores_d, num_det * sizeof(float));
+	cudaMalloc(&boxes_d, num_det * 4 * sizeof(float));
+	cudaMalloc(&classes_d, num_det * sizeof(float));
+
+	unique_ptr<float[]> scores(new float[num_det]);
+	unique_ptr<float[]> boxes(new float[num_det * 4]);
+	unique_ptr<float[]> classes(new float[num_det]);
+
+	vector<float> mean {0.485, 0.456, 0.406};
+	vector<float> std {0.229, 0.224, 0.225};
+
+	vector<uint8_t> blues {0,63,127,191,255,0}; //colors for bonuding boxes
+	vector<uint8_t> greens {0,255,191,127,63,0};
+	vector<uint8_t> reds {191,255,0,0,63,127};
+
+	int channels = 3;
+	vector<float> img;
+	vector<float> data (channels * inputSize[0] * inputSize[1]);
+
+	while(1)
+	{
+		src >> frame;
+		if (frame.empty()){
+			cout << "Finished inference!" << endl;
+			break;
+		}
+
+		cv::resize(frame, resized_frame, Size(inputSize[1], inputSize[0]));
+		cv::Mat pixels;
+		resized_frame.convertTo(pixels, CV_32FC3, 1.0 / 255, 0);
+
+		img.assign((float*)pixels.datastart, (float*)pixels.dataend);
+
+		for (int c = 0; c < channels; c++) {
+			for (int j = 0, hw = inputSize[0] * inputSize[1]; j < hw; j++) {
+				data[c * hw + j] = (img[channels * j + 2 - c] - mean[c]) / std[c];
+			}
+		}
+
+		// Copy image to device
+		size_t dataSize = data.size() * sizeof(float);
+		cudaMemcpy(data_d, data.data(), dataSize, cudaMemcpyHostToDevice);
+
+		//Do inference
+		cout << "Inferring on frame: " << count <<"/" << nframes << endl;
+		count++;
+		vector<void *> buffers = { data_d, scores_d, boxes_d, classes_d };
+		engine.infer(buffers, 1);
+
+		cudaMemcpy(scores.get(), scores_d, sizeof(float) * num_det, cudaMemcpyDeviceToHost);
+		cudaMemcpy(boxes.get(), boxes_d, sizeof(float) * num_det * 4, cudaMemcpyDeviceToHost);
+		cudaMemcpy(classes.get(), classes_d, sizeof(float) * num_det, cudaMemcpyDeviceToHost);
+
+		// Get back the bounding boxes
+		for (int i = 0; i < num_det; i++) {
+			// Show results over confidence threshold
+			if (scores[i] >= 0.2f) {
+				float x1 = boxes[i*4+0];
+				float y1 = boxes[i*4+1];
+				float x2 = boxes[i*4+2];
+				float y2 = boxes[i*4+3];
+				int cls=classes[i];
+				// Draw bounding box on image
+				cv::rectangle(resized_frame, Point(x1, y1), Point(x2, y2), cv::Scalar(blues[cls], greens[cls], reds[cls]));
+			}
+		}
+		cv::resize(resized_frame, inferred_frame, Size(fw, fh));
+		sink.write(inferred_frame);
+	}
+	src.release();
+	sink.release();
+	cudaFree(data_d);
+	cudaFree(scores_d);
+	cudaFree(boxes_d);
+	cudaFree(classes_d);
+	return 0;
+}

+# Deploying RetinaNet in DeepStream 4.0
+
+This shows how to export a trained RetinaNet model to TensorRT and deploy it in a video analytics application using NVIDIA DeepStream 4.0.
+
+## Prerequisites
+* A GPU supported by DeepStream: Jetson Xavier, Tesla P4/P40/V100/T4
+* A trained PyTorch RetinaNet model.
+* A video source, either `.mp4` files or a webcam.
+
+## Tesla GPUs
+Setup instructions:
+
+#### 1. Download DeepStream 4.0 
+Download DeepStream 4.0 SDK for Tesla "Download .tar" from [https://developer.nvidia.com/deepstream-download](https://developer.nvidia.com/deepstream-download) and place in the `extras/deepstream` directory. 
+
+This file should be called `deepstream_sdk_v4.0.2_x86_64.tbz2`.
+
+#### 2. Unpack DeepStream
+You may need to adjust the permissions on the `.tbz2` file before you can extract it. 
+
+```
+cd extras/deepstream
+mkdir DeepStream_Release
+tar -xvf deepstream_sdk_v4.0.2_x86_64.tbz2 -C DeepStream_Release/
+```
+
+#### 3. Build and enter the DeepStream docker container
+```
+docker build -f <your_path>/retinanet-examples/Dockerfile.deepstream -t ds_odtk:latest <your_path>/retinanet-examples
+docker run --gpus all -it --rm --ipc=host -v <dir containing your data>:/data ds_odtk:latest
+```
+
+#### 4. Export your trained PyTorch RetinaNet model to TensorRT per the [INFERENCE](https://github.com/NVIDIA/retinanet-examples/blob/master/INFERENCE.md) instructions:
+```
+odtk export <PyTorch model> <engine> --batch n
+
+OR
+
+odtk export <PyTorch model> <engine> --int8 --calibration-images <example images> --batch n
+```
+
+#### 5. Run deepstream-app
+Once all of the config files have been modified, launch the DeepStream application: 
+```
+cd /workspace/retinanet-examples/extras/deepstream/deepstream-sample/
+LD_PRELOAD=build/libnvdsparsebbox_odtk.so deepstream-app -c <config file>
+```
+
+## Jetson AGX Xavier
+Setup instructions.
+
+#### 1. Flash Jetson Xavier with [Jetpack 4.3](https://developer.nvidia.com/embedded/jetpack)
+
+**Ensure that you tick the DeepStream box, under Additional SDKs**
+
+#### 2. (on host) Covert PyTorch model to ONNX.
+
+```bash
+odtk export model.pth model.onnx
+```
+
+#### 3. Copy ONNX RetinaNet model and config files to Jetson Xavier
+
+Use `scp` or a memory card.
+
+#### 4. (on Jetson) Make the C++ API
+
+```bash
+cd extras/cppapi
+mkdir build && cd build
+cmake -DCMAKE_CUDA_FLAGS="--expt-extended-lambda -std=c++14" ..
+make
+```
+
+#### 5. (on Jetson) Make the RetinaNet plugin
+
+```bash
+cd extras/deepstream/deepstream-sample
+mkdir build && cd build
+cmake -DDeepStream_DIR=/opt/nvidia/deepstream/deepstream-4.0 .. && make -j
+```
+
+#### 6. (on Jetson) Build the TensorRT Engine
+
+```bash
+cd extras/cppapi/build
+./export model.onnx engine.plan
+```
+
+#### 7. (on Jetson) Modify the DeepStream config files
+As described in the "preparing the DeepStream config file" section below. 
+
+#### 8. (on Jetson) Run deepstream-app
+Once all of the config files have been modified, launch the DeepStream application: 
+```
+cd extras/deepstream/deepstream-sample
+LD_PRELOAD=build/libnvdsparsebbox_odtk.so deepstream-app -c <config file>
+```
+
+## Preparing the DeepStream config file:
+We have included two example DeepStream config files in `deepstream-sample`.
+- `ds_config_1vids.txt`: Performs detection on a single video, using the detector specified by `infer_config_batch1.txt`.
+- `ds_config_8vids.txt`: Performs detection on multiple video streams simultaneously, using the detector specified by `infer_config_batch8.txt`. Frames from each video are combined into a single batch and passed to the detector for inference.
+
+The `ds_config_*` files are DeepStream config files. They describe the overall processing. `infer_config_*` files define the individual detectors, which can be chained in series.
+
+Before they can be used, these config files must be modified to specify the correct paths to the input and output videos files, and the TensorRT engines.
+
+* **Input files** are specified in the deepstream config files by the `uri=file://<path>` parameter.
+
+* **Output files** are specified in the deepstream config files by the `output-file=<path>` parameter.
+
+* **TensorRT engines** are specified in both the DeepStream config files, and also the detector config files, by the `model-engine-file=<path>` parameters. 
+
+On Xavier, you can optionally set `enable=1` to `[sink1]` in `ds_config_*` files to display the processed video stream.
+
+
+## Convert output video file to mp4
+You can convert the outputted `.mkv` file to `.mp4` using `ffmpeg`.
+```
+ffmpeg -i /data/output/file1.mkv -c copy /data/output/file2.mp4
+```

+cmake_minimum_required(VERSION 3.5.1)
+
+project(deepstream-odtk)
+enable_language(CXX)
+include(FindCUDA)
+
+set(CMAKE_CXX_STANDARD 14)
+find_package(CUDA REQUIRED)
+find_package(OpenCV REQUIRED)
+
+if(DEFINED TensorRT_DIR)
+  include_directories("${TensorRT_DIR}/include")
+  link_directories("${TensorRT_DIR}/lib")
+endif(DEFINED TensorRT_DIR)
+if(DEFINED DeepStream_DIR)
+  include_directories("${DeepStream_DIR}/sources/includes")
+endif(DEFINED DeepStream_DIR)
+include_directories(${CUDA_INCLUDE_DIRS})
+
+if(NOT DEFINED ARCH)
+  set(ARCH "sm_70")
+endif(NOT DEFINED ARCH)
+
+cuda_add_library(nvdsparsebbox_odtk SHARED
+  ../../../csrc/cuda/decode.h
+  ../../../csrc/cuda/decode.cu
+  ../../../csrc/cuda/nms.h
+  ../../../csrc/cuda/nms.cu
+  ../../../csrc/cuda/utils.h
+  ../../../csrc/engine.cpp
+  nvdsparsebbox_odtk.cpp
+  OPTIONS -arch ${ARCH} -std=c++14 --expt-extended-lambda
+)
+include_directories(${OpenCV_INCLUDE_DIRS})
+target_link_libraries(nvdsparsebbox_odtk ${CUDA_LIBRARIES} nvinfer nvinfer_plugin nvonnxparser ${OpenCV_LIBS})

+# Copyright (c) 2018 NVIDIA Corporation.  All rights reserved.
+#
+# NVIDIA Corporation and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA Corporation is strictly prohibited.
+
+[application]
+enable-perf-measurement=1
+perf-measurement-interval-sec=1
+
+[tiled-display]
+enable=0
+rows=1
+columns=1
+width=1280
+height=720
+gpu-id=0
+
+[source0]
+enable=1
+type=2
+num-sources=1
+uri=file://<path>
+gpu-id=0
+
+[streammux]
+gpu-id=0
+batch-size=1
+batched-push-timeout=-1
+## Set muxer output width and height
+width=1280
+height=720
+cuda-memory-type=1
+enable-padding=1
+
+[sink0]
+enable=1
+type=3
+#1=mp4 2=mkv
+container=1
+#1=h264 2=h265 3=mpeg4
+## only SW mpeg4 is supported right now.
+codec=3
+sync=1
+bitrate=80000000
+output-file=<path>
+source-id=0
+
+[sink1]
+enable=0
+#Type - 1=FakeSink 2=EglSink 3=File
+type=2
+sync=1
+source-id=0
+gpu-id=0
+cuda-memory-type=1
+
+
+[osd]
+enable=1
+gpu-id=0
+border-width=2
+text-size=12
+text-color=1;1;1;1;
+text-bg-color=0.3;0.3;0.3;1
+font=Arial
+show-clock=0
+clock-x-offset=800
+clock-y-offset=820
+clock-text-size=12
+clock-color=1;0;0;0
+
+[primary-gie]
+enable=1
+gpu-id=0
+batch-size=1
+gie-unique-id=1
+interval=0
+labelfile-path=labels_coco.txt
+model-engine-file=<path>
+config-file=infer_config_batch1.txt

+# Copyright (c) 2018 NVIDIA Corporation.  All rights reserved.
+#
+# NVIDIA Corporation and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA Corporation is strictly prohibited.
+
+[application]
+enable-perf-measurement=1
+perf-measurement-interval-sec=5
+
+[tiled-display]
+enable=1
+rows=2
+columns=4
+width=1280
+height=720
+gpu-id=0
+cuda-memory-type=1
+
+[source0]
+enable=1
+type=3
+num-sources=4
+uri=file://<path>
+gpu-id=0
+cuda-memory-type=1
+
+[source1]
+enable=1
+type=3
+num-sources=4
+uri=file://<path>
+gpu-id=0
+cuda-memory-type=1
+
+[streammux]
+gpu-id=0
+batched-push-timeout=-1
+## Set muxer output width and height
+width=1280
+height=720
+cuda-memory-type=1
+enable-padding=1
+batch-size=8
+
+[sink0]
+enable=1
+type=3
+#1=mp4 2=mkv
+container=1
+#1=h264 2=h265 3=mpeg4
+## only SW mpeg4 is supported right now.
+codec=3
+sync=0
+bitrate=32000000
+output-file=<path>
+source-id=0
+cuda-memory-type=1
+
+[sink1]
+enable=0
+#Type - 1=FakeSink 2=EglSink 3=File
+type=2
+sync=1
+source-id=0
+gpu-id=0
+cuda-memory-type=1
+
+
+[osd]
+enable=1
+gpu-id=0
+border-width=2
+text-size=12
+text-color=1;1;1;1;
+text-bg-color=0.3;0.3;0.3;1
+font=Arial
+show-clock=0
+clock-x-offset=800
+clock-y-offset=820
+clock-text-size=12
+clock-color=1;0;0;0
+
+[primary-gie]
+enable=1
+gpu-id=0
+batch-size=8
+gie-unique-id=1
+interval=0
+labelfile-path=labels_coco.txt
+model-engine-file=<path>
+config-file=infer_config_batch8.txt
+cuda-memory-type=1

+# Copyright (c) 2018 NVIDIA Corporation.  All rights reserved.
+# NVIDIA Corporation and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA Corporation is strictly prohibited.
+
+# Following properties are mandatory when engine files are not specified:
+#   int8-calib-file(Only in INT8)
+#   Caffemodel mandatory properties: model-file, proto-file, output-blob-names
+#   UFF: uff-file, input-dims, uff-input-blob-name, output-blob-names
+#   ONNX: onnx-file
+#
+# Mandatory properties for detectors:
+#   parse-func, num-detected-classes,
+#   custom-lib-path (when parse-func=0 i.e. custom),
+#   parse-bbox-func-name (when parse-func=0)
+#
+# Optional properties for detectors:
+#   enable-dbscan(Default=false), interval(Primary mode only, Default=0)
+#
+# Mandatory properties for classifiers:
+#   classifier-threshold, is-classifier
+#
+# Optional properties for classifiers:
+#   classifier-async-mode(Secondary mode only, Default=false)
+#
+# Optional properties in secondary mode:
+#   operate-on-gie-id(Default=0), operate-on-class-ids(Defaults to all classes),
+#   input-object-min-width, input-object-min-height, input-object-max-width,
+#   input-object-max-height
+#
+# Following properties are always recommended:
+#   batch-size(Default=1)
+#
+# Other optional properties:
+#   net-scale-factor(Default=1), network-mode(Default=0 i.e FP32),
+#   model-color-format(Default=0 i.e. RGB) model-engine-file, labelfile-path,
+#   mean-file, gie-unique-id(Default=0), offsets, gie-mode (Default=1 i.e. primary),
+#   custom-lib-path, network-mode(Default=0 i.e FP32)
+#
+# The values in the config file are overridden by values set through GObject
+# properties.
+
+[property]
+gpu-id=0
+net-scale-factor=0.017352074
+offsets=123.675;116.28;103.53
+model-engine-file=<path>
+labelfile-path=labels_coco.txt
+batch-size=1
+## 0=FP32, 1=INT8, 2=FP16 mode
+network-mode=2
+num-detected-classes=80
+interval=0
+gie-unique-id=1
+parse-func=0
+is-classifier=0
+output-blob-names=boxes;scores;classes
+parse-bbox-func-name=NvDsInferParseRetinaNet
+custom-lib-path=build/libnvdsparsebbox_odtk.so
+#enable-dbscan=1
+
+
+[class-attrs-all]
+threshold=0.5
+group-threshold=0
+## Set eps=0.7 and minBoxes for enable-dbscan=1
+#eps=0.2
+##minBoxes=3
+#roi-top-offset=0
+#roi-bottom-offset=0
+detected-min-w=4
+detected-min-h=4
+#detected-max-w=0
+#detected-max-h=0
+
+## Per class configuration
+#[class-attrs-2]
+#threshold=0.6
+#eps=0.5
+#group-threshold=3
+#roi-top-offset=20
+#roi-bottom-offset=10
+#detected-min-w=40
+#detected-min-h=40
+#detected-max-w=400
+#detected-max-h=800

+# Copyright (c) 2018 NVIDIA Corporation.  All rights reserved.
+# NVIDIA Corporation and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA Corporation is strictly prohibited.
+
+# Following properties are mandatory when engine files are not specified:
+#   int8-calib-file(Only in INT8)
+#   Caffemodel mandatory properties: model-file, proto-file, output-blob-names
+#   UFF: uff-file, input-dims, uff-input-blob-name, output-blob-names
+#   ONNX: onnx-file
+#
+# Mandatory properties for detectors:
+#   parse-func, num-detected-classes,
+#   custom-lib-path (when parse-func=0 i.e. custom),
+#   parse-bbox-func-name (when parse-func=0)
+#
+# Optional properties for detectors:
+#   enable-dbscan(Default=false), interval(Primary mode only, Default=0)
+#
+# Mandatory properties for classifiers:
+#   classifier-threshold, is-classifier
+#
+# Optional properties for classifiers:
+#   classifier-async-mode(Secondary mode only, Default=false)
+#
+# Optional properties in secondary mode:
+#   operate-on-gie-id(Default=0), operate-on-class-ids(Defaults to all classes),
+#   input-object-min-width, input-object-min-height, input-object-max-width,
+#   input-object-max-height
+#
+# Following properties are always recommended:
+#   batch-size(Default=1)
+#
+# Other optional properties:
+#   net-scale-factor(Default=1), network-mode(Default=0 i.e FP32),
+#   model-color-format(Default=0 i.e. RGB) model-engine-file, labelfile-path,
+#   mean-file, gie-unique-id(Default=0), offsets, gie-mode (Default=1 i.e. primary),
+#   custom-lib-path, network-mode(Default=0 i.e FP32)
+#
+# The values in the config file are overridden by values set through GObject
+# properties.
+
+[property]
+gpu-id=0
+net-scale-factor=0.017352074
+offsets=123.675;116.28;103.53
+model-engine-file=<path>
+labelfile-path=labels_coco.txt
+#int8-calib-file=cal_trt4.bin
+batch-size=8
+## 0=FP32, 1=INT8, 2=FP16 mode
+network-mode=2
+num-detected-classes=80
+interval=0
+gie-unique-id=1
+parse-func=0
+is-classifier=0
+output-blob-names=boxes;scores;classes
+parse-bbox-func-name=NvDsInferParseRetinaNet
+custom-lib-path=build/libnvdsparsebbox_odtk.so
+#enable-dbscan=1
+
+
+[class-attrs-all]
+threshold=0.5
+group-threshold=0
+## Set eps=0.7 and minBoxes for enable-dbscan=1
+#eps=0.2
+##minBoxes=3
+#roi-top-offset=0
+#roi-bottom-offset=0
+detected-min-w=4
+detected-min-h=4
+#detected-max-w=0
+#detected-max-h=0
+
+## Per class configuration
+#[class-attrs-2]
+#threshold=0.6
+#eps=0.5
+#group-threshold=3
+#roi-top-offset=20
+#roi-bottom-offset=10
+#detected-min-w=40
+#detected-min-h=40
+#detected-max-w=400
+#detected-max-h=800

+person
+bicycle
+car
+motorcycle
+airplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+couch
+potted plant
+bed
+dining table
+toilet
+tv
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush
\ No newline at end of file

+/**
+ * Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * NVIDIA Corporation and its licensors retain all intellectual property
+ * and proprietary rights in and to this software, related documentation
+ * and any modifications thereto.  Any use, reproduction, disclosure or
+ * distridbution of this software and related documentation without an express
+ * license agreement from NVIDIA Corporation is strictly prohibited.
+ *
+ */
+
+#include <cstring>
+#include <iostream>
+#include "nvdsinfer_custom_impl.h"
+
+#define MIN(a,b) ((a) < (b) ? (a) : (b))
+
+/* This is a sample bounding box parsing function for the sample Resnet10
+ * detector model provided with the SDK. */
+
+/* C-linkage to prevent name-mangling */
+extern "C"
+bool NvDsInferParseRetinaNet (std::vector<NvDsInferLayerInfo> const &outputLayersInfo,
+        NvDsInferNetworkInfo  const &networkInfo,
+        NvDsInferParseDetectionParams const &detectionParams,
+        std::vector<NvDsInferParseObjectInfo> &objectList)
+{
+  static int bboxLayerIndex = -1;
+  static int classesLayerIndex = -1;
+  static int scoresLayerIndex = -1;
+  static NvDsInferDimsCHW scoresLayerDims;
+  int numDetsToParse;
+
+  /* Find the bbox layer */
+  if (bboxLayerIndex == -1) {
+    for (unsigned int i = 0; i < outputLayersInfo.size(); i++) {
+      if (strcmp(outputLayersInfo[i].layerName, "boxes") == 0) {
+        bboxLayerIndex = i;
+        break;
+      }
+    }
+    if (bboxLayerIndex == -1) {
+    std::cerr << "Could not find bbox layer buffer while parsing" << std::endl;
+    return false;
+    }
+  }
+
+  /* Find the scores layer */
+  if (scoresLayerIndex == -1) {
+    for (unsigned int i = 0; i < outputLayersInfo.size(); i++) {
+      if (strcmp(outputLayersInfo[i].layerName, "scores") == 0) {
+        scoresLayerIndex = i;
+        getDimsCHWFromDims(scoresLayerDims, outputLayersInfo[i].dims);
+        break;
+      }
+    }
+    if (scoresLayerIndex == -1) {
+    std::cerr << "Could not find scores layer buffer while parsing" << std::endl;
+    return false;
+    }
+  }
+
+  /* Find the classes layer */
+  if (classesLayerIndex == -1) {
+    for (unsigned int i = 0; i < outputLayersInfo.size(); i++) {
+      if (strcmp(outputLayersInfo[i].layerName, "classes") == 0) {
+        classesLayerIndex = i;
+        break;
+      }
+    }
+    if (classesLayerIndex == -1) {
+    std::cerr << "Could not find classes layer buffer while parsing" << std::endl;
+    return false;
+    }
+  }  
+
+  
+  /* Calculate the number of detections to parse */
+  numDetsToParse = scoresLayerDims.c;
+
+  float *bboxes = (float *) outputLayersInfo[bboxLayerIndex].buffer;
+  float *classes = (float *) outputLayersInfo[classesLayerIndex].buffer;
+  float *scores = (float *) outputLayersInfo[scoresLayerIndex].buffer;
+  
+  for (int indx = 0; indx < numDetsToParse; indx++)
+  {
+    float outputX1 = bboxes[indx * 4];
+    float outputY1 = bboxes[indx * 4 + 1];
+    float outputX2 = bboxes[indx * 4 + 2];
+    float outputY2 = bboxes[indx * 4 + 3];
+    float this_class = classes[indx];
+    float this_score = scores[indx];
+    float threshold = detectionParams.perClassThreshold[this_class];
+    
+    if (this_score >= threshold)
+    {
+      NvDsInferParseObjectInfo object;
+      
+      object.classId = this_class;
+      object.detectionConfidence = this_score;
+
+      object.left = outputX1;
+      object.top = outputY1;
+      object.width = outputX2 - outputX1;
+      object.height = outputY2 - outputY1;
+
+      objectList.push_back(object);
+    }
+  }
+  return true;
+}
+
+/* Check that the custom function has been defined correctly */
+CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseRetinaNet);

+#!/usr/bin/env bash
+
+if [ $# -ne 2 ]; then
+    echo "Usage: $0 images_path annotations.json"
+    exit 1
+fi
+
+tmp="/tmp/odtk"
+
+tests=(
+    "odtk train ${tmp}/model.pth --images $1 --annotations $2 --max-size 640 --override --iters 100 --backbone ResNet18FPN ResNet50FPN"
+    "odtk train ${tmp}/model.pth --images $1 --annotations $2 --max-size 640 --override --iters 100"
+    "odtk train ${tmp}/model.pth --fine-tune ${tmp}/model.pth --images $1 --annotations $2 --max-size 640 --override --iters 100"
+    "odtk infer ${tmp}/model.pth --images ${tmp}/test_images --max-size 640"
+    "odtk export ${tmp}/model.pth ${tmp}/engine.plan --size 640"
+    "odtk infer ${tmp}/engine.plan --images ${tmp}/test_images --max-size 640"
+)
+
+start=`date +%s`
+
+# Prepare small image folder for inference
+if [ ! -d ${tmp}/test_images ]; then
+    mkdir -p ${tmp}/test_images
+    cp $(find $1 | tail -n 10) ${tmp}/test_images
+fi
+
+# Run all tests
+for test in "${tests[@]}"; do
+    echo "Running \"${test}\""
+    ${test}
+    if [ $? -ne 0 ]; then
+        echo "Test failed!"
+        exit 1
+    fi
+done
+
+end=`date +%s`
+
+echo "All test succeeded in $((end-start)) seconds!"
\ No newline at end of file

+import sys
+
+from .resnet import *
+from .mobilenet import *
+from .fpn import *

+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import resnet as vrn
+from torchvision.models import mobilenet as vmn
+
+from .resnet import ResNet
+from .mobilenet import MobileNet
+from .utils import register
+
+
+class FPN(nn.Module):
+    'Feature Pyramid Network - https://arxiv.org/abs/1612.03144'
+
+    def __init__(self, features):
+        super().__init__()
+
+        self.stride = 128
+        self.features = features
+
+        if isinstance(features, ResNet):
+            is_light = features.bottleneck == vrn.BasicBlock
+            channels = [128, 256, 512] if is_light else [512, 1024, 2048]
+        elif isinstance(features, MobileNet):
+            channels = [32, 96, 320]
+
+        self.lateral3 = nn.Conv2d(channels[0], 256, 1)
+        self.lateral4 = nn.Conv2d(channels[1], 256, 1)
+        self.lateral5 = nn.Conv2d(channels[2], 256, 1)
+        self.pyramid6 = nn.Conv2d(channels[2], 256, 3, stride=2, padding=1)
+        self.pyramid7 = nn.Conv2d(256, 256, 3, stride=2, padding=1)
+        self.smooth3 = nn.Conv2d(256, 256, 3, padding=1)
+        self.smooth4 = nn.Conv2d(256, 256, 3, padding=1)
+        self.smooth5 = nn.Conv2d(256, 256, 3, padding=1)
+
+    def initialize(self):
+        def init_layer(layer):
+            if isinstance(layer, nn.Conv2d):
+                nn.init.xavier_uniform_(layer.weight)
+                if layer.bias is not None:
+                    nn.init.constant_(layer.bias, val=0)
+        self.apply(init_layer)
+
+        self.features.initialize()
+
+    def forward(self, x):
+        c3, c4, c5 = self.features(x)
+
+        p5 = self.lateral5(c5)
+        p4 = self.lateral4(c4)
+        p4 = F.interpolate(p5, scale_factor=2) + p4
+        p3 = self.lateral3(c3)
+        p3 = F.interpolate(p4, scale_factor=2) + p3
+
+        p6 = self.pyramid6(c5)
+        p7 = self.pyramid7(F.relu(p6))
+
+        p3 = self.smooth3(p3)
+        p4 = self.smooth4(p4)
+        p5 = self.smooth5(p5)
+
+        return [p3, p4, p5, p6, p7]
+
+@register
+def ResNet18FPN():
+    return FPN(ResNet(layers=[2, 2, 2, 2], bottleneck=vrn.BasicBlock, outputs=[3, 4, 5], url=vrn.model_urls['resnet18']))
+
+@register
+def ResNet34FPN():
+    return FPN(ResNet(layers=[3, 4, 6, 3], bottleneck=vrn.BasicBlock, outputs=[3, 4, 5], url=vrn.model_urls['resnet34']))
+
+@register
+def ResNet50FPN():
+    return FPN(ResNet(layers=[3, 4, 6, 3], bottleneck=vrn.Bottleneck, outputs=[3, 4, 5], url=vrn.model_urls['resnet50']))
+
+@register
+def ResNet101FPN():
+    return FPN(ResNet(layers=[3, 4, 23, 3], bottleneck=vrn.Bottleneck, outputs=[3, 4, 5], url=vrn.model_urls['resnet101']))
+
+@register
+def ResNet152FPN():
+    return FPN(ResNet(layers=[3, 8, 36, 3], bottleneck=vrn.Bottleneck, outputs=[3, 4, 5], url=vrn.model_urls['resnet152']))
+
+@register
+def ResNeXt50_32x4dFPN():
+    return FPN(ResNet(layers=[3, 4, 6, 3], bottleneck=vrn.Bottleneck, outputs=[3, 4, 5], groups=32, width_per_group=4, url=vrn.model_urls['resnext50_32x4d']))
+
+@register
+def ResNeXt101_32x8dFPN():
+    return FPN(ResNet(layers=[3, 4, 23, 3], bottleneck=vrn.Bottleneck, outputs=[3, 4, 5], groups=32, width_per_group=8, url=vrn.model_urls['resnext101_32x8d']))
+
+@register
+def MobileNetV2FPN():
+    return FPN(MobileNet(outputs=[6, 13, 17], url=vmn.model_urls['mobilenet_v2']))

+import torch
+from torch import nn
+import torch.nn.functional as F
+
+class FixedBatchNorm2d(nn.Module):
+    'BatchNorm2d where the batch statistics and the affine parameters are fixed'
+
+    def __init__(self, n):
+        super().__init__()
+        self.register_buffer("weight", torch.ones(n))
+        self.register_buffer("bias", torch.zeros(n))
+        self.register_buffer("running_mean", torch.zeros(n))
+        self.register_buffer("running_var", torch.ones(n))
+
+    def forward(self, x):
+        return F.batch_norm(x, running_mean=self.running_mean, running_var=self.running_var, weight=self.weight, bias=self.bias)
+
+def convert_fixedbn_model(module):
+    'Convert batch norm layers to fixed'
+
+    mod = module
+    if isinstance(module, nn.BatchNorm2d):
+        mod = FixedBatchNorm2d(module.num_features)
+        mod.running_mean = module.running_mean
+        mod.running_var = module.running_var
+        if module.affine:
+            mod.weight.data = module.weight.data.clone().detach()
+            mod.bias.data = module.bias.data.clone().detach()
+    for name, child in module.named_children():
+        mod.add_module(name, convert_fixedbn_model(child))
+
+    return mod

+import torch.nn as nn
+from torchvision.models import mobilenet as vmn
+import torch.utils.model_zoo as model_zoo
+
+class MobileNet(vmn.MobileNetV2):
+    'MobileNetV2: Inverted Residuals and Linear Bottlenecks - https://arxiv.org/abs/1801.04381'
+
+    def __init__(self, outputs=[18], url=None):
+        self.stride = 128
+        self.url = url
+        super().__init__()
+        self.outputs = outputs
+        self.unused_modules = ['features.18', 'classifier']
+
+    def initialize(self):
+        if self.url:
+            self.load_state_dict(model_zoo.load_url(self.url))
+
+    def forward(self, x):
+        outputs = []
+        for indx, feat in enumerate(self.features[:-1]):
+            x = feat(x)
+            if indx in self.outputs:
+                outputs.append(x)
+        return outputs

+import torchvision
+from torchvision.models import resnet as vrn
+import torch.utils.model_zoo as model_zoo
+
+from .utils import register
+
+class ResNet(vrn.ResNet):
+    'Deep Residual Network - https://arxiv.org/abs/1512.03385'
+
+    def __init__(self, layers=[3, 4, 6, 3], bottleneck=vrn.Bottleneck, outputs=[5], groups=1, width_per_group=64, url=None):
+        self.stride = 128        
+        self.bottleneck = bottleneck
+        self.outputs = outputs
+        self.url = url
+
+        kwargs = {'block': bottleneck, 'layers': layers, 'groups': groups, 'width_per_group': width_per_group}
+        super().__init__(**kwargs)
+        self.unused_modules = ['fc']
+
+    def initialize(self):
+        if self.url:
+            self.load_state_dict(model_zoo.load_url(self.url))
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        outputs = []
+        for i, layer in enumerate([self.layer1, self.layer2, self.layer3, self.layer4]):
+            level = i + 2
+            if level > max(self.outputs):
+                break
+            x = layer(x)
+            if level in self.outputs:
+                outputs.append(x)
+
+        return outputs
+
+@register
+def ResNet18C4():
+    return ResNet(layers=[2, 2, 2, 2], bottleneck=vrn.BasicBlock, outputs=[4], url=vrn.model_urls['resnet18'])
+
+@register
+def ResNet34C4():
+    return ResNet(layers=[3, 4, 6, 3], bottleneck=vrn.BasicBlock, outputs=[4], url=vrn.model_urls['resnet34'])

+import sys
+import torchvision
+
+def register(f):
+    all = sys.modules[f.__module__].__dict__.setdefault('__all__', [])
+    if f.__name__ in all:
+        raise RuntimeError('{} already exist!'.format(f.__name__))
+    all.append(f.__name__)
+    return f

+from contextlib import redirect_stdout
+from math import ceil
+import ctypes
+import torch
+from nvidia.dali import pipeline, ops, types
+from pycocotools.coco import COCO
+
+class COCOPipeline(pipeline.Pipeline):
+    'Dali pipeline for COCO'
+
+    def __init__(self, batch_size, num_threads, path, training, annotations, world, device_id, mean, std, resize,
+                 max_size, stride, rotate_augment=False,
+                 augment_brightness=0.0,
+                 augment_contrast=0.0, augment_hue=0.0,
+                 augment_saturation=0.0):
+        super().__init__(batch_size=batch_size, num_threads=num_threads, device_id=device_id,
+                         prefetch_queue_depth=num_threads, seed=42)
+        self.path = path
+        self.training = training
+        self.stride = stride
+        self.iter = 0
+
+        self.rotate_augment = rotate_augment
+        self.augment_brightness = augment_brightness
+        self.augment_contrast = augment_contrast
+        self.augment_hue = augment_hue
+        self.augment_saturation = augment_saturation
+
+        self.reader = ops.COCOReader(annotations_file=annotations, file_root=path, num_shards=world,
+                                     shard_id=torch.cuda.current_device(),
+                                     ltrb=True, ratio=True, shuffle_after_epoch=True, save_img_ids=True)
+
+        self.decode_train = ops.ImageDecoderSlice(device="mixed", output_type=types.RGB)
+        self.decode_infer = ops.ImageDecoder(device="mixed", output_type=types.RGB)
+        self.bbox_crop = ops.RandomBBoxCrop(device='cpu', bbox_layout="xyXY", scaling=[0.3, 1.0],
+                                            thresholds=[0.1, 0.3, 0.5, 0.7, 0.9])
+
+        self.bbox_flip = ops.BbFlip(device='cpu', ltrb=True)
+        self.img_flip = ops.Flip(device='gpu')
+        self.coin_flip = ops.CoinFlip(probability=0.5)
+        self.bc = ops.BrightnessContrast(device='gpu')
+        self.hsv = ops.Hsv(device='gpu')
+
+        # Random number generation for augmentation
+        self.brightness_dist = ops.NormalDistribution(mean=1.0, stddev=augment_brightness)
+        self.contrast_dist = ops.NormalDistribution(mean=1.0, stddev=augment_contrast)
+        self.hue_dist = ops.NormalDistribution(mean=0.0, stddev=augment_hue)
+        self.saturation_dist = ops.NormalDistribution(mean=1.0, stddev=augment_saturation)
+
+        if rotate_augment:
+            raise RuntimeWarning("--augment-rotate current has no effect when using the DALI data loader.")
+
+        if isinstance(resize, list): resize = max(resize)
+        self.rand_resize = ops.Uniform(range=[resize, float(max_size)])
+
+        self.resize_train = ops.Resize(device='gpu', interp_type=types.DALIInterpType.INTERP_CUBIC, save_attrs=True)
+        self.resize_infer = ops.Resize(device='gpu', interp_type=types.DALIInterpType.INTERP_CUBIC,
+                                       resize_longer=max_size, save_attrs=True)
+
+        padded_size = max_size + ((self.stride - max_size % self.stride) % self.stride)
+
+        self.pad = ops.Paste(device='gpu', fill_value=0, ratio=1.1, min_canvas_size=padded_size, paste_x=0, paste_y=0)
+        self.normalize = ops.CropMirrorNormalize(device='gpu', mean=mean, std=std, crop=(padded_size, padded_size),
+                                                 crop_pos_x=0, crop_pos_y=0)
+
+    def define_graph(self):
+
+        images, bboxes, labels, img_ids = self.reader()
+
+        if self.training:
+            crop_begin, crop_size, bboxes, labels = self.bbox_crop(bboxes, labels)
+            images = self.decode_train(images, crop_begin, crop_size)
+            resize = self.rand_resize()
+            images, attrs = self.resize_train(images, resize_longer=resize)
+
+            flip = self.coin_flip()
+            bboxes = self.bbox_flip(bboxes, horizontal=flip)
+            images = self.img_flip(images, horizontal=flip)
+
+            if self.augment_brightness or self.augment_contrast:
+                images = self.bc(images, brightness=self.brightness_dist(), contrast=self.contrast_dist())
+            if self.augment_hue or self.augment_saturation:
+                images = self.hsv(images, hue=self.hue_dist(), saturation=self.saturation_dist())
+
+        else:
+            images = self.decode_infer(images)
+            images, attrs = self.resize_infer(images)
+
+        resized_images = images
+        images = self.normalize(self.pad(images))
+
+        return images, bboxes, labels, img_ids, attrs, resized_images
+
+
+class DaliDataIterator():
+    'Data loader for data parallel using Dali'
+
+    def __init__(self, path, resize, max_size, batch_size, stride, world, annotations, training=False,
+                 rotate_augment=False, augment_brightness=0.0,
+                 augment_contrast=0.0, augment_hue=0.0, augment_saturation=0.0):
+        self.training = training
+        self.resize = resize
+        self.max_size = max_size
+        self.stride = stride
+        self.batch_size = batch_size // world
+        self.mean = [255. * x for x in [0.485, 0.456, 0.406]]
+        self.std = [255. * x for x in [0.229, 0.224, 0.225]]
+        self.world = world
+        self.path = path
+
+        # Setup COCO
+        with redirect_stdout(None):
+            self.coco = COCO(annotations)
+        self.ids = list(self.coco.imgs.keys())
+        if 'categories' in self.coco.dataset:
+            self.categories_inv = {k: i for i, k in enumerate(self.coco.getCatIds())}
+
+        self.pipe = COCOPipeline(batch_size=self.batch_size, num_threads=2,
+                                 path=path, training=training, annotations=annotations, world=world,
+                                 device_id=torch.cuda.current_device(), mean=self.mean, std=self.std, resize=resize,
+                                 max_size=max_size, stride=self.stride, rotate_augment=rotate_augment,
+                                 augment_brightness=augment_brightness,
+                                 augment_contrast=augment_contrast, augment_hue=augment_hue,
+                                 augment_saturation=augment_saturation)
+
+        self.pipe.build()
+
+    def __repr__(self):
+        return '\n'.join([
+            '    loader: dali',
+            '    resize: {}, max: {}'.format(self.resize, self.max_size),
+        ])
+
+    def __len__(self):
+        return ceil(len(self.ids) // self.world / self.batch_size)
+
+    def __iter__(self):
+        for _ in range(self.__len__()):
+
+            data, ratios, ids, num_detections = [], [], [], []
+            dali_data, dali_boxes, dali_labels, dali_ids, dali_attrs, dali_resize_img = self.pipe.run()
+
+            for l in range(len(dali_boxes)):
+                num_detections.append(dali_boxes.at(l).shape[0])
+
+            pyt_targets = -1 * torch.ones([len(dali_boxes), max(max(num_detections), 1), 5])
+
+            for batch in range(self.batch_size):
+                id = int(dali_ids.at(batch)[0])
+
+                # Convert dali tensor to pytorch
+                dali_tensor = dali_data[batch]
+                tensor_shape = dali_tensor.shape()
+
+                datum = torch.zeros(dali_tensor.shape(), dtype=torch.float, device=torch.device('cuda'))
+                c_type_pointer = ctypes.c_void_p(datum.data_ptr())
+                dali_tensor.copy_to_external(c_type_pointer)
+
+                # Calculate image resize ratio to rescale boxes
+                prior_size = dali_attrs.as_cpu().at(batch)
+                resized_size = dali_resize_img[batch].shape()
+                ratio = max(resized_size) / max(prior_size)
+
+                if self.training:
+                    # Rescale boxes
+                    b_arr = dali_boxes.at(batch)
+                    num_dets = b_arr.shape[0]
+                    if num_dets!=0:
+                        pyt_bbox = torch.from_numpy(b_arr).float()
+
+                        pyt_bbox[:, 0] *= float(prior_size[1])
+                        pyt_bbox[:, 1] *= float(prior_size[0])
+                        pyt_bbox[:, 2] *= float(prior_size[1])
+                        pyt_bbox[:, 3] *= float(prior_size[0])
+                        # (l,t,r,b) ->  (x,y,w,h) == (l,r, r-l, b-t)
+                        pyt_bbox[:, 2] -= pyt_bbox[:, 0]
+                        pyt_bbox[:, 3] -= pyt_bbox[:, 1]
+                        pyt_targets[batch, :num_dets, :4] = pyt_bbox * ratio
+
+                    # Arrange labels in target tensor
+                    l_arr = dali_labels.at(batch)
+                    if num_dets!=0:
+                        pyt_label = torch.from_numpy(l_arr).float()
+                        pyt_label -= 1  # Rescale labels to [0,79] instead of [1,80]
+                        pyt_targets[batch, :num_dets, 4] = pyt_label.squeeze()
+
+                ids.append(id)
+                data.append(datum.unsqueeze(0))
+                ratios.append(ratio)
+
+            data = torch.cat(data, dim=0)
+
+            if self.training:
+                pyt_targets = pyt_targets.cuda(non_blocking=True)
+                yield data, pyt_targets
+
+            else:
+                ids = torch.Tensor(ids).int().cuda(non_blocking=True)
+                ratios = torch.Tensor(ratios).cuda(non_blocking=True)
+                yield data, ids, ratios
+

+import os
+import json
+import tempfile
+from contextlib import redirect_stdout
+import torch
+from apex import amp
+from apex.parallel import DistributedDataParallel as ADDP
+from torch.nn.parallel import DistributedDataParallel
+from pycocotools.cocoeval import COCOeval
+import numpy as np
+
+from .data import DataIterator, RotatedDataIterator
+from .dali import DaliDataIterator
+from .model import Model
+from .utils import Profiler, rotate_box
+
+
+def infer(model, path, detections_file, resize, max_size, batch_size, mixed_precision=True, is_master=True, world=0,
+          annotations=None, no_apex=False, use_dali=True, is_validation=False, verbose=True, rotated_bbox=False):
+    'Run inference on images from path'
+
+    DDP = DistributedDataParallel if no_apex else ADDP
+    backend = 'pytorch' if isinstance(model, Model) or isinstance(model, DDP) else 'tensorrt'
+
+    stride = model.module.stride if isinstance(model, DDP) else model.stride
+
+    # Create annotations if none was provided
+    if not annotations:
+        annotations = tempfile.mktemp('.json')
+        images = [{'id': i, 'file_name': f} for i, f in enumerate(os.listdir(path))]
+        json.dump({'images': images}, open(annotations, 'w'))
+
+    # TensorRT only supports fixed input sizes, so override input size accordingly
+    if backend == 'tensorrt': max_size = max(model.input_size)
+
+    # Prepare dataset
+    if verbose: print('Preparing dataset...')
+    if rotated_bbox:
+        if use_dali: raise NotImplementedError("This repo does not currently support DALI for rotated bbox.")
+        data_iterator = RotatedDataIterator(path, resize, max_size, batch_size, stride,
+                                            world, annotations, training=False)
+    else:
+        data_iterator = (DaliDataIterator if use_dali else DataIterator)(
+            path, resize, max_size, batch_size, stride,
+            world, annotations, training=False)
+    if verbose: print(data_iterator)
+
+    # Prepare model
+    if backend == 'pytorch':
+        # If we are doing validation during training,
+        # no need to register model with AMP again
+        if not is_validation:
+            if torch.cuda.is_available(): model = model.to(memory_format=torch.channels_last).cuda()
+            if not no_apex:
+                model = amp.initialize(model, None,
+                                    opt_level='O2' if mixed_precision else 'O0',
+                                    keep_batchnorm_fp32=True,
+                                    verbosity=0)
+
+        model.eval()
+
+    if verbose:
+        print('   backend: {}'.format(backend))
+        print('    device: {} {}'.format(
+            world, 'cpu' if not torch.cuda.is_available() else 'GPU' if world == 1 else 'GPUs'))
+        print('     batch: {}, precision: {}'.format(batch_size,
+                                                     'unknown' if backend == 'tensorrt' else 'mixed' if mixed_precision else 'full'))
+        print(' BBOX type:', 'rotated' if rotated_bbox else 'axis aligned')
+        print('Running inference...')
+
+    results = []
+    profiler = Profiler(['infer', 'fw'])
+    with torch.no_grad():
+        for i, (data, ids, ratios) in enumerate(data_iterator):
+            # Forward pass
+            if backend=='pytorch': data = data.contiguous(memory_format=torch.channels_last)
+            profiler.start('fw')
+            scores, boxes, classes = model(data, rotated_bbox) #Need to add model size (B, 3, W, H)
+            profiler.stop('fw')
+
+            results.append([scores, boxes, classes, ids, ratios])
+
+            profiler.bump('infer')
+            if verbose and (profiler.totals['infer'] > 60 or i == len(data_iterator) - 1):
+                size = len(data_iterator.ids)
+                msg = '[{:{len}}/{}]'.format(min((i + 1) * batch_size,
+                                                 size), size, len=len(str(size)))
+                msg += ' {:.3f}s/{}-batch'.format(profiler.means['infer'], batch_size)
+                msg += ' (fw: {:.3f}s)'.format(profiler.means['fw'])
+                msg += ', {:.1f} im/s'.format(batch_size / profiler.means['infer'])
+                print(msg, flush=True)
+
+                profiler.reset()
+
+    # Gather results from all devices
+    if verbose: print('Gathering results...')
+    results = [torch.cat(r, dim=0) for r in zip(*results)]
+    if world > 1:
+        for r, result in enumerate(results):
+            all_result = [torch.ones_like(result, device=result.device) for _ in range(world)]
+            torch.distributed.all_gather(list(all_result), result)
+            results[r] = torch.cat(all_result, dim=0)
+
+    if is_master:
+        # Copy buffers back to host
+        results = [r.cpu() for r in results]
+
+        # Collect detections
+        detections = []
+        processed_ids = set()
+        for scores, boxes, classes, image_id, ratios in zip(*results):
+            image_id = image_id.item()
+            if image_id in processed_ids:
+                continue
+            processed_ids.add(image_id)
+
+            keep = (scores > 0).nonzero(as_tuple=False)
+            scores = scores[keep].view(-1)
+            if rotated_bbox:
+                boxes = boxes[keep, :].view(-1, 6)
+                boxes[:, :4] /= ratios
+            else:
+                boxes = boxes[keep, :].view(-1, 4) / ratios
+            classes = classes[keep].view(-1).int()
+
+            for score, box, cat in zip(scores, boxes, classes):
+                if rotated_bbox:
+                    x1, y1, x2, y2, sin, cos = box.data.tolist()
+                    theta = np.arctan2(sin, cos)
+                    w = x2 - x1 + 1
+                    h = y2 - y1 + 1
+                    seg = rotate_box([x1, y1, w, h, theta])
+                else:
+                    x1, y1, x2, y2 = box.data.tolist()
+                cat = cat.item()
+                if 'annotations' in data_iterator.coco.dataset:
+                    cat = data_iterator.coco.getCatIds()[cat]
+                this_det = {
+                    'image_id': image_id,
+                    'score': score.item(),
+                    'category_id': cat}
+                if rotated_bbox:
+                    this_det['bbox'] = [x1, y1, x2 - x1 + 1, y2 - y1 + 1, theta]
+                    this_det['segmentation'] = [seg]
+                else:
+                    this_det['bbox'] = [x1, y1, x2 - x1 + 1, y2 - y1 + 1]
+
+                detections.append(this_det)
+
+        if detections:
+            # Save detections
+            if detections_file and verbose: print('Writing {}...'.format(detections_file))
+            detections = {'annotations': detections}
+            detections['images'] = data_iterator.coco.dataset['images']
+            if 'categories' in data_iterator.coco.dataset:
+                detections['categories'] = data_iterator.coco.dataset['categories']
+            if detections_file:
+                for d_file in detections_file:
+                    json.dump(detections, open(d_file, 'w'), indent=4)
+
+            # Evaluate model on dataset
+            if 'annotations' in data_iterator.coco.dataset:
+                if verbose: print('Evaluating model...')
+                with redirect_stdout(None):
+                    coco_pred = data_iterator.coco.loadRes(detections['annotations'])
+                    if rotated_bbox:
+                        coco_eval = COCOeval(data_iterator.coco, coco_pred, 'segm')
+                    else:
+                        coco_eval = COCOeval(data_iterator.coco, coco_pred, 'bbox')
+                    coco_eval.evaluate()
+                    coco_eval.accumulate()
+                coco_eval.summarize()
+                return coco_eval.stats # mAP and mAR
+        else:
+            print('No detections!')
+            return None
+    return 0

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class FocalLoss(nn.Module):
+    'Focal Loss - https://arxiv.org/abs/1708.02002'
+
+    def __init__(self, alpha=0.25, gamma=2):
+        super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+
+    def forward(self, pred_logits, target):
+        pred = pred_logits.sigmoid()
+        ce = F.binary_cross_entropy_with_logits(pred_logits, target, reduction='none')
+        alpha = target * self.alpha + (1. - target) * (1. - self.alpha)
+        pt = torch.where(target == 1,  pred, 1 - pred)
+        return alpha * (1. - pt) ** self.gamma * ce
+
+class SmoothL1Loss(nn.Module):
+    'Smooth L1 Loss'
+
+    def __init__(self, beta=0.11):
+        super().__init__()
+        self.beta = beta
+
+    def forward(self, pred, target):
+        x = (pred - target).abs()
+        l1 = x - 0.5 * self.beta
+        l2 = 0.5 * x ** 2 / self.beta
+        return torch.where(x >= self.beta, l1, l2)

+import os.path
+import time
+import json
+import warnings
+import signal
+from datetime import datetime
+from contextlib import contextmanager
+from PIL import Image, ImageDraw
+import requests
+import numpy as np
+import math
+import torch
+
+
+def order_points(pts):
+    pts_reorder = []
+
+    for idx, pt in enumerate(pts):
+        idx = torch.argsort(pt[:, 0])
+        xSorted = pt[idx, :]
+        leftMost = xSorted[:2, :]
+        rightMost = xSorted[2:, :]
+
+        leftMost = leftMost[torch.argsort(leftMost[:, 1]), :]
+        (tl, bl) = leftMost
+
+        D = torch.cdist(tl[np.newaxis], rightMost)[0]
+        (br, tr) = rightMost[torch.argsort(D, descending=True), :]
+        pts_reorder.append(torch.stack([tl, tr, br, bl]))
+
+    return torch.stack([p for p in pts_reorder])
+
+def rotate_boxes(boxes, points=False):
+    '''
+    Rotate target bounding boxes
+
+    Input:
+        Target boxes (xmin_ymin, width_height, theta)
+    Output:
+        boxes_axis (xmin_ymin, xmax_ymax, theta)
+        boxes_rotated (xy0, xy1, xy2, xy3)
+    '''
+
+    u = torch.stack([torch.cos(boxes[:,4]), torch.sin(boxes[:,4])], dim=1)
+    l = torch.stack([-torch.sin(boxes[:,4]), torch.cos(boxes[:,4])], dim=1)
+    R = torch.stack([u, l], dim=1)
+
+    if points:
+        cents = torch.stack([(boxes[:,0]+boxes[:,2])/2, (boxes[:,1]+boxes[:,3])/2],1).transpose(1,0)
+        boxes_rotated = torch.stack([boxes[:,0],boxes[:,1],
+            boxes[:,2], boxes[:,1],
+            boxes[:,2], boxes[:,3],
+            boxes[:,0], boxes[:,3],
+            boxes[:,-2],
+            boxes[:,-1]],1)
+
+    else:
+        cents = torch.stack([boxes[:,0]+(boxes[:,2])/2, boxes[:,1]+(boxes[:,3])/2],1).transpose(1,0)
+        boxes_rotated = torch.stack([boxes[:,0],boxes[:,1],
+            (boxes[:,0]+boxes[:,2]), boxes[:,1],
+            (boxes[:,0]+boxes[:,2]), (boxes[:,1]+boxes[:,3]),
+            boxes[:,0], (boxes[:,1]+boxes[:,3]),
+            boxes[:,-2],
+            boxes[:,-1]],1)
+
+    xy0R = torch.matmul(R,boxes_rotated[:,:2].transpose(1,0) - cents) + cents
+    xy1R = torch.matmul(R,boxes_rotated[:,2:4].transpose(1,0) - cents) + cents
+    xy2R = torch.matmul(R,boxes_rotated[:,4:6].transpose(1,0) - cents) + cents
+    xy3R = torch.matmul(R,boxes_rotated[:,6:8].transpose(1,0) - cents) + cents
+
+    xy0R = torch.stack([xy0R[i,:,i] for i in range(xy0R.size(0))])
+    xy1R = torch.stack([xy1R[i,:,i] for i in range(xy1R.size(0))])
+    xy2R = torch.stack([xy2R[i,:,i] for i in range(xy2R.size(0))])
+    xy3R = torch.stack([xy3R[i,:,i] for i in range(xy3R.size(0))])
+
+    boxes_axis = torch.cat([boxes[:, :2], boxes[:, :2] + boxes[:, 2:4] - 1,
+        torch.sin(boxes[:,-1, None]), torch.cos(boxes[:,-1, None])], 1)
+    boxes_rotated = order_points(torch.stack([xy0R,xy1R,xy2R,xy3R],dim = 1)).view(-1,8)
+
+    return boxes_axis, boxes_rotated
+
+
+def rotate_box(bbox):
+    xmin, ymin, width, height, theta = bbox
+
+    xy1 = xmin, ymin
+    xy2 = xmin, ymin + height - 1
+    xy3 = xmin + width - 1, ymin + height - 1
+    xy4 = xmin + width - 1, ymin
+
+    cents = np.array([xmin + (width - 1) / 2, ymin + (height - 1) / 2])
+
+    corners = np.stack([xy1, xy2, xy3, xy4])
+
+    u = np.stack([np.cos(theta), -np.sin(theta)])
+    l = np.stack([np.sin(theta), np.cos(theta)])
+    R = np.vstack([u, l])
+
+    corners = np.matmul(R, (corners - cents).transpose(1, 0)).transpose(1, 0) + cents
+
+    return corners.reshape(-1).tolist()
+
+
+def show_detections(detections):
+    'Show image with drawn detections'
+
+    for image, detections in detections.items():
+        im = Image.open(image).convert('RGBA')
+        overlay = Image.new('RGBA', im.size, (255, 255, 255, 0))
+        draw = ImageDraw.Draw(overlay)
+        detections.sort(key=lambda d: d['score'])
+        for detection in detections:
+            box = detection['bbox']
+            alpha = int(detection['score'] * 255)
+            draw.rectangle(box, outline=(255, 255, 255, alpha))
+            draw.text((box[0] + 2, box[1]), '[{}]'.format(detection['class']),
+                      fill=(255, 255, 255, alpha))
+            draw.text((box[0] + 2, box[1] + 10), '{:.2}'.format(detection['score']),
+                      fill=(255, 255, 255, alpha))
+        im = Image.alpha_composite(im, overlay)
+        im.show()
+
+
+def save_detections(path, detections):
+    print('Writing detections to {}...'.format(os.path.basename(path)))
+    with open(path, 'w') as f:
+        json.dump(detections, f)
+
+
+@contextmanager
+def ignore_sigint():
+    handler = signal.getsignal(signal.SIGINT)
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+    try:
+        yield
+    finally:
+        signal.signal(signal.SIGINT, handler)
+
+
+class Profiler(object):
+    def __init__(self, names=['main']):
+        self.names = names
+        self.lasts = {k: 0 for k in names}
+        self.totals = self.lasts.copy()
+        self.counts = self.lasts.copy()
+        self.means = self.lasts.copy()
+        self.reset()
+
+    def reset(self):
+        last = time.time()
+        for name in self.names:
+            self.lasts[name] = last
+            self.totals[name] = 0
+            self.counts[name] = 0
+            self.means[name] = 0
+
+    def start(self, name='main'):
+        self.lasts[name] = time.time()
+
+    def stop(self, name='main'):
+        self.totals[name] += time.time() - self.lasts[name]
+        self.counts[name] += 1
+        self.means[name] = self.totals[name] / self.counts[name]
+
+    def bump(self, name='main'):
+        self.stop(name)
+        self.start(name)
+
+
+def post_metrics(url, metrics):
+    try:
+        for k, v in metrics.items():
+            requests.post(url,
+                          data={'time': int(datetime.now().timestamp() * 1e9),
+                                'metric': k, 'value': v})
+    except Exception as e:
+        warnings.warn('Warning: posting metrics failed: {}'.format(e))

+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+    name='odtk',
+    version='0.2.6',
+    description='Fast and accurate single shot object detector',
+    author = 'NVIDIA Corporation',
+    packages=['odtk', 'odtk.backbones'],
+    ext_modules=[CUDAExtension('odtk._C',
+        ['csrc/extensions.cpp', 'csrc/engine.cpp', 'csrc/cuda/decode.cu', 'csrc/cuda/decode_rotate.cu', 'csrc/cuda/nms.cu', 'csrc/cuda/nms_iou.cu'],
+        extra_compile_args={
+            'cxx': ['-std=c++14', '-O2', '-Wall'],
+            'nvcc': [
+                '-std=c++14', '--expt-extended-lambda', '--use_fast_math', '-Xcompiler', '-Wall,-fno-gnu-unique',
+                '-gencode=arch=compute_60,code=sm_60', '-gencode=arch=compute_61,code=sm_61',
+                '-gencode=arch=compute_70,code=sm_70', '-gencode=arch=compute_72,code=sm_72',
+                '-gencode=arch=compute_75,code=sm_75', '-gencode=arch=compute_80,code=sm_80',
+                '-gencode=arch=compute_86,code=sm_86', '-gencode=arch=compute_86,code=compute_86'
+            ],
+        },
+        libraries=['nvinfer', 'nvinfer_plugin', 'nvonnxparser'])
+    ],
+    cmdclass={'build_ext': BuildExtension.with_options(no_python_abi_suffix=True)},
+    install_requires=[
+        'torch>=1.0.0a0',
+        'torchvision',
+        'apex @ git+https://github.com/NVIDIA/apex',
+        'pycocotools @ git+https://github.com/nvidia/cocoapi.git#subdirectory=PythonAPI',
+        'pillow',
+        'requests',
+    ],
+    entry_points = {'console_scripts': ['odtk=odtk.main:main']}
+)

+# voc2coco
+
+This is script for converting VOC format XMLs to COCO format json(ex. coco_eval.json).
+
+### Why we need to convert VOC xmls to COCO format json ?
+
+We can use COCO API, this is very useful(ex. calculating mAP).
+
+## How to use
+
+### 1. Make labels.txt
+
+labels.txt if need for making dictionary for converting label to id.
+
+**Sample labels.txt**
+
+```txt
+Label1
+Label2
+...
+```
+
+### 2. Run script
+
+##### 2.1 Usage 1(Use ids list)
+
+```bash
+$ python voc2coco.py \
+    --ann_dir /path/to/annotation/dir \
+    --ann_ids /path/to/annotations/ids/list.txt \
+    --labels /path/to/labels.txt \
+    --output /path/to/output.json \
+    <option> --ext xml
+```
+
+##### 2.2 Usage 2(Use annotation paths list)
+
+**Sample paths.txt**
+
+```txt
+/path/to/annotation/file.xml
+/path/to/annotation/file2.xml
+...
+```
+
+```bash
+$ python voc2coco.py \
+    --ann_paths_list /path/to/annotation/paths.txt \
+    --labels /path/to/labels.txt \
+    --output /path/to/output.json \
+    <option> --ext xml
+```
+
+### 3. Example of usage
+
+In this case, you can convert [Shenggan/BCCD_Dataset: BCCD Dataset is a small-scale dataset for blood cells detection.](https://github.com/Shenggan/BCCD_Dataset) by this script.
+
+```bash
+$ python voc2coco.py
+    --ann_dir sample/Annotations \
+    --ann_ids sample/dataset_ids/test.txt \
+    --labels sample/labels.txt \
+    --output sample/bccd_test_cocoformat.json \
+    --ext xml
+
+# Check output
+$ ls sample/ | grep bccd_test_cocoformat.json
+bccd_test_cocoformat.json
+
+# Check output
+cut -f -4 -d , sample/bccd_test_cocoformat.json
+{"images": [{"file_name": "BloodImage_00007.jpg", "height": 480, "width": 640, "id": "BloodImage_00007"}
+```