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.
+ */
+
+#include "nms_iou.h"
+#include "utils.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <iostream>
+#include <stdexcept>
+#include <vector>
+#include <cuda.h>
+
+#include <thrust/device_ptr.h>
+#include <thrust/execution_policy.h>
+#include <thrust/gather.h>
+#include <thrust/sequence.h>
+#include <cub/device/device_radix_sort.cuh>
+#include <cub/iterator/counting_input_iterator.cuh>
+
+constexpr int   kTPB     = 64;  // threads per block
+constexpr int   kCorners = 4;
+constexpr int   kPoints  = 8;
+
+namespace odtk {
+namespace cuda {
+
+class Vector {
+public:
+    __host__ __device__ Vector( );  // Default constructor
+    __host__ __device__ ~Vector( );  // Deconstructor
+    __host__ __device__ Vector( float2 const point );
+    float2 const p;
+    friend class Line;
+
+private:
+    __host__ __device__ float cross( Vector const v ) const;
+};
+
+Vector::Vector( ) : p( make_float2( 0.0f, 0.0f ) ) {}
+
+Vector::~Vector( ) {}
+
+Vector::Vector( float2 const point ) : p( point ) {}
+
+float Vector::cross( Vector const v ) const {
+    return ( p.x * v.p.y - p.y * v.p.x );
+}
+
+class Line {
+public:
+    __host__ __device__ Line( );  // Default constructor
+    __host__ __device__ ~Line( );  // Deconstructor
+    __host__ __device__ Line( Vector const v1, Vector const v2 );
+    __host__ __device__ float call( Vector const v ) const;
+    __host__ __device__ float2 intersection( Line const l ) const;
+
+private:
+    float const a;
+    float const b;
+    float const c;
+};
+
+Line::Line( ) : a( 0.0f ), b( 0.0f ), c( 0.0f ) {}
+
+Line::~Line( ) {}
+
+Line::Line( Vector const v1, Vector const v2 ) : a( v2.p.y - v1.p.y ), b( v1.p.x - v2.p.x ), c( v2.cross( v1 ) ) {}
+
+float Line::call( Vector const v ) const {
+    return ( a * v.p.x + b * v.p.y + c );
+}
+
+float2 Line::intersection( Line const l ) const {
+    float w { a * l.b - b * l.a };
+    return ( make_float2( ( b * l.c - c * l.b ) / w, ( c * l.a - a * l.c ) / w ) );
+}
+
+template<typename T>
+__host__ __device__ void rotateLeft( T *array, int const &count ) {
+    T temp = array[0];
+    for ( int i = 0; i < count - 1; i++ )
+        array[i] = array[i + 1];
+    array[count - 1] = temp;
+}
+
+__host__ __device__ static __inline__ float2 padfloat2( float2 a, float2 b ) {
+    float2 res;
+    res.x = a.x + b.x;
+    res.y = a.y + b.y;
+    return res;
+}
+
+__device__ float IntersectionArea( float2 *mrect, float2 *mrect_shift, float2 *intersection ) {
+    int count = kCorners;
+    for ( int i = 0; i < kCorners; i++ ) {
+        float2 intersection_shift[kPoints] {};
+        for ( int k = 0; k < count; k++ )
+            intersection_shift[k] = intersection[k];
+        float line_values[kPoints] {};
+        Vector const r1( mrect[i] );
+        Vector const r2( mrect_shift[i] );
+        Line const   line1( r1, r2 );
+        for ( int j = 0; j < count; j++ ) {
+            Vector const inter( intersection[j] );
+            line_values[j] = line1.call( inter );
+        }
+        float line_values_shift[kPoints] {};
+
+#pragma unroll
+        for ( int k = 0; k < kPoints; k++ )
+            line_values_shift[k] = line_values[k];
+        rotateLeft( line_values_shift, count );
+        rotateLeft( intersection_shift, count );
+        float2 new_intersection[kPoints] {};
+        int temp = count;
+        count = 0;
+        for ( int j = 0; j < temp; j++ ) {
+            if ( line_values[j] <= 0 ) {
+                new_intersection[count] = intersection[j];
+                count++;
+            }
+            if ( ( line_values[j] * line_values_shift[j] ) <= 0 ) {
+                Vector const r3( intersection[j] );
+                Vector const r4( intersection_shift[j] );
+                Line const Line( r3, r4 );
+                new_intersection[count] = line1.intersection( Line );
+                count++;
+            }
+        }
+        for ( int k = 0; k < count; k++ )
+            intersection[k] = new_intersection[k];
+    }
+
+    float2 intersection_shift[kPoints] {};
+
+    for ( int k = 0; k < count; k++ )
+        intersection_shift[k] = intersection[k];
+    rotateLeft( intersection_shift, count );
+
+    // Intersection
+    float intersection_area = 0.0f;
+    if ( count > 2 ) {
+        for ( int k = 0; k < count; k++ )
+            intersection_area +=
+                intersection[k].x * intersection_shift[k].y - intersection[k].y * intersection_shift[k].x;
+    }
+    return ( abs( intersection_area / 2.0f ) );
+}
+
+__global__ void nms_rotate_kernel(const int num_per_thread, const float threshold, const int num_detections, 
+  const int *indices, float *scores, const float *classes, const float6 *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 ii = threadIdx.x * num_per_thread + n;
+      if ( ii < num_detections && m < ii && scores[m] > 0.0f ) {
+        int idx     = indices[ii];
+        int max_idx = indices[m];
+        int icls    = classes[idx];
+        int mcls    = classes[max_idx];
+        if ( mcls == icls ) {
+          float6 ibox = make_float6( make_float4( boxes[idx].x1,
+                                                  boxes[idx].y1,
+                                                  boxes[idx].x2,
+                                                  boxes[idx].y2 ),
+                                  make_float2( boxes[idx].s, boxes[idx].c ) );
+          float6 mbox = make_float6( make_float4( boxes[max_idx].x1,
+                                                  boxes[max_idx].y1,
+                                                  boxes[max_idx].x2,
+                                                  boxes[max_idx].y2 ),
+                                  make_float2( boxes[idx].s, boxes[idx].c ) );
+          float2 intersection[kPoints] { -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f,
+                                      -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f };
+          float2 irect[kPoints] {};
+          float2 irect_shift[kPoints] {};
+          float2 mrect[kPoints] {};
+          float2 mrect_shift[kPoints] {};
+          float2 icent = { ( ibox.x1 + ibox.x2 ) / 2.0f, ( ibox.y1 + ibox.y2 ) / 2.0f };
+          float2 mcent = { ( mbox.x1 + mbox.x2 ) / 2.0f, ( mbox.y1 + mbox.y2 ) / 2.0f };
+          float2 iboxc[kCorners] = { ibox.x1 - icent.x, ibox.y1 - icent.y, ibox.x2 - icent.x,
+                                  ibox.y1 - icent.y, ibox.x2 - icent.x, ibox.y2 - icent.y,
+                                  ibox.x1 - icent.x, ibox.y2 - icent.y };
+          float2 mboxc[kCorners] = { mbox.x1 - mcent.x, mbox.y1 - mcent.y, mbox.x2 - mcent.x,
+                                  mbox.y1 - mcent.y, mbox.x2 - mcent.x, mbox.y2 - mcent.y,
+                                  mbox.x1 - mcent.x, mbox.y2 - mcent.y };
+          float2 pad;
+#pragma unroll
+          for ( int b = 0; b < kCorners; b++ ) {
+            if ((iboxc[b].x * ibox.c - iboxc[b].y * ibox.s) + icent.x == (mboxc[b].x * mbox.c - mboxc[b].y * mbox.s) + mcent.x)
+              pad.x = 0.001f;
+            else
+              pad.x = 0.0f;
+            if ((iboxc[b].y * ibox.c + iboxc[b].x * ibox.s) + icent.y == (mboxc[b].y * mbox.c + mboxc[b].x * mbox.s) + mcent.y)
+              pad.y = 0.001f;
+            else
+              pad.y = 0.0f;
+            intersection[b] = { ( iboxc[b].x * ibox.c - iboxc[b].y * ibox.s ) + icent.x + pad.x,
+                                ( iboxc[b].y * ibox.c + iboxc[b].x * ibox.s ) + icent.y + pad.y};
+            irect[b]        = { ( iboxc[b].x * ibox.c - iboxc[b].y * ibox.s ) + icent.x,
+                        ( iboxc[b].y * ibox.c + iboxc[b].x * ibox.s ) + icent.y };
+            irect_shift[b]  = { ( iboxc[b].x * ibox.c - iboxc[b].y * ibox.s ) + icent.x,
+                            ( iboxc[b].y * ibox.c + iboxc[b].x * ibox.s ) + icent.y };
+            mrect[b]        = { ( mboxc[b].x * mbox.c - mboxc[b].y * mbox.s ) + mcent.x,
+                        ( mboxc[b].y * mbox.c + mboxc[b].x * mbox.s ) + mcent.y };
+            mrect_shift[b]  = { ( mboxc[b].x * mbox.c - mboxc[b].y * mbox.s ) + mcent.x,
+                            ( mboxc[b].y * mbox.c + mboxc[b].x * mbox.s ) + mcent.y };
+          }
+          rotateLeft( irect_shift, 4 );
+          rotateLeft( mrect_shift, 4 );
+          float intersection_area = IntersectionArea( mrect, mrect_shift, intersection );
+          // Union
+          float irect_area = 0.0f;
+          float mrect_area = 0.0f;
+#pragma unroll
+          for ( int k = 0; k < kCorners; k++ ) {
+            irect_area += irect[k].x * irect_shift[k].y - irect[k].y * irect_shift[k].x;
+            mrect_area += mrect[k].x * mrect_shift[k].y - mrect[k].y * mrect_shift[k].x;
+          }
+          float union_area = ( abs( irect_area ) + abs( mrect_area ) ) / 2.0f;
+          float overlap;
+          if ( isnan( intersection_area ) && isnan( union_area ) ) {
+            overlap = 1.0f;
+          } else if ( isnan( intersection_area ) ) {
+            overlap = 0.0f;
+          } else {
+            overlap = intersection_area / ( union_area - intersection_area );  // Check nans and inf
+          }
+          if ( overlap > threshold ) {
+            scores[ii] = 0.0f;
+          }
+        }
+      }
+    }
+    // Sync discarded detections
+    __syncthreads( );
+  }
+}
+
+int nms_rotate(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 float6 *>( 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<float6 *>( 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 ); // From 8
+    // 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_rotate_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 );  // From 8
+    // 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;
+}
+
+__global__ void iou_cuda_kernel(int const numBoxes, int const numAnchors,
+  float2 const *b_box_vals, float2 const *a_box_vals, float *iou_vals ) {
+  int t      = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = blockDim.x * gridDim.x;
+  int combos = numBoxes * numAnchors;
+  for ( int tid = t; tid < combos; tid += stride ) {
+    float2 intersection[kPoints] { -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f,
+                                -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f };
+    float2 rect1[kPoints] {};
+    float2 rect1_shift[kPoints] {};
+    float2 rect2[kPoints] {};
+    float2 rect2_shift[kPoints] {};
+    float2 pad;
+#pragma unroll
+    for ( int b = 0; b < kCorners; b++ ) {
+      if (b_box_vals[(static_cast<int>(tid/numAnchors) * kCorners + b)].x == a_box_vals[(tid * kCorners + b) % (numAnchors * kCorners)].x)
+        pad.x = 0.001f;
+      else
+        pad.x = 0.0f;
+      if (b_box_vals[(static_cast<int>(tid/numAnchors) * kCorners + b)].y == a_box_vals[(tid * kCorners + b) % (numAnchors * kCorners)].y)
+        pad.y = 0.001f;
+      else
+        pad.y = 0.0f;
+      intersection[b] = padfloat2( b_box_vals[( static_cast<int>( tid / numAnchors ) * kCorners + b )], pad);
+      rect1[b]        = b_box_vals[( static_cast<int>( tid / numAnchors ) * kCorners + b )];
+      rect1_shift[b]  = b_box_vals[( static_cast<int>( tid / numAnchors ) * kCorners + b )];
+      rect2[b]        = a_box_vals[( tid * kCorners + b ) % ( numAnchors * kCorners )];
+      rect2_shift[b]  = a_box_vals[( tid * kCorners + b ) % ( numAnchors * kCorners )];
+    }
+    rotateLeft( rect1_shift, 4 );
+    rotateLeft( rect2_shift, 4 );
+    float intersection_area = IntersectionArea( rect2, rect2_shift, intersection );
+    // Union
+    float rect1_area = 0.0f;
+    float rect2_area = 0.0f;
+#pragma unroll
+    for ( int k = 0; k < kCorners; k++ ) {
+        rect1_area += rect1[k].x * rect1_shift[k].y - rect1[k].y * rect1_shift[k].x;
+        rect2_area += rect2[k].x * rect2_shift[k].y - rect2[k].y * rect2_shift[k].x;
+    }
+    float union_area = ( abs( rect1_area ) + abs( rect2_area ) ) / 2.0f;
+    float iou_val = intersection_area / ( union_area - intersection_area );
+    // Write out answer
+    if ( isnan( intersection_area ) && isnan( union_area ) ) {
+        iou_vals[tid] = 1.0f;
+    } else if ( isnan( intersection_area ) ) {
+        iou_vals[tid] = 0.0f;
+    } else {
+        iou_vals[tid] = iou_val;
+    }
+  }
+}
+
+int iou( const void *const *inputs, void *const *outputs, int num_boxes, int num_anchors, cudaStream_t stream ) {
+  auto boxes    = static_cast<const float2 *>( inputs[0] );
+  auto anchors  = static_cast<const float2 *>( inputs[1] );
+  auto iou_vals = static_cast<float *>( outputs[0] );
+  int numSMs;
+  cudaDeviceGetAttribute( &numSMs, cudaDevAttrMultiProcessorCount, 0 );
+  int threadsPerBlock = kTPB;
+  int blocksPerGrid   = numSMs * 10;
+  iou_cuda_kernel<<<blocksPerGrid, threadsPerBlock, 0, stream>>>( num_anchors, num_boxes, anchors, boxes, iou_vals );
+  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_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)
+
+
+
+

+#include <iostream>
+#include <stdexcept>
+#include <fstream>
+#include <vector>
+#include <glob.h>
+
+#include "../../csrc/engine.h"
+
+#define ROTATED false // Change to true for Rotated Bounding Box export
+#define COCO_PATH "/coco/coco2017/val2017" // Path to calibration images
+
+using namespace std;
+
+// Sample program to build a TensorRT Engine from an ONNX model from RetinaNet
+//
+// By default TensorRT will target FP16 precision (supported on Pascal, Volta, and Turing GPUs)
+//
+// You can optionally provide an INT8CalibrationTable file created during RetinaNet INT8 calibration
+// to build a TensorRT engine with INT8 precision
+
+inline vector<string> glob(int batch){
+	glob_t glob_result;
+	string path = string(COCO_PATH);
+	if(path.back()!='/') path+="/";
+	glob((path+"*").c_str(), (GLOB_TILDE | GLOB_NOSORT), NULL, &glob_result);
+	vector<string> calibration_files;
+	for(int i=0; i<batch; i++){
+		calibration_files.push_back(string(glob_result.gl_pathv[i]));
+	}
+	globfree(&glob_result);
+	return calibration_files;
+}
+
+int main(int argc, char *argv[]) {
+	if (argc != 3 && argc != 4) {
+		cerr << "Usage: " << argv[0] << " core_model.onnx engine.plan {Int8CalibrationTable}" << endl;
+		return 1;
+	}
+
+	ifstream onnxFile;
+	onnxFile.open(argv[1], ios::in | ios::binary); 
+
+	if (!onnxFile.good()) {
+		cerr << "\nERROR: Unable to read specified ONNX model " << argv[1] << endl;
+		return -1;
+	}
+
+	onnxFile.seekg (0, onnxFile.end);
+	size_t size = onnxFile.tellg();
+	onnxFile.seekg (0, onnxFile.beg);
+
+	auto *buffer = new char[size];
+	onnxFile.read(buffer, size);
+	onnxFile.close();
+
+	// Define default RetinaNet parameters to use for TRT export
+	const vector<int> dynamic_batch_opts{1, 8, 16};
+	int calibration_batches = 2; // must be >= 1
+	float score_thresh = 0.05f;
+	int top_n = 1000;
+	size_t workspace_size =(1ULL << 30);
+	float nms_thresh = 0.5;
+	int detections_per_im = 100;
+	bool verbose = false;
+	// Generated from generate_anchors.py
+	vector<vector<float>> anchors;
+	if(!ROTATED) {
+		// Axis-aligned
+		anchors = {
+			{-12.0, -12.0, 20.0, 20.0, -7.31, -18.63, 15.31, 26.63, -18.63, -7.31, 26.63, 15.31, -16.16, -16.16, 24.16, 24.16, -10.25, -24.51, 18.25, 32.51, -24.51, -10.25, 32.51, 18.25, -21.4, -21.4, 29.4, 29.4, -13.96, -31.92, 21.96, 39.92, -31.92, -13.96, 39.92, 21.96},
+			{-24.0, -24.0, 40.0, 40.0, -14.63, -37.25, 30.63, 53.25, -37.25, -14.63, 53.25, 30.63, -32.32, -32.32, 48.32, 48.32, -20.51, -49.02, 36.51, 65.02, -49.02, -20.51, 65.02, 36.51, -42.8, -42.8, 58.8, 58.8, -27.92, -63.84, 43.92, 79.84, -63.84, -27.92, 79.84, 43.92},
+			{-48.0, -48.0, 80.0, 80.0, -29.25, -74.51, 61.25, 106.51, -74.51, -29.25, 106.51, 61.25, -64.63, -64.63, 96.63, 96.63, -41.02, -98.04, 73.02, 130.04, -98.04, -41.02, 130.04, 73.02, -85.59, -85.59, 117.59, 117.59, -55.84, -127.68, 87.84, 159.68, -127.68, -55.84, 159.68, 87.84},
+			{-96.0, -96.0, 160.0, 160.0, -58.51, -149.02, 122.51, 213.02, -149.02, -58.51, 213.02, 122.51, -129.27, -129.27, 193.27, 193.27, -82.04, -196.07, 146.04, 260.07, -196.07, -82.04, 260.07, 146.04, -171.19, -171.19, 235.19, 235.19, -111.68, -255.35, 175.68, 319.35, -255.35, -111.68, 319.35, 175.68},
+			{-192.0, -192.0, 320.0, 320.0, -117.02, -298.04, 245.02, 426.04, -298.04, -117.02, 426.04, 245.02, -258.54, -258.54, 386.54, 386.54, -164.07, -392.14, 292.07, 520.14, -392.14, -164.07, 520.14, 292.07, -342.37, -342.37, 470.37, 470.37, -223.35, -510.7, 351.35, 638.7, -510.7, -223.35, 638.7, 351.35}
+		};
+	}
+	else {
+		// Rotated-bboxes
+		anchors = {
+			{-12.0, 0.0, 19.0, 7.0, -7.0, -2.0, 14.0, 9.0, -4.0, -4.0, 11.0, 11.0, -2.0, -8.0, 9.0, 15.0, 0.0, -12.0, 7.0, 19.0, -21.4, -2.35, 28.4, 9.35, -13.46, -5.52, 20.46, 12.52, -8.7, -8.7, 15.7, 15.7, -5.52, -15.05, 12.52, 22.05, -2.35, -21.4, 9.35, 28.4, -36.32, -6.08, 43.32, 13.08, -23.72, -11.12, 30.72, 18.12, -16.16, -16.16, 23.16, 23.16, -11.12, -26.24, 18.12, 33.24, -6.08, -36.32, 13.08, 43.32, -12.0, 0.0, 19.0, 7.0, -7.0, -2.0, 14.0, 9.0, -4.0, -4.0, 11.0, 11.0, -2.0, -8.0, 9.0, 15.0, 0.0, -12.0, 7.0, 19.0, -21.4, -2.35, 28.4, 9.35, -13.46, -5.52, 20.46, 12.52, -8.7, -8.7, 15.7, 15.7, -5.52, -15.05, 12.52, 22.05, -2.35, -21.4, 9.35, 28.4, -36.32, -6.08, 43.32, 13.08, -23.72, -11.12, 30.72, 18.12, -16.16, -16.16, 23.16, 23.16, -11.12, -26.24, 18.12, 33.24, -6.08, -36.32, 13.08, 43.32, -12.0, 0.0, 19.0, 7.0, -7.0, -2.0, 14.0, 9.0, -4.0, -4.0, 11.0, 11.0, -2.0, -8.0, 9.0, 15.0, 0.0, -12.0, 7.0, 19.0, -21.4, -2.35, 28.4, 9.35, -13.46, -5.52, 20.46, 12.52, -8.7, -8.7, 15.7, 15.7, -5.52, -15.05, 12.52, 22.05, -2.35, -21.4, 9.35, 28.4, -36.32, -6.08, 43.32, 13.08, -23.72, -11.12, 30.72, 18.12, -16.16, -16.16, 23.16, 23.16, -11.12, -26.24, 18.12, 33.24, -6.08, -36.32, 13.08, 43.32},
+			{-24.0, 0.0, 39.0, 15.0, -15.0, -4.0, 30.0, 19.0, -8.0, -8.0, 23.0, 23.0, -3.0, -14.0, 18.0, 29.0, 0.0, -24.0, 15.0, 39.0, -42.8, -4.7, 57.8, 19.7, -28.51, -11.05, 43.51, 26.05, -17.4, -17.4, 32.4, 32.4, -9.46, -26.92, 24.46, 41.92, -4.7, -42.8, 19.7, 57.8, -72.63, -12.16, 87.63, 27.16, -49.96, -22.24, 64.96, 37.24, -32.32, -32.32, 47.32, 47.32, -19.72, -47.44, 34.72, 62.44, -12.16, -72.63, 27.16, 87.63, -24.0, 0.0, 39.0, 15.0, -15.0, -4.0, 30.0, 19.0, -8.0, -8.0, 23.0, 23.0, -3.0, -14.0, 18.0, 29.0, 0.0, -24.0, 15.0, 39.0, -42.8, -4.7, 57.8, 19.7, -28.51, -11.05, 43.51, 26.05, -17.4, -17.4, 32.4, 32.4, -9.46, -26.92, 24.46, 41.92, -4.7, -42.8, 19.7, 57.8, -72.63, -12.16, 87.63, 27.16, -49.96, -22.24, 64.96, 37.24, -32.32, -32.32, 47.32, 47.32, -19.72, -47.44, 34.72, 62.44, -12.16, -72.63, 27.16, 87.63, -24.0, 0.0, 39.0, 15.0, -15.0, -4.0, 30.0, 19.0, -8.0, -8.0, 23.0, 23.0, -3.0, -14.0, 18.0, 29.0, 0.0, -24.0, 15.0, 39.0, -42.8, -4.7, 57.8, 19.7, -28.51, -11.05, 43.51, 26.05, -17.4, -17.4, 32.4, 32.4, -9.46, -26.92, 24.46, 41.92, -4.7, -42.8, 19.7, 57.8, -72.63, -12.16, 87.63, 27.16, -49.96, -22.24, 64.96, 37.24, -32.32, -32.32, 47.32, 47.32, -19.72, -47.44, 34.72, 62.44, -12.16, -72.63, 27.16, 87.63},
+			{-48.0, 0.0, 79.0, 31.0, -29.0, -6.0, 60.0, 37.0, -16.0, -16.0, 47.0, 47.0, -7.0, -30.0, 38.0, 61.0, 0.0, -48.0, 31.0, 79.0, -85.59, -9.4, 116.59, 40.4, -55.43, -18.92, 86.43, 49.92, -34.8, -34.8, 65.8, 65.8, -20.51, -57.02, 51.51, 88.02, -9.4, -85.59, 40.4, 116.59, -145.27, -24.32, 176.27, 55.32, -97.39, -39.44, 128.39, 70.44, -64.63, -64.63, 95.63, 95.63, -41.96, -99.91, 72.96, 130.91, -24.32, -145.27, 55.32, 176.27, -48.0, 0.0, 79.0, 31.0, -29.0, -6.0, 60.0, 37.0, -16.0, -16.0, 47.0, 47.0, -7.0, -30.0, 38.0, 61.0, 0.0, -48.0, 31.0, 79.0, -85.59, -9.4, 116.59, 40.4, -55.43, -18.92, 86.43, 49.92, -34.8, -34.8, 65.8, 65.8, -20.51, -57.02, 51.51, 88.02, -9.4, -85.59, 40.4, 116.59, -145.27, -24.32, 176.27, 55.32, -97.39, -39.44, 128.39, 70.44, -64.63, -64.63, 95.63, 95.63, -41.96, -99.91, 72.96, 130.91, -24.32, -145.27, 55.32, 176.27, -48.0, 0.0, 79.0, 31.0, -29.0, -6.0, 60.0, 37.0, -16.0, -16.0, 47.0, 47.0, -7.0, -30.0, 38.0, 61.0, 0.0, -48.0, 31.0, 79.0, -85.59, -9.4, 116.59, 40.4, -55.43, -18.92, 86.43, 49.92, -34.8, -34.8, 65.8, 65.8, -20.51, -57.02, 51.51, 88.02, -9.4, -85.59, 40.4, 116.59, -145.27, -24.32, 176.27, 55.32, -97.39, -39.44, 128.39, 70.44, -64.63, -64.63, 95.63, 95.63, -41.96, -99.91, 72.96, 130.91, -24.32, -145.27, 55.32, 176.27},
+			{-96.0, 0.0, 159.0, 63.0, -59.0, -14.0, 122.0, 77.0, -32.0, -32.0, 95.0, 95.0, -13.0, -58.0, 76.0, 121.0, 0.0, -96.0, 63.0, 159.0, -171.19, -18.8, 234.19, 81.8, -112.45, -41.02, 175.45, 104.02, -69.59, -69.59, 132.59, 132.59, -39.43, -110.87, 102.43, 173.87, -18.8, -171.19, 81.8, 234.19, -290.54, -48.63, 353.54, 111.63, -197.31, -83.91, 260.31, 146.91, -129.27, -129.27, 192.27, 192.27, -81.39, -194.79, 144.39, 257.79, -48.63, -290.54, 111.63, 353.54, -96.0, 0.0, 159.0, 63.0, -59.0, -14.0, 122.0, 77.0, -32.0, -32.0, 95.0, 95.0, -13.0, -58.0, 76.0, 121.0, 0.0, -96.0, 63.0, 159.0, -171.19, -18.8, 234.19, 81.8, -112.45, -41.02, 175.45, 104.02, -69.59, -69.59, 132.59, 132.59, -39.43, -110.87, 102.43, 173.87, -18.8, -171.19, 81.8, 234.19, -290.54, -48.63, 353.54, 111.63, -197.31, -83.91, 260.31, 146.91, -129.27, -129.27, 192.27, 192.27, -81.39, -194.79, 144.39, 257.79, -48.63, -290.54, 111.63, 353.54, -96.0, 0.0, 159.0, 63.0, -59.0, -14.0, 122.0, 77.0, -32.0, -32.0, 95.0, 95.0, -13.0, -58.0, 76.0, 121.0, 0.0, -96.0, 63.0, 159.0, -171.19, -18.8, 234.19, 81.8, -112.45, -41.02, 175.45, 104.02, -69.59, -69.59, 132.59, 132.59, -39.43, -110.87, 102.43, 173.87, -18.8, -171.19, 81.8, 234.19, -290.54, -48.63, 353.54, 111.63, -197.31, -83.91, 260.31, 146.91, -129.27, -129.27, 192.27, 192.27, -81.39, -194.79, 144.39, 257.79, -48.63, -290.54, 111.63, 353.54},
+			{-192.0, 0.0, 319.0, 127.0, -117.0, -26.0, 244.0, 153.0, -64.0, -64.0, 191.0, 191.0, -27.0, -118.0, 154.0, 245.0, 0.0, -192.0, 127.0, 319.0, -342.37, -37.59, 469.37, 164.59, -223.32, -78.87, 350.32, 205.87, -139.19, -139.19, 266.19, 266.19, -80.45, -224.91, 207.45, 351.91, -37.59, -342.37, 164.59, 469.37, -581.08, -97.27, 708.08, 224.27, -392.09, -162.79, 519.09, 289.79, -258.54, -258.54, 385.54, 385.54, -165.31, -394.61, 292.31, 521.61, -97.27, -581.08, 224.27, 708.08, -192.0, 0.0, 319.0, 127.0, -117.0, -26.0, 244.0, 153.0, -64.0, -64.0, 191.0, 191.0, -27.0, -118.0, 154.0, 245.0, 0.0, -192.0, 127.0, 319.0, -342.37, -37.59, 469.37, 164.59, -223.32, -78.87, 350.32, 205.87, -139.19, -139.19, 266.19, 266.19, -80.45, -224.91, 207.45, 351.91, -37.59, -342.37, 164.59, 469.37, -581.08, -97.27, 708.08, 224.27, -392.09, -162.79, 519.09, 289.79, -258.54, -258.54, 385.54, 385.54, -165.31, -394.61, 292.31, 521.61, -97.27, -581.08, 224.27, 708.08, -192.0, 0.0, 319.0, 127.0, -117.0, -26.0, 244.0, 153.0, -64.0, -64.0, 191.0, 191.0, -27.0, -118.0, 154.0, 245.0, 0.0, -192.0, 127.0, 319.0, -342.37, -37.59, 469.37, 164.59, -223.32, -78.87, 350.32, 205.87, -139.19, -139.19, 266.19, 266.19, -80.45, -224.91, 207.45, 351.91, -37.59, -342.37, 164.59, 469.37, -581.08, -97.27, 708.08, 224.27, -392.09, -162.79, 519.09, 289.79, -258.54, -258.54, 385.54, 385.54, -165.31, -394.61, 292.31, 521.61, -97.27, -581.08, 224.27, 708.08}
+		};
+	}
+
+	// For INT8 calibration, after setting COCO_PATH on line 10:
+	// const vector<string> calibration_files = glob(calibration_batches*dynamic_batch_opts[1]);
+	const vector<string> calibration_files;
+	string model_name = "";
+	string calibration_table = argc == 4 ? string(argv[3]) : "";
+
+	// Use FP16 precision by default, use INT8 if calibration table is provided
+	string precision = "FP16";
+	if (argc == 4)
+		precision = "INT8";
+
+	cout << "Building engine..." << endl;
+	auto engine = odtk::Engine(buffer, size, dynamic_batch_opts, precision, score_thresh, top_n,
+		anchors, ROTATED, nms_thresh, detections_per_im, calibration_files, model_name, calibration_table, verbose, workspace_size);
+	engine.save(string(argv[2]));
+
+	delete [] buffer;
+
+	return 0;
+}

+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

+import torch
+from ._C import decode as decode_cuda
+from ._C import iou as iou_cuda
+from ._C import nms as nms_cuda
+import numpy as np
+from .utils import order_points, rotate_boxes
+
+def generate_anchors(stride, ratio_vals, scales_vals, angles_vals=None):
+    'Generate anchors coordinates from scales/ratios'
+
+    scales = torch.FloatTensor(scales_vals).repeat(len(ratio_vals), 1)
+    scales = scales.transpose(0, 1).contiguous().view(-1, 1)
+    ratios = torch.FloatTensor(ratio_vals * len(scales_vals))
+
+    wh = torch.FloatTensor([stride]).repeat(len(ratios), 2)
+    ws = torch.sqrt(wh[:, 0] * wh[:, 1] / ratios)
+    dwh = torch.stack([ws, ws * ratios], dim=1)
+    xy1 = 0.5 * (wh - dwh * scales)
+    xy2 = 0.5 * (wh + dwh * scales)
+    return torch.cat([xy1, xy2], dim=1)
+
+
+def generate_anchors_rotated(stride, ratio_vals, scales_vals, angles_vals):
+    'Generate anchors coordinates from scales/ratios/angles'
+    scales = torch.FloatTensor(scales_vals).repeat(len(ratio_vals), 1) 
+    scales = scales.transpose(0, 1).contiguous().view(-1, 1)
+    ratios = torch.FloatTensor(ratio_vals * len(scales_vals))
+
+    wh = torch.FloatTensor([stride]).repeat(len(ratios), 2)
+    ws = torch.round(torch.sqrt(wh[:, 0] * wh[:, 1] / ratios))
+    dwh = torch.stack([ws, torch.round(ws * ratios)], dim=1)
+    
+    xy0 = 0.5 * (wh - dwh * scales)
+    xy2 = 0.5 * (wh + dwh * scales) - 1
+    xy1 = xy0 + (xy2 - xy0) * torch.FloatTensor([0,1])
+    xy3 = xy0 + (xy2 - xy0) * torch.FloatTensor([1,0])
+    
+    angles = torch.FloatTensor(angles_vals)
+    theta = angles.repeat(xy0.size(0),1)
+    theta = theta.transpose(0,1).contiguous().view(-1,1)
+
+    xmin_ymin = xy0.repeat(int(theta.size(0)/xy0.size(0)),1)
+    xmax_ymax = xy2.repeat(int(theta.size(0)/xy2.size(0)),1)
+    widths_heights = dwh * scales
+    widths_heights = widths_heights.repeat(int(theta.size(0)/widths_heights.size(0)),1)
+
+    u = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)
+    l = torch.stack([-torch.sin(angles), torch.cos(angles)], dim=1)
+    R = torch.stack([u, l], dim=1)
+
+    xy0R = torch.matmul(R,xy0.transpose(1,0) - stride/2 + 0.5) + stride/2 - 0.5
+    xy1R = torch.matmul(R,xy1.transpose(1,0) - stride/2 + 0.5) + stride/2 - 0.5
+    xy2R = torch.matmul(R,xy2.transpose(1,0) - stride/2 + 0.5) + stride/2 - 0.5
+    xy3R = torch.matmul(R,xy3.transpose(1,0) - stride/2 + 0.5) + stride/2 - 0.5
+    
+    xy0R = xy0R.permute(0,2,1).contiguous().view(-1,2)
+    xy1R = xy1R.permute(0,2,1).contiguous().view(-1,2)
+    xy2R = xy2R.permute(0,2,1).contiguous().view(-1,2)
+    xy3R = xy3R.permute(0,2,1).contiguous().view(-1,2)
+
+    anchors_axis = torch.cat([xmin_ymin, xmax_ymax], dim=1)
+    anchors_rotated = order_points(torch.stack([xy0R,xy1R,xy2R,xy3R],dim = 1)).view(-1,8)
+
+    return anchors_axis, anchors_rotated
+
+
+def box2delta(boxes, anchors):
+    'Convert boxes to deltas from anchors'
+
+    anchors_wh = anchors[:, 2:] - anchors[:, :2] + 1
+    anchors_ctr = anchors[:, :2] + 0.5 * anchors_wh
+    boxes_wh = boxes[:, 2:] - boxes[:, :2] + 1
+    boxes_ctr = boxes[:, :2] + 0.5 * boxes_wh
+
+    return torch.cat([
+        (boxes_ctr - anchors_ctr) / anchors_wh,
+        torch.log(boxes_wh / anchors_wh)
+    ], 1)
+
+
+def box2delta_rotated(boxes, anchors):
+    'Convert boxes to deltas from anchors'
+
+    anchors_wh = anchors[:, 2:4] - anchors[:, :2] + 1
+    anchors_ctr = anchors[:, :2] + 0.5 * anchors_wh
+    boxes_wh = boxes[:, 2:4] - boxes[:, :2] + 1
+    boxes_ctr = boxes[:, :2] + 0.5 * boxes_wh
+    boxes_sin = boxes[:, 4]
+    boxes_cos = boxes[:, 5]
+
+    return torch.cat([
+        (boxes_ctr - anchors_ctr) / anchors_wh,
+        torch.log(boxes_wh / anchors_wh), boxes_sin[:, None], boxes_cos[:, None]
+    ], 1)
+
+
+def delta2box(deltas, anchors, size, stride):
+    'Convert deltas from anchors to boxes'
+
+    anchors_wh = anchors[:, 2:] - anchors[:, :2] + 1
+    ctr = anchors[:, :2] + 0.5 * anchors_wh
+    pred_ctr = deltas[:, :2] * anchors_wh + ctr
+    pred_wh = torch.exp(deltas[:, 2:]) * anchors_wh
+
+    m = torch.zeros([2], device=deltas.device, dtype=deltas.dtype)
+    M = (torch.tensor([size], device=deltas.device, dtype=deltas.dtype) * stride - 1)
+    clamp = lambda t: torch.max(m, torch.min(t, M))
+    return torch.cat([
+        clamp(pred_ctr - 0.5 * pred_wh),
+        clamp(pred_ctr + 0.5 * pred_wh - 1)
+    ], 1)
+
+
+def delta2box_rotated(deltas, anchors, size, stride):
+    'Convert deltas from anchors to boxes'
+
+    anchors_wh = anchors[:, 2:4] - anchors[:, :2] + 1
+    ctr = anchors[:, :2] + 0.5 * anchors_wh
+    pred_ctr = deltas[:, :2] * anchors_wh + ctr
+    pred_wh = torch.exp(deltas[:, 2:4]) * anchors_wh
+    pred_sin = deltas[:, 4]
+    pred_cos = deltas[:, 5]
+
+    m = torch.zeros([2], device=deltas.device, dtype=deltas.dtype)
+    M = (torch.tensor([size], device=deltas.device, dtype=deltas.dtype) * stride - 1)
+    clamp = lambda t: torch.max(m, torch.min(t, M))
+    return torch.cat([
+        clamp(pred_ctr - 0.5 * pred_wh),
+        clamp(pred_ctr + 0.5 * pred_wh - 1),
+        torch.atan2(pred_sin, pred_cos)[:, None]
+    ], 1)
+
+
+def snap_to_anchors(boxes, size, stride, anchors, num_classes, device, anchor_ious):
+    'Snap target boxes (x, y, w, h) to anchors'
+
+    num_anchors = anchors.size()[0] if anchors is not None else 1
+    width, height = (int(size[0] / stride), int(size[1] / stride))
+
+    if boxes.nelement() == 0:
+        return (torch.zeros([num_anchors, num_classes, height, width], device=device),
+                torch.zeros([num_anchors, 4, height, width], device=device),
+                torch.zeros([num_anchors, 1, height, width], device=device))
+
+    boxes, classes = boxes.split(4, dim=1)
+
+    # Generate anchors
+    x, y = torch.meshgrid([torch.arange(0, size[i], stride, device=device, dtype=classes.dtype) for i in range(2)])
+    xyxy = torch.stack((x, y, x, y), 2).unsqueeze(0)
+    anchors = anchors.view(-1, 1, 1, 4).to(dtype=classes.dtype)
+    anchors = (xyxy + anchors).contiguous().view(-1, 4)
+
+    # Compute overlap between boxes and anchors
+    boxes = torch.cat([boxes[:, :2], boxes[:, :2] + boxes[:, 2:] - 1], 1)
+    xy1 = torch.max(anchors[:, None, :2], boxes[:, :2])
+    xy2 = torch.min(anchors[:, None, 2:], boxes[:, 2:])
+    inter = torch.prod((xy2 - xy1 + 1).clamp(0), 2)
+    boxes_area = torch.prod(boxes[:, 2:] - boxes[:, :2] + 1, 1)
+    anchors_area = torch.prod(anchors[:, 2:] - anchors[:, :2] + 1, 1)
+    overlap = inter / (anchors_area[:, None] + boxes_area - inter)
+
+    # Keep best box per anchor
+    overlap, indices = overlap.max(1)
+    box_target = box2delta(boxes[indices], anchors)
+    box_target = box_target.view(num_anchors, 1, width, height, 4)
+    box_target = box_target.transpose(1, 4).transpose(2, 3)
+    box_target = box_target.squeeze().contiguous()
+
+    depth = torch.ones_like(overlap) * -1
+    depth[overlap < anchor_ious[0]] = 0  # background
+    depth[overlap >= anchor_ious[1]] = classes[indices][overlap >= anchor_ious[1]].squeeze() + 1  # objects
+    depth = depth.view(num_anchors, width, height).transpose(1, 2).contiguous()
+
+    # Generate target classes
+    cls_target = torch.zeros((anchors.size()[0], num_classes + 1), device=device, dtype=boxes.dtype)
+    if classes.nelement() == 0:
+        classes = torch.LongTensor([num_classes], device=device).expand_as(indices)
+    else:
+        classes = classes[indices].long()
+    classes = classes.view(-1, 1)
+    classes[overlap < anchor_ious[0]] = num_classes  # background has no class
+    cls_target.scatter_(1, classes, 1)
+    cls_target = cls_target[:, :num_classes].view(-1, 1, width, height, num_classes)
+    cls_target = cls_target.transpose(1, 4).transpose(2, 3)
+    cls_target = cls_target.squeeze().contiguous()
+
+    return (cls_target.view(num_anchors, num_classes, height, width),
+            box_target.view(num_anchors, 4, height, width),
+            depth.view(num_anchors, 1, height, width))
+
+
+def snap_to_anchors_rotated(boxes, size, stride, anchors, num_classes, device, anchor_ious):
+    'Snap target boxes (x, y, w, h, a) to anchors'
+
+    anchors_axis, anchors_rotated = anchors
+
+    num_anchors = anchors_rotated.size()[0] if anchors_rotated is not None else 1
+    width, height = (int(size[0] / stride), int(size[1] / stride))
+
+    if boxes.nelement() == 0:
+        return (torch.zeros([num_anchors, num_classes, height, width], device=device),
+                torch.zeros([num_anchors, 6, height, width], device=device),
+                torch.zeros([num_anchors, 1, height, width], device=device))
+
+    boxes, classes = boxes.split(5, dim=1)
+    boxes_axis, boxes_rotated = rotate_boxes(boxes)
+    
+    boxes_axis = boxes_axis.to(device)
+    boxes_rotated = boxes_rotated.to(device)
+    anchors_axis = anchors_axis.to(device)
+    anchors_rotated = anchors_rotated.to(device)
+
+    # Generate anchors
+    x, y = torch.meshgrid([torch.arange(0, size[i], stride, device=device, dtype=classes.dtype) for i in range(2)])
+    xy_2corners = torch.stack((x, y, x, y), 2).unsqueeze(0)
+    xy_4corners = torch.stack((x, y, x, y, x, y, x, y), 2).unsqueeze(0)
+    anchors_axis = (xy_2corners.to(torch.float) + anchors_axis.view(-1, 1, 1, 4)).contiguous().view(-1, 4)
+    anchors_rotated = (xy_4corners.to(torch.float) + anchors_rotated.view(-1, 1, 1, 8)).contiguous().view(-1, 8)
+
+    if torch.cuda.is_available():
+        iou = iou_cuda
+
+    overlap = iou(boxes_rotated.contiguous().view(-1), anchors_rotated.contiguous().view(-1))[0]
+
+    # Keep best box per anchor
+    overlap, indices = overlap.max(1)
+    box_target = box2delta_rotated(boxes_axis[indices], anchors_axis)
+    box_target = box_target.view(num_anchors, 1, width, height, 6)
+    box_target = box_target.transpose(1, 4).transpose(2, 3)
+    box_target = box_target.squeeze().contiguous()
+
+    depth = torch.ones_like(overlap, device=device) * -1
+    depth[overlap < anchor_ious[0]] = 0  # background
+    depth[overlap >= anchor_ious[1]] = classes[indices][overlap >= anchor_ious[1]].squeeze() + 1  # objects
+    depth = depth.view(num_anchors, width, height).transpose(1, 2).contiguous()
+
+    # Generate target classes
+    cls_target = torch.zeros((anchors_axis.size()[0], num_classes + 1), device=device, dtype=boxes_axis.dtype)
+    if classes.nelement() == 0:
+        classes = torch.LongTensor([num_classes], device=device).expand_as(indices)
+    else:
+        classes = classes[indices].long()
+    classes = classes.view(-1, 1)
+    classes[overlap < anchor_ious[0]] = num_classes  # background has no class
+    cls_target.scatter_(1, classes, 1)
+    cls_target = cls_target[:, :num_classes].view(-1, 1, width, height, num_classes)
+    cls_target = cls_target.transpose(1, 4).transpose(2, 3)
+    cls_target = cls_target.squeeze().contiguous()
+
+    return (cls_target.view(num_anchors, num_classes, height, width),
+            box_target.view(num_anchors, 6, height, width),
+            depth.view(num_anchors, 1, height, width))
+
+
+def decode(all_cls_head, all_box_head, stride=1, threshold=0.05, top_n=1000, anchors=None, rotated=False):
+    'Box Decoding and Filtering'
+
+    if rotated:
+        anchors = anchors[0]
+    num_boxes = 4 if not rotated else 6
+
+    if torch.cuda.is_available():
+        return decode_cuda(all_cls_head.float(), all_box_head.float(),
+            anchors.view(-1).tolist(), stride, threshold, top_n, rotated)
+
+    device = all_cls_head.device
+    anchors = anchors.to(device).type(all_cls_head.type())
+    num_anchors = anchors.size()[0] if anchors is not None else 1
+    num_classes = all_cls_head.size()[1] // num_anchors
+    height, width = all_cls_head.size()[-2:]
+
+    batch_size = all_cls_head.size()[0]
+    out_scores = torch.zeros((batch_size, top_n), device=device)
+    out_boxes = torch.zeros((batch_size, top_n, num_boxes), device=device)
+    out_classes = torch.zeros((batch_size, top_n), device=device)
+
+    # Per item in batch
+    for batch in range(batch_size):
+        cls_head = all_cls_head[batch, :, :, :].contiguous().view(-1)
+        box_head = all_box_head[batch, :, :, :].contiguous().view(-1, num_boxes)
+
+        # Keep scores over threshold
+        keep = (cls_head >= threshold).nonzero().view(-1)
+        if keep.nelement() == 0:
+            continue
+
+        # Gather top elements
+        scores = torch.index_select(cls_head, 0, keep)
+        scores, indices = torch.topk(scores, min(top_n, keep.size()[0]), dim=0)
+        indices = torch.index_select(keep, 0, indices).view(-1)
+        classes = (indices / width / height) % num_classes
+        classes = classes.type(all_cls_head.type())
+
+        # Infer kept bboxes
+        x = indices % width
+        y = (indices / width) % height
+        a = indices / num_classes / height / width
+        box_head = box_head.view(num_anchors, num_boxes, height, width)
+        boxes = box_head[a, :, y, x]
+
+        if anchors is not None:
+            grid = torch.stack([x, y, x, y], 1).type(all_cls_head.type()) * stride + anchors[a, :]
+            boxes = delta2box(boxes, grid, [width, height], stride)
+
+        out_scores[batch, :scores.size()[0]] = scores
+        out_boxes[batch, :boxes.size()[0], :] = boxes
+        out_classes[batch, :classes.size()[0]] = classes
+
+    return out_scores, out_boxes, out_classes
+
+
+def nms(all_scores, all_boxes, all_classes, nms=0.5, ndetections=100):
+    'Non Maximum Suppression'
+
+    if torch.cuda.is_available():
+        return nms_cuda(all_scores.float(), all_boxes.float(), all_classes.float(), 
+            nms, ndetections, False)
+
+    device = all_scores.device
+    batch_size = all_scores.size()[0]
+    out_scores = torch.zeros((batch_size, ndetections), device=device)
+    out_boxes = torch.zeros((batch_size, ndetections, 4), device=device)
+    out_classes = torch.zeros((batch_size, ndetections), device=device)
+
+    # Per item in batch
+    for batch in range(batch_size):
+        # Discard null scores
+        keep = (all_scores[batch, :].view(-1) > 0).nonzero()
+        scores = all_scores[batch, keep].view(-1)
+        boxes = all_boxes[batch, keep, :].view(-1, 4)
+        classes = all_classes[batch, keep].view(-1)
+
+        if scores.nelement() == 0:
+            continue
+
+        # Sort boxes
+        scores, indices = torch.sort(scores, descending=True)
+        boxes, classes = boxes[indices], classes[indices]
+        areas = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1).view(-1)
+        keep = torch.ones(scores.nelement(), device=device, dtype=torch.uint8).view(-1)
+
+        for i in range(ndetections):
+            if i >= keep.nonzero().nelement() or i >= scores.nelement():
+                i -= 1
+                break
+
+            # Find overlapping boxes with lower score
+            xy1 = torch.max(boxes[:, :2], boxes[i, :2])
+            xy2 = torch.min(boxes[:, 2:], boxes[i, 2:])
+            inter = torch.prod((xy2 - xy1 + 1).clamp(0), 1)
+            criterion = ((scores > scores[i]) |
+                         (inter / (areas + areas[i] - inter) <= nms) |
+                         (classes != classes[i]))
+            criterion[i] = 1
+
+            # Only keep relevant boxes
+            scores = scores[criterion.nonzero()].view(-1)
+            boxes = boxes[criterion.nonzero(), :].view(-1, 4)
+            classes = classes[criterion.nonzero()].view(-1)
+            areas = areas[criterion.nonzero()].view(-1)
+            keep[(~criterion).nonzero()] = 0
+
+        out_scores[batch, :i + 1] = scores[:i + 1]
+        out_boxes[batch, :i + 1, :] = boxes[:i + 1, :]
+        out_classes[batch, :i + 1] = classes[:i + 1]
+
+    return out_scores, out_boxes, out_classes
+
+
+def nms_rotated(all_scores, all_boxes, all_classes, nms=0.5, ndetections=100):
+    'Non Maximum Suppression'
+
+    if torch.cuda.is_available():
+        return nms_cuda(all_scores.float(), all_boxes.float(), all_classes.float(), 
+            nms, ndetections, True)
+
+    device = all_scores.device
+    batch_size = all_scores.size()[0]
+    out_scores = torch.zeros((batch_size, ndetections), device=device)
+    out_boxes = torch.zeros((batch_size, ndetections, 6), device=device)
+    out_classes = torch.zeros((batch_size, ndetections), device=device)
+
+    # Per item in batch
+    for batch in range(batch_size):
+        # Discard null scores
+        keep = (all_scores[batch, :].view(-1) > 0).nonzero()
+        scores = all_scores[batch, keep].view(-1)
+        boxes = all_boxes[batch, keep, :].view(-1, 6)
+        classes = all_classes[batch, keep].view(-1)
+        theta = torch.atan2(boxes[:, -2], boxes[:, -1])
+        boxes_theta = torch.cat([boxes[:, :-2], theta[:, None]], dim=1)
+
+        if scores.nelement() == 0:
+            continue
+
+        # Sort boxes
+        scores, indices = torch.sort(scores, descending=True)
+        boxes, boxes_theta, classes = boxes[indices], boxes_theta[indices], classes[indices]
+        areas = (boxes_theta[:, 2] - boxes_theta[:, 0] + 1) * (boxes_theta[:, 3] - boxes_theta[:, 1] + 1).view(-1)
+        keep = torch.ones(scores.nelement(), device=device, dtype=torch.uint8).view(-1)
+
+        for i in range(ndetections):
+            if i >= keep.nonzero().nelement() or i >= scores.nelement():
+                i -= 1
+                break
+
+            boxes_axis, boxes_rotated = rotate_boxes(boxes_theta, points=True)
+            overlap, inter = iou(boxes_rotated.contiguous().view(-1), boxes_rotated[i, :].contiguous().view(-1))
+            inter = inter.squeeze()
+            criterion = ((scores > scores[i]) |
+                         (inter / (areas + areas[i] - inter) <= nms) |
+                         (classes != classes[i]))
+            criterion[i] = 1
+
+            # Only keep relevant boxes
+            scores = scores[criterion.nonzero()].view(-1)
+            boxes = boxes[criterion.nonzero(), :].view(-1, 6)
+            boxes_theta = boxes_theta[criterion.nonzero(), :].view(-1, 5)
+            classes = classes[criterion.nonzero()].view(-1)
+            areas = areas[criterion.nonzero()].view(-1)
+            keep[(~criterion).nonzero()] = 0
+
+        out_scores[batch, :i + 1] = scores[:i + 1]
+        out_boxes[batch, :i + 1, :] = boxes[:i + 1, :]
+        out_classes[batch, :i + 1] = classes[:i + 1]
+
+    return out_scores, out_boxes, out_classes

+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 random
+from contextlib import redirect_stdout
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from torch.utils import data
+from pycocotools.coco import COCO
+import math
+from torchvision.transforms.functional import adjust_brightness, adjust_contrast, adjust_hue, adjust_saturation
+
+
+class CocoDataset(data.dataset.Dataset):
+    'Dataset looping through a set of images'
+
+    def __init__(self, path, resize, max_size, stride, annotations=None, training=False, rotate_augment=False,
+                 augment_brightness=0.0, augment_contrast=0.0,
+                 augment_hue=0.0, augment_saturation=0.0):
+        super().__init__()
+
+        self.path = os.path.expanduser(path)
+        self.resize = resize
+        self.max_size = max_size
+        self.stride = stride
+        self.mean = [0.485, 0.456, 0.406]
+        self.std = [0.229, 0.224, 0.225]
+        self.training = training
+        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
+
+        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())}
+
+    def __len__(self):
+        return len(self.ids)
+
+    def __getitem__(self, index):
+        ' Get sample'
+
+        # Load image
+        id = self.ids[index]
+        if self.coco:
+            image = self.coco.loadImgs(id)[0]['file_name']
+        im = Image.open('{}/{}'.format(self.path, image)).convert("RGB")
+
+        # Randomly sample scale for resize during training
+        resize = self.resize
+        if isinstance(resize, list):
+            resize = random.randint(self.resize[0], self.resize[-1])
+
+        ratio = resize / min(im.size)
+        if ratio * max(im.size) > self.max_size:
+            ratio = self.max_size / max(im.size)
+        im = im.resize((int(ratio * d) for d in im.size), Image.BILINEAR)
+
+        if self.training:
+            # Get annotations
+            boxes, categories = self._get_target(id)
+            boxes *= ratio
+
+            # Random rotation, if self.rotate_augment
+            random_angle = random.randint(0, 3) * 90
+            if self.rotate_augment and random_angle != 0:
+                # rotate by random_angle degrees.
+                im = im.rotate(random_angle)
+                x, y, w, h = boxes[:, 0].clone(), boxes[:, 1].clone(), boxes[:, 2].clone(), boxes[:, 3].clone()
+                if random_angle == 90:
+                    boxes[:, 0] = y - im.size[1] / 2 + im.size[0] / 2
+                    boxes[:, 1] = im.size[0] / 2 + im.size[1] / 2 - x - w
+                    boxes[:, 2] = h
+                    boxes[:, 3] = w
+                elif random_angle == 180:
+                    boxes[:, 0] = im.size[0] - x - w
+                    boxes[:, 1] = im.size[1] - y - h
+                elif random_angle == 270:
+                    boxes[:, 0] = im.size[0] / 2 + im.size[1] / 2 - y - h
+                    boxes[:, 1] = x - im.size[0] / 2 + im.size[1] / 2
+                    boxes[:, 2] = h
+                    boxes[:, 3] = w
+
+            # Random horizontal flip
+            if random.randint(0, 1):
+                im = im.transpose(Image.FLIP_LEFT_RIGHT)
+                boxes[:, 0] = im.size[0] - boxes[:, 0] - boxes[:, 2]
+
+            # Apply image brightness, contrast etc augmentation
+            if self.augment_brightness:
+                brightness_factor = random.normalvariate(1, self.augment_brightness)
+                brightness_factor = max(0, brightness_factor)
+                im = adjust_brightness(im, brightness_factor)
+            if self.augment_contrast:
+                contrast_factor = random.normalvariate(1, self.augment_contrast)
+                contrast_factor = max(0, contrast_factor)
+                im = adjust_contrast(im, contrast_factor)
+            if self.augment_hue:
+                hue_factor = random.normalvariate(0, self.augment_hue)
+                hue_factor = max(-0.5, hue_factor)
+                hue_factor = min(0.5, hue_factor)
+                im = adjust_hue(im, hue_factor)
+            if self.augment_saturation:
+                saturation_factor = random.normalvariate(1, self.augment_saturation)
+                saturation_factor = max(0, saturation_factor)
+                im = adjust_saturation(im, saturation_factor)
+
+            target = torch.cat([boxes, categories], dim=1)
+
+        # Convert to tensor and normalize
+        data = torch.ByteTensor(torch.ByteStorage.from_buffer(im.tobytes()))
+        data = data.float().div(255).view(*im.size[::-1], len(im.mode))
+        data = data.permute(2, 0, 1)
+
+        for t, mean, std in zip(data, self.mean, self.std):
+            t.sub_(mean).div_(std)
+
+        # Apply padding
+        pw, ph = ((self.stride - d % self.stride) % self.stride for d in im.size)
+        data = F.pad(data, (0, pw, 0, ph))
+
+        if self.training:
+            return data, target
+
+        return data, id, ratio
+
+    def _get_target(self, id):
+        'Get annotations for sample'
+
+        ann_ids = self.coco.getAnnIds(imgIds=id)
+        annotations = self.coco.loadAnns(ann_ids)
+
+        boxes, categories = [], []
+        for ann in annotations:
+            if ann['bbox'][2] < 1 and ann['bbox'][3] < 1:
+                continue
+            boxes.append(ann['bbox'])
+            cat = ann['category_id']
+            if 'categories' in self.coco.dataset:
+                cat = self.categories_inv[cat]
+            categories.append(cat)
+
+        if boxes:
+            target = (torch.FloatTensor(boxes),
+                      torch.FloatTensor(categories).unsqueeze(1))
+        else:
+            target = (torch.ones([1, 4]), torch.ones([1, 1]) * -1)
+
+        return target
+
+    def collate_fn(self, batch):
+        'Create batch from multiple samples'
+
+        if self.training:
+            data, targets = zip(*batch)
+            max_det = max([t.size()[0] for t in targets])
+            targets = [torch.cat([t, torch.ones([max_det - t.size()[0], 5]) * -1]) for t in targets]
+            targets = torch.stack(targets, 0)
+        else:
+            data, indices, ratios = zip(*batch)
+
+        # Pad data to match max batch dimensions
+        sizes = [d.size()[-2:] for d in data]
+        w, h = (max(dim) for dim in zip(*sizes))
+
+        data_stack = []
+        for datum in data:
+            pw, ph = w - datum.size()[-2], h - datum.size()[-1]
+            data_stack.append(
+                F.pad(datum, (0, ph, 0, pw)) if max(ph, pw) > 0 else datum)
+
+        data = torch.stack(data_stack)
+
+        if self.training:
+            return data, targets
+
+        ratios = torch.FloatTensor(ratios).view(-1, 1, 1)
+        return data, torch.IntTensor(indices), ratios
+
+
+class DataIterator():
+    'Data loader for data parallel'
+
+    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.resize = resize
+        self.max_size = max_size
+
+        self.dataset = CocoDataset(path, resize=resize, max_size=max_size,
+                                   stride=stride, annotations=annotations, training=training,
+                                   rotate_augment=rotate_augment,
+                                   augment_brightness=augment_brightness,
+                                   augment_contrast=augment_contrast, augment_hue=augment_hue,
+                                   augment_saturation=augment_saturation)
+        self.ids = self.dataset.ids
+        self.coco = self.dataset.coco
+
+        self.sampler = data.distributed.DistributedSampler(self.dataset) if world > 1 else None
+        self.dataloader = data.DataLoader(self.dataset, batch_size=batch_size // world,
+                                          sampler=self.sampler, collate_fn=self.dataset.collate_fn, num_workers=2,
+                                          pin_memory=True)
+
+    def __repr__(self):
+        return '\n'.join([
+            '    loader: pytorch',
+            '    resize: {}, max: {}'.format(self.resize, self.max_size),
+        ])
+
+    def __len__(self):
+        return len(self.dataloader)
+
+    def __iter__(self):
+        for output in self.dataloader:
+            if self.dataset.training:
+                data, target = output
+            else:
+                data, ids, ratio = output
+
+            if torch.cuda.is_available():
+                data = data.cuda(non_blocking=True)
+
+            if self.dataset.training:
+                if torch.cuda.is_available():
+                    target = target.cuda(non_blocking=True)
+                yield data, target
+            else:
+                if torch.cuda.is_available():
+                    ids = ids.cuda(non_blocking=True)
+                    ratio = ratio.cuda(non_blocking=True)
+                yield data, ids, ratio
+
+
+class RotatedCocoDataset(data.dataset.Dataset):
+    'Dataset looping through a set of images'
+
+    def __init__(self, path, resize, max_size, stride, annotations=None, training=False, rotate_augment=False,
+                 augment_brightness=0.0, augment_contrast=0.0,
+                 augment_hue=0.0, augment_saturation=0.0, absolute_angle=False):
+        super().__init__()
+
+        self.path = os.path.expanduser(path)
+        self.resize = resize
+        self.max_size = max_size
+        self.stride = stride
+        self.mean = [0.485, 0.456, 0.406]
+        self.std = [0.229, 0.224, 0.225]
+        self.training = training
+        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.absolute_angle=absolute_angle
+
+        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())}
+
+    def __len__(self):
+        return len(self.ids)
+
+    def __getitem__(self, index):
+        ' Get sample'
+
+        # Load image
+        id = self.ids[index]
+        if self.coco:
+            image = self.coco.loadImgs(id)[0]['file_name']
+        im = Image.open('{}/{}'.format(self.path, image)).convert("RGB")
+
+        # Randomly sample scale for resize during training
+        resize = self.resize
+        if isinstance(resize, list):
+            resize = random.randint(self.resize[0], self.resize[-1])
+
+        ratio = resize / min(im.size)
+        if ratio * max(im.size) > self.max_size:
+            ratio = self.max_size / max(im.size)
+        im = im.resize((int(ratio * d) for d in im.size), Image.BILINEAR)
+
+        if self.training:
+            # Get annotations
+            boxes, categories = self._get_target(id)
+            # boxes *= ratio
+            boxes[:, :4] *= ratio
+
+            # Random rotation, if self.rotate_augment
+            random_angle = random.randint(0, 3) * 90
+            if self.rotate_augment and random_angle != 0:
+                # rotate by random_angle degrees.
+                original_size = im.size
+                im = im.rotate(random_angle, expand=True)
+                x, y, w, h, t = boxes[:, 0].clone(), boxes[:, 1].clone(), boxes[:, 2].clone(), \
+                                boxes[:, 3].clone(), boxes[:, 4].clone()
+                if random_angle == 90:
+                    boxes[:, 0] = y
+                    boxes[:, 1] = original_size[0] - x - w
+                    if not self.absolute_angle:
+                        boxes[:, 2] = h
+                        boxes[:, 3] = w
+                elif random_angle == 180:
+                    boxes[:, 0] = original_size[0] - x - w
+                    boxes[:, 1] = original_size[1] - y - h
+
+                elif random_angle == 270:
+                    boxes[:, 0] = original_size[1] - y - h
+                    boxes[:, 1] = x
+                    if not self.absolute_angle:
+                        boxes[:, 2] = h
+                        boxes[:, 3] = w
+
+                    pass
+
+                # Adjust theta
+                if self.absolute_angle:
+                    # This is only needed in absolute angle mode.
+                    t += math.radians(random_angle)
+                    rem = torch.remainder(torch.abs(t), math.pi)
+                    sign = torch.sign(t)
+                    t = rem * sign
+
+                boxes[:, 4] = t
+
+            # Random horizontal flip
+            if random.randint(0, 1):
+                im = im.transpose(Image.FLIP_LEFT_RIGHT)
+                boxes[:, 0] = im.size[0] - boxes[:, 0] - boxes[:, 2]
+                boxes[:, 1] = boxes[:, 1]
+                boxes[:, 4] = -boxes[:, 4]
+
+            # Apply image brightness, contrast etc augmentation
+            if self.augment_brightness:
+                brightness_factor = random.normalvariate(1, self.augment_brightness)
+                brightness_factor = max(0, brightness_factor)
+                im = adjust_brightness(im, brightness_factor)
+            if self.augment_contrast:
+                contrast_factor = random.normalvariate(1, self.augment_contrast)
+                contrast_factor = max(0, contrast_factor)
+                im = adjust_contrast(im, contrast_factor)
+            if self.augment_hue:
+                hue_factor = random.normalvariate(0, self.augment_hue)
+                hue_factor = max(-0.5, hue_factor)
+                hue_factor = min(0.5, hue_factor)
+                im = adjust_hue(im, hue_factor)
+            if self.augment_saturation:
+                saturation_factor = random.normalvariate(1, self.augment_saturation)
+                saturation_factor = max(0, saturation_factor)
+                im = adjust_saturation(im, saturation_factor)
+
+            target = torch.cat([boxes, categories], dim=1)
+
+        # Convert to tensor and normalize
+        data = torch.ByteTensor(torch.ByteStorage.from_buffer(im.tobytes()))
+        data = data.float().div(255).view(*im.size[::-1], len(im.mode))
+        data = data.permute(2, 0, 1)
+
+        for t, mean, std in zip(data, self.mean, self.std):
+            t.sub_(mean).div_(std)
+
+        # Apply padding
+        pw, ph = ((self.stride - d % self.stride) % self.stride for d in im.size)
+        data = F.pad(data, (0, pw, 0, ph))
+
+        if self.training:
+            return data, target
+
+        return data, id, ratio
+
+    def _get_target(self, id):
+        'Get annotations for sample'
+
+        ann_ids = self.coco.getAnnIds(imgIds=id)
+        annotations = self.coco.loadAnns(ann_ids)
+
+        boxes, categories = [], []
+        for ann in annotations:
+            if ann['bbox'][2] < 1 and ann['bbox'][3] < 1:
+                continue
+            final_bbox = ann['bbox']
+            if len(final_bbox) == 4:
+                final_bbox.append(0.0)  # add theta of zero.
+            assert len(ann['bbox']) == 5, "Bounding box for id %i does not contain five entries." % id
+            boxes.append(final_bbox)
+            cat = ann['category_id']
+            if 'categories' in self.coco.dataset:
+                cat = self.categories_inv[cat]
+            categories.append(cat)
+
+        if boxes:
+            target = (torch.FloatTensor(boxes),
+                      torch.FloatTensor(categories).unsqueeze(1))
+        else:
+            target = (torch.ones([1, 5]), torch.ones([1, 1]) * -1)
+
+        return target
+
+    def collate_fn(self, batch):
+        'Create batch from multiple samples'
+
+        if self.training:
+            data, targets = zip(*batch)
+            max_det = max([t.size()[0] for t in targets])
+            targets = [torch.cat([t, torch.ones([max_det - t.size()[0], 6]) * -1]) for t in targets]
+            targets = torch.stack(targets, 0)
+        else:
+            data, indices, ratios = zip(*batch)
+
+        # Pad data to match max batch dimensions
+        sizes = [d.size()[-2:] for d in data]
+        w, h = (max(dim) for dim in zip(*sizes))
+
+        data_stack = []
+        for datum in data:
+            pw, ph = w - datum.size()[-2], h - datum.size()[-1]
+            data_stack.append(
+                F.pad(datum, (0, ph, 0, pw)) if max(ph, pw) > 0 else datum)
+
+        data = torch.stack(data_stack)
+
+        if self.training:
+            return data, targets
+
+        ratios = torch.FloatTensor(ratios).view(-1, 1, 1)
+        return data, torch.IntTensor(indices), ratios
+
+
+class RotatedDataIterator():
+    'Data loader for data parallel'
+
+    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, absolute_angle=False
+                 ):
+        self.resize = resize
+        self.max_size = max_size
+
+        self.dataset = RotatedCocoDataset(path, resize=resize, max_size=max_size,
+                                          stride=stride, annotations=annotations, training=training,
+                                          rotate_augment=rotate_augment,
+                                          augment_brightness=augment_brightness,
+                                          augment_contrast=augment_contrast, augment_hue=augment_hue,
+                                          augment_saturation=augment_saturation, absolute_angle=absolute_angle)
+        self.ids = self.dataset.ids
+        self.coco = self.dataset.coco
+
+        self.sampler = data.distributed.DistributedSampler(self.dataset) if world > 1 else None
+        self.dataloader = data.DataLoader(self.dataset, batch_size=batch_size // world,
+                                          sampler=self.sampler, collate_fn=self.dataset.collate_fn, num_workers=2,
+                                          pin_memory=True)
+
+    def __repr__(self):
+        return '\n'.join([
+            '    loader: pytorch',
+            '    resize: {}, max: {}'.format(self.resize, self.max_size),
+        ])
+
+    def __len__(self):
+        return len(self.dataloader)
+
+    def __iter__(self):
+        for output in self.dataloader:
+            if self.dataset.training:
+                data, target = output
+            else:
+                data, ids, ratio = output
+
+            if torch.cuda.is_available():
+                data = data.cuda(non_blocking=True)
+
+            if self.dataset.training:
+                if torch.cuda.is_available():
+                    target = target.cuda(non_blocking=True)
+                yield data, target
+            else:
+                if torch.cuda.is_available():
+                    ids = ids.cuda(non_blocking=True)
+                    ratio = ratio.cuda(non_blocking=True)
+                yield data, ids, ratio

+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)

+#!/usr/bin/env python3
+import sys
+import os
+import argparse
+import random
+import torch.cuda
+import torch.distributed
+import torch.multiprocessing
+
+from odtk import infer, train, utils
+from odtk.model import Model
+from odtk._C import Engine
+
+
+def parse(args):
+    parser = argparse.ArgumentParser(description='ODTK: Object Detection Toolkit.')
+    parser.add_argument('--master', metavar='address:port', type=str, help='Address and port of the master worker',
+                        default='127.0.0.1:29500')
+
+    subparsers = parser.add_subparsers(help='sub-command', dest='command')
+    subparsers.required = True
+
+    devcount = max(1, torch.cuda.device_count())
+
+    parser_train = subparsers.add_parser('train', help='train a network')
+    parser_train.add_argument('model', type=str, help='path to output model or checkpoint to resume from')
+    parser_train.add_argument('--annotations', metavar='path', type=str, help='path to COCO style annotations',
+                              required=True)
+    parser_train.add_argument('--images', metavar='path', type=str, help='path to images', default='.')
+    parser_train.add_argument('--backbone', action='store', type=str, nargs='+', help='backbone model (or list of)',
+                              default=['ResNet50FPN'])
+    parser_train.add_argument('--classes', metavar='num', type=int, help='number of classes', default=80)
+    parser_train.add_argument('--batch', metavar='size', type=int, help='batch size', default=2 * devcount)
+    parser_train.add_argument('--resize', metavar='scale', type=int, help='resize to given size', default=800)
+    parser_train.add_argument('--max-size', metavar='max', type=int, help='maximum resizing size', default=1333)
+    parser_train.add_argument('--jitter', metavar='min max', type=int, nargs=2, help='jitter size within range',
+                              default=[640, 1024])
+    parser_train.add_argument('--iters', metavar='number', type=int, help='number of iterations to train for',
+                              default=90000)
+    parser_train.add_argument('--milestones', action='store', type=int, nargs='*',
+                              help='list of iteration indices where learning rate decays', default=[60000, 80000])
+    parser_train.add_argument('--schedule', metavar='scale', type=float,
+                              help='scale schedule (affecting iters and milestones)', default=1)
+    parser_train.add_argument('--full-precision', help='train in full precision', action='store_true')
+    parser_train.add_argument('--lr', metavar='value', help='learning rate', type=float, default=0.01)
+    parser_train.add_argument('--warmup', metavar='iterations', help='numer of warmup iterations', type=int,
+                              default=1000)
+    parser_train.add_argument('--gamma', metavar='value', type=float,
+                              help='multiplicative factor of learning rate decay', default=0.1)
+    parser_train.add_argument('--override', help='override model', action='store_true')
+    parser_train.add_argument('--val-annotations', metavar='path', type=str,
+                              help='path to COCO style validation annotations')
+    parser_train.add_argument('--val-images', metavar='path', type=str, help='path to validation images')
+    parser_train.add_argument('--post-metrics', metavar='url', type=str, help='post metrics to specified url')
+    parser_train.add_argument('--fine-tune', metavar='path', type=str, help='fine tune a pretrained model')
+    parser_train.add_argument('--logdir', metavar='logdir', type=str, help='directory where to write logs')
+    parser_train.add_argument('--val-iters', metavar='number', type=int,
+                              help='number of iterations between each validation', default=8000)
+    parser_train.add_argument('--no-apex', help='use Pytorch native AMP and DDP', action='store_true')
+    parser_train.add_argument('--with-dali', help='use dali for data loading', action='store_true')
+    parser_train.add_argument('--augment-rotate', help='use four-fold rotational augmentation', action='store_true')
+    parser_train.add_argument('--augment-free-rotate', type=float, metavar='value value', nargs=2, default=[0, 0],
+                              help='rotate images by an arbitrary angle, between min and max (in degrees)')
+    parser_train.add_argument('--augment-brightness', metavar='value', type=float,
+                              help='adjust the brightness of the image.', default=0.002)
+    parser_train.add_argument('--augment-contrast', metavar='value', type=float,
+                              help='adjust the contrast of the image.', default=0.002)
+    parser_train.add_argument('--augment-hue', metavar='value', type=float,
+                              help='adjust the hue of the image.', default=0.0002)
+    parser_train.add_argument('--augment-saturation', metavar='value', type=float,
+                              help='adjust the saturation of the image.', default=0.002)
+    parser_train.add_argument('--regularization-l2', metavar='value', type=float, help='L2 regularization for optim',
+                              default=0.0001)
+    parser_train.add_argument('--rotated-bbox', help='detect rotated bounding boxes [x, y, w, h, theta]',
+                              action='store_true')
+    parser_train.add_argument('--anchor-ious', metavar='value value', type=float, nargs=2,
+                              help='anchor/bbox overlap threshold', default=[0.4, 0.5])
+    parser_train.add_argument('--absolute-angle', help='regress absolute angle (rather than -45 to 45 degrees.',
+                              action='store_true')
+
+    parser_infer = subparsers.add_parser('infer', help='run inference')
+    parser_infer.add_argument('model', type=str, help='path to model')
+    parser_infer.add_argument('--images', metavar='path', type=str, help='path to images', default='.')
+    parser_infer.add_argument('--annotations', metavar='annotations', type=str,
+                              help='evaluate using provided annotations')
+    parser_infer.add_argument('--output', metavar='file', type=str, nargs='+',
+                              help='save detections to specified JSON file(s)', default=['detections.json'])
+    parser_infer.add_argument('--batch', metavar='size', type=int, help='batch size', default=2 * devcount)
+    parser_infer.add_argument('--resize', metavar='scale', type=int, help='resize to given size', default=800)
+    parser_infer.add_argument('--max-size', metavar='max', type=int, help='maximum resizing size', default=1333)
+    parser_infer.add_argument('--no-apex', help='use Pytorch native AMP and DDP', action='store_true')
+    parser_infer.add_argument('--with-dali', help='use dali for data loading', action='store_true')
+    parser_infer.add_argument('--full-precision', help='inference in full precision', action='store_true')
+    parser_infer.add_argument('--rotated-bbox', help='inference using a rotated bounding box model',
+                              action='store_true')
+
+    parser_export = subparsers.add_parser('export', help='export a model into a TensorRT engine')
+    parser_export.add_argument('model', type=str, help='path to model')
+    parser_export.add_argument('export', type=str, help='path to exported output')
+    parser_export.add_argument('--size', metavar='height width', type=int, nargs='+',
+                               help='input size (square) or sizes (h w) to use when generating TensorRT engine',
+                               default=[1280])
+    parser_export.add_argument('--full-precision', help='export in full instead of half precision', action='store_true')
+    parser_export.add_argument('--int8', help='calibrate model and export in int8 precision', action='store_true')
+    parser_export.add_argument('--calibration-batches', metavar='size', type=int,
+                               help='number of batches to use for int8 calibration', default=2)
+    parser_export.add_argument('--calibration-images', metavar='path', type=str,
+                               help='path to calibration images to use for int8 calibration', default="")
+    parser_export.add_argument('--calibration-table', metavar='path', type=str,
+                               help='path of existing calibration table to load from, or name of new calibration table',
+                               default="")
+    parser_export.add_argument('--verbose', help='enable verbose logging', action='store_true')
+    parser_export.add_argument('--rotated-bbox', help='inference using a rotated bounding box model',
+                               action='store_true')
+    parser_export.add_argument('--dynamic-batch-opts', help='Profile batch sizes for tensorrt engine export (min, opt, max)',
+                               metavar='value value value', type=int, nargs=3, default=[1,8,16])
+
+    return parser.parse_args(args)
+
+
+def load_model(args, verbose=False):
+    if args.command != 'train' and not os.path.isfile(args.model):
+        raise RuntimeError('Model file {} does not exist!'.format(args.model))
+
+    model = None
+    state = {}
+    _, ext = os.path.splitext(args.model)
+
+    if args.command == 'train' and (not os.path.exists(args.model) or args.override):
+        if verbose: print('Initializing model...')
+        model = Model(backbones=args.backbone, classes=args.classes, rotated_bbox=args.rotated_bbox,
+                      anchor_ious=args.anchor_ious)
+        model.initialize(args.fine_tune)
+        # Freeze unused params from training
+        for n, p in model.named_parameters():
+            if any(i in n for i in model.unused_modules):
+                p.requires_grad = False
+        if verbose: print(model)
+
+    elif ext == '.pth' or ext == '.torch':
+        if verbose: print('Loading model from {}...'.format(os.path.basename(args.model)))
+        model, state = Model.load(filename=args.model, rotated_bbox=args.rotated_bbox)
+        if verbose: print(model)
+
+    elif args.command == 'infer' and ext in ['.engine', '.plan']:
+        model = None
+
+    else:
+        raise RuntimeError('Invalid model format "{}"!'.format(ext))
+
+    state['path'] = args.model
+    return model, state
+
+
+def worker(rank, args, world, model, state):
+    'Per-device distributed worker'
+
+    if torch.cuda.is_available():
+        os.environ.update({
+            'MASTER_PORT': args.master.split(':')[-1],
+            'MASTER_ADDR': ':'.join(args.master.split(':')[:-1]),
+            'WORLD_SIZE': str(world),
+            'RANK': str(rank),
+            'CUDA_DEVICE': str(rank)
+        })
+
+        torch.cuda.set_device(rank)
+        torch.distributed.init_process_group(backend='nccl', init_method='env://')
+
+        if (args.command != 'export') and (args.batch % world != 0):
+            raise RuntimeError('Batch size should be a multiple of the number of GPUs')
+
+    if model and model.angles is not None:
+        args.rotated_bbox = True
+
+    if args.command == 'train':
+        train.train(model, state, args.images, args.annotations,
+                    args.val_images or args.images, args.val_annotations, args.resize, args.max_size, args.jitter,
+                    args.batch, int(args.iters * args.schedule), args.val_iters, not args.full_precision, args.lr,
+                    args.warmup, [int(m * args.schedule) for m in args.milestones], args.gamma,
+                    rank, world=world, no_apex=args.no_apex, use_dali=args.with_dali,
+                    metrics_url=args.post_metrics, logdir=args.logdir, verbose=(rank == 0),
+                    rotate_augment=args.augment_rotate,
+                    augment_brightness=args.augment_brightness, augment_contrast=args.augment_contrast,
+                    augment_hue=args.augment_hue, augment_saturation=args.augment_saturation,
+                    regularization_l2=args.regularization_l2, rotated_bbox=args.rotated_bbox, absolute_angle=args.absolute_angle)
+
+    elif args.command == 'infer':
+        if model is None:
+            if rank == 0: print('Loading CUDA engine from {}...'.format(os.path.basename(args.model)))
+            model = Engine.load(args.model)
+
+        infer.infer(model, args.images, args.output, args.resize, args.max_size, args.batch,
+                    annotations=args.annotations, mixed_precision=not args.full_precision,
+                    is_master=(rank == 0), world=world, no_apex=args.no_apex, use_dali=args.with_dali,
+                    verbose=(rank == 0), rotated_bbox=args.rotated_bbox)
+
+    elif args.command == 'export':
+        onnx_only = args.export.split('.')[-1] == 'onnx'
+        input_size = args.size * 2 if len(args.size) == 1 else args.size
+
+        calibration_files = []
+        if args.int8:
+            # Get list of images to use for calibration
+            if os.path.isdir(args.calibration_images):
+                import glob
+                file_extensions = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG']
+                for ex in file_extensions:
+                    calibration_files += glob.glob("{}/*{}".format(args.calibration_images, ex), recursive=True)
+                # Only need enough images for specified num of calibration batches
+                if len(calibration_files) >= args.calibration_batches * args.dynamic_batch_opts[1]:
+                    calibration_files = calibration_files[:(args.calibration_batches * args.dynamic_batch_opts[1])]
+                else:
+                    # Number of images for calibration must be greater than or equal to the kOPT optimization profile
+                    if len(calibration_files) >= args.dynamic_batch_opts[1]:
+                        print('Only found enough images for {} batches. Continuing anyway...'.format(
+                            len(calibration_files) // args.dynamic_batch_opts[1]))
+                    else:
+                        raise RuntimeError('Not enough images found for calibration. ({} < {})'
+                                            .format(len(calibration_files), args.dynamic_batch_opts[1]))
+
+                random.shuffle(calibration_files)
+
+        precision = "FP32"
+        if args.int8:
+            precision = "INT8"
+        elif not args.full_precision:
+            precision = "FP16"
+
+        exported = model.export(input_size, args.dynamic_batch_opts, precision, calibration_files,
+                                args.calibration_table, args.verbose, onnx_only=onnx_only)
+        if onnx_only:
+            with open(args.export, 'wb') as out:
+                out.write(exported)
+        else:
+            exported.save(args.export)
+
+
+def main(args=None):
+    'Entry point for the odtk command'
+
+    args = parse(args or sys.argv[1:])
+
+    model, state = load_model(args, verbose=True)
+    if model: model.share_memory()
+
+    world = torch.cuda.device_count()
+    if args.command == 'export' or world <= 1:
+        worker(0, args, 1, model, state)
+    else:
+        torch.multiprocessing.spawn(worker, args=(args, world, model, state), nprocs=world)
+
+
+if __name__ == '__main__':
+    main()

+import os.path
+import io
+import numpy as np
+import math
+import torch
+import torch.nn as nn
+
+from . import backbones as backbones_mod
+from ._C import Engine
+from .box import generate_anchors, snap_to_anchors, decode, nms
+from .box import generate_anchors_rotated, snap_to_anchors_rotated, nms_rotated
+from .loss import FocalLoss, SmoothL1Loss
+
+
+class Model(nn.Module):
+    'RetinaNet - https://arxiv.org/abs/1708.02002'
+
+    def __init__(
+        self, 
+        backbones='ResNet50FPN', 
+        classes=80, 
+        ratios=[1.0, 2.0, 0.5], 
+        scales=[4 * 2 ** (i / 3) for i in range(3)],
+        angles=None, 
+        rotated_bbox=False, 
+        anchor_ious=[0.4, 0.5], 
+        config={}
+    ):
+        super().__init__()
+
+        if not isinstance(backbones, list):
+            backbones = [backbones]
+
+        self.backbones = nn.ModuleDict({b: getattr(backbones_mod, b)() for b in backbones})
+        self.name = 'RetinaNet'
+        self.unused_modules = []
+        for b in backbones: self.unused_modules.extend(getattr(self.backbones, b).features.unused_modules)
+        self.exporting = False
+        self.rotated_bbox = rotated_bbox
+        self.anchor_ious = anchor_ious
+
+        self.ratios = ratios
+        self.scales = scales
+        self.angles = angles if angles is not None else \
+                    [-np.pi / 6, 0, np.pi / 6] if self.rotated_bbox else None
+        self.anchors = {}
+        self.classes = classes
+
+        self.threshold = config.get('threshold', 0.05)
+        self.top_n = config.get('top_n', 1000)
+        self.nms = config.get('nms', 0.5)
+        self.detections = config.get('detections', 100)
+
+        self.stride = max([b.stride for _, b in self.backbones.items()])
+
+        # classification and box regression heads
+        def make_head(out_size):
+            layers = []
+            for _ in range(4):
+                layers += [nn.Conv2d(256, 256, 3, padding=1), nn.ReLU()]
+            layers += [nn.Conv2d(256, out_size, 3, padding=1)]
+            return nn.Sequential(*layers)
+
+        self.num_anchors = len(self.ratios) * len(self.scales)
+        self.num_anchors = self.num_anchors if not self.rotated_bbox else (self.num_anchors * len(self.angles))
+        self.cls_head = make_head(classes * self.num_anchors)
+        self.box_head = make_head(4 * self.num_anchors) if not self.rotated_bbox \
+                        else make_head(6 * self.num_anchors)  # theta -> cos(theta), sin(theta)
+
+        self.cls_criterion = FocalLoss()
+        self.box_criterion = SmoothL1Loss(beta=0.11)
+
+    def __repr__(self):
+        return '\n'.join([
+            '     model: {}'.format(self.name),
+            '  backbone: {}'.format(', '.join([k for k, _ in self.backbones.items()])),
+            '   classes: {}, anchors: {}'.format(self.classes, self.num_anchors)
+        ])
+
+    def initialize(self, pre_trained):
+        if pre_trained:
+            # Initialize using weights from pre-trained model
+            if not os.path.isfile(pre_trained):
+                raise ValueError('No checkpoint {}'.format(pre_trained))
+
+            print('Fine-tuning weights from {}...'.format(os.path.basename(pre_trained)))
+            state_dict = self.state_dict()
+            chk = torch.load(pre_trained, map_location=lambda storage, loc: storage)
+            ignored = ['cls_head.8.bias', 'cls_head.8.weight']
+            if self.rotated_bbox:
+                ignored += ['box_head.8.bias', 'box_head.8.weight']
+            weights = {k: v for k, v in chk['state_dict'].items() if k not in ignored}
+            state_dict.update(weights)
+            self.load_state_dict(state_dict)
+
+            del chk, weights
+            torch.cuda.empty_cache()
+
+        else:
+            # Initialize backbone(s)
+            for _, backbone in self.backbones.items():
+                backbone.initialize()
+
+            # Initialize heads
+            def initialize_layer(layer):
+                if isinstance(layer, nn.Conv2d):
+                    nn.init.normal_(layer.weight, std=0.01)
+                    if layer.bias is not None:
+                        nn.init.constant_(layer.bias, val=0)
+
+            self.cls_head.apply(initialize_layer)
+            self.box_head.apply(initialize_layer)
+
+        # Initialize class head prior
+        def initialize_prior(layer):
+            pi = 0.01
+            b = - math.log((1 - pi) / pi)
+            nn.init.constant_(layer.bias, b)
+            nn.init.normal_(layer.weight, std=0.01)
+
+        self.cls_head[-1].apply(initialize_prior)
+        if self.rotated_bbox:
+            self.box_head[-1].apply(initialize_prior)
+
+    def forward(self, x, rotated_bbox=None):
+        if self.training: x, targets = x
+
+        # Backbones forward pass
+        features = []
+        for _, backbone in self.backbones.items():
+            features.extend(backbone(x))
+
+        # Heads forward pass
+        cls_heads = [self.cls_head(t) for t in features]
+        box_heads = [self.box_head(t) for t in features]
+
+        if self.training:
+            return self._compute_loss(x, cls_heads, box_heads, targets.float())
+
+        cls_heads = [cls_head.sigmoid() for cls_head in cls_heads]
+
+        if self.exporting:
+            self.strides = [x.shape[-1] // cls_head.shape[-1] for cls_head in cls_heads]
+            return cls_heads, box_heads
+
+        global nms, generate_anchors
+        if self.rotated_bbox:
+            nms = nms_rotated
+            generate_anchors = generate_anchors_rotated
+
+        # Inference post-processing
+        decoded = []
+        for cls_head, box_head in zip(cls_heads, box_heads):
+            # Generate level's anchors
+            stride = x.shape[-1] // cls_head.shape[-1]
+            if stride not in self.anchors:
+                self.anchors[stride] = generate_anchors(stride, self.ratios, self.scales, self.angles)
+
+            # Decode and filter boxes
+            decoded.append(decode(cls_head.contiguous(), box_head.contiguous(), stride, self.threshold, 
+                                self.top_n, self.anchors[stride], self.rotated_bbox))
+
+        # Perform non-maximum suppression
+        decoded = [torch.cat(tensors, 1) for tensors in zip(*decoded)]
+        return nms(*decoded, self.nms, self.detections)
+
+    def _extract_targets(self, targets, stride, size):
+        global generate_anchors, snap_to_anchors
+        if self.rotated_bbox:
+            generate_anchors = generate_anchors_rotated
+            snap_to_anchors = snap_to_anchors_rotated
+        cls_target, box_target, depth = [], [], []
+        for target in targets:
+            target = target[target[:, -1] > -1]
+            if stride not in self.anchors:
+                self.anchors[stride] = generate_anchors(stride, self.ratios, self.scales, self.angles)
+
+            anchors = self.anchors[stride]
+            if not self.rotated_bbox:
+                anchors = anchors.to(targets.device)
+            snapped = snap_to_anchors(target, [s * stride for s in size[::-1]], stride, 
+                                    anchors, self.classes, targets.device, self.anchor_ious)
+            for l, s in zip((cls_target, box_target, depth), snapped): l.append(s)
+        return torch.stack(cls_target), torch.stack(box_target), torch.stack(depth)
+
+    def _compute_loss(self, x, cls_heads, box_heads, targets):
+        cls_losses, box_losses, fg_targets = [], [], []
+        for cls_head, box_head in zip(cls_heads, box_heads):
+            size = cls_head.shape[-2:]
+            stride = x.shape[-1] / cls_head.shape[-1]
+
+            cls_target, box_target, depth = self._extract_targets(targets, stride, size)
+            fg_targets.append((depth > 0).sum().float().clamp(min=1))
+
+            cls_head = cls_head.view_as(cls_target).float()
+            cls_mask = (depth >= 0).expand_as(cls_target).float()
+            cls_loss = self.cls_criterion(cls_head, cls_target)
+            cls_loss = cls_mask * cls_loss
+            cls_losses.append(cls_loss.sum())
+
+            box_head = box_head.view_as(box_target).float()
+            box_mask = (depth > 0).expand_as(box_target).float()
+            box_loss = self.box_criterion(box_head, box_target)
+            box_loss = box_mask * box_loss
+            box_losses.append(box_loss.sum())
+
+        fg_targets = torch.stack(fg_targets).sum()
+        cls_loss = torch.stack(cls_losses).sum() / fg_targets
+        box_loss = torch.stack(box_losses).sum() / fg_targets
+        return cls_loss, box_loss
+
+    def save(self, state):
+        checkpoint = {
+            'backbone': [k for k, _ in self.backbones.items()],
+            'classes': self.classes,
+            'state_dict': self.state_dict(),
+            'ratios': self.ratios,
+            'scales': self.scales
+        }
+        if self.rotated_bbox and self.angles:
+            checkpoint['angles'] = self.angles
+
+        for key in ('iteration', 'optimizer', 'scheduler'):
+            if key in state:
+                checkpoint[key] = state[key]
+
+        torch.save(checkpoint, state['path'])
+
+    @classmethod
+    def load(cls, filename, rotated_bbox=False):
+        if not os.path.isfile(filename):
+            raise ValueError('No checkpoint {}'.format(filename))
+
+        checkpoint = torch.load(filename, map_location=lambda storage, loc: storage)
+        kwargs = {}
+        for i in ['ratios', 'scales', 'angles']:
+            if i in checkpoint:
+                kwargs[i] = checkpoint[i]
+        if ('angles' in checkpoint) or rotated_bbox:
+            kwargs['rotated_bbox'] = True
+        # Recreate model from checkpoint instead of from individual backbones
+        model = cls(backbones=checkpoint['backbone'], classes=checkpoint['classes'], **kwargs)
+        model.load_state_dict(checkpoint['state_dict'])
+
+        state = {}
+        for key in ('iteration', 'optimizer', 'scheduler'):
+            if key in checkpoint:
+                state[key] = checkpoint[key]
+
+        del checkpoint
+        torch.cuda.empty_cache()
+
+        return model, state
+
+    def export(self, size, dynamic_batch_opts, precision, calibration_files, calibration_table, verbose, onnx_only=False):
+
+        # import torch.onnx.symbolic_opset11 as onnx_symbolic
+        # def upsample_nearest2d(g, input, output_size, *args):
+        #     # Currently, TRT 7.1 ONNX Parser does not support all ONNX ops
+        #     # needed to support dynamic upsampling ONNX forumlation
+        #     # Here we hardcode scale=2 as a temporary workaround
+        #     scales = g.op("Constant", value_t=torch.tensor([1., 1., 2., 2.]))
+        #     empty_tensor = g.op("Constant", value_t=torch.tensor([], dtype=torch.float32))
+        #     return g.op("Resize", input, empty_tensor, scales, mode_s="nearest", nearest_mode_s="floor")
+
+        # onnx_symbolic.upsample_nearest2d = upsample_nearest2d
+
+        # Export to ONNX
+        print('Exporting to ONNX...')
+        self.exporting = True
+        onnx_bytes = io.BytesIO()
+        zero_input = torch.zeros([1, 3, *size]).cuda()
+        input_names = ['input_1']
+        output_names = ['score_1', 'score_2', 'score_3', 'score_4', 'score_5',
+                        'box_1', 'box_2', 'box_3', 'box_4', 'box_5']
+        dynamic_axes = {input_names[0]: {0:'batch'}}
+        for _, name in enumerate(output_names):
+            dynamic_axes[name] = dynamic_axes[input_names[0]]
+        extra_args = {'opset_version': 12, 'verbose': verbose,
+                      'input_names': input_names, 'output_names': output_names,
+                      'dynamic_axes': dynamic_axes} 
+        torch.onnx.export(self.cuda(), zero_input, onnx_bytes, **extra_args)
+        self.exporting = False
+
+        if onnx_only:
+            return onnx_bytes.getvalue()
+
+        # Build TensorRT engine
+        model_name = '_'.join([k for k, _ in self.backbones.items()])
+        anchors = []
+        if not self.rotated_bbox:
+            anchors = [generate_anchors(stride, self.ratios, self.scales, 
+                    self.angles).view(-1).tolist() for stride in self.strides]
+        else:
+            anchors = [generate_anchors_rotated(stride, self.ratios, self.scales, 
+                    self.angles)[0].view(-1).tolist() for stride in self.strides]
+
+        return Engine(onnx_bytes.getvalue(), len(onnx_bytes.getvalue()), dynamic_batch_opts, precision,
+                      self.threshold, self.top_n, anchors, self.rotated_bbox, self.nms, self.detections,
+                      calibration_files, model_name, calibration_table, verbose)

+from statistics import mean
+from math import isfinite
+import torch
+from torch.optim import SGD, AdamW
+from torch.optim.lr_scheduler import LambdaLR, SAVE_STATE_WARNING
+from apex import amp, optimizers
+from apex.parallel import DistributedDataParallel as ADDP
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.cuda.amp import GradScaler, autocast
+from .backbones.layers import convert_fixedbn_model
+
+from .data import DataIterator, RotatedDataIterator
+from .dali import DaliDataIterator
+from .utils import ignore_sigint, post_metrics, Profiler
+from .infer import infer
+
+import warnings
+warnings.filterwarnings('ignore', message=SAVE_STATE_WARNING, category=UserWarning)
+
+
+def train(model, state, path, annotations, val_path, val_annotations, resize, max_size, jitter, batch_size, iterations,
+          val_iterations, mixed_precision, lr, warmup, milestones, gamma, rank=0, world=1, no_apex=False, use_dali=True,
+          verbose=True, metrics_url=None, logdir=None, rotate_augment=False, augment_brightness=0.0,
+          augment_contrast=0.0, augment_hue=0.0, augment_saturation=0.0, regularization_l2=0.0001, rotated_bbox=False,
+          absolute_angle=False):
+    'Train the model on the given dataset'
+
+    # Prepare model
+    nn_model = model
+    stride = model.stride
+
+    model = convert_fixedbn_model(model)
+    if torch.cuda.is_available():
+        model = model.to(memory_format=torch.channels_last).cuda()
+
+    # Setup optimizer and schedule
+    optimizer = SGD(model.parameters(), lr=lr, weight_decay=regularization_l2, momentum=0.9)
+
+    is_master = rank==0
+    if not no_apex:
+        loss_scale = "dynamic" if use_dali else "128.0"
+        model, optimizer = amp.initialize(model, optimizer,
+                                        opt_level='O2' if mixed_precision else 'O0',
+                                        keep_batchnorm_fp32=True,
+                                        loss_scale=loss_scale,
+                                        verbosity=is_master)
+
+    if world > 1:
+        model = DDP(model, device_ids=[rank]) if no_apex else ADDP(model)
+    model.train()
+
+    if 'optimizer' in state:
+        optimizer.load_state_dict(state['optimizer'])
+
+    def schedule(train_iter):
+        if warmup and train_iter <= warmup:
+            return 0.9 * train_iter / warmup + 0.1
+        return gamma ** len([m for m in milestones if m <= train_iter])
+
+    scheduler = LambdaLR(optimizer, schedule)
+    if 'scheduler' in state:
+        scheduler.load_state_dict(state['scheduler'])
+
+    # 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 detections.")
+        data_iterator = RotatedDataIterator(path, jitter, max_size, batch_size, stride,
+                                            world, annotations, training=True, rotate_augment=rotate_augment,
+                                            augment_brightness=augment_brightness,
+                                            augment_contrast=augment_contrast, augment_hue=augment_hue,
+                                            augment_saturation=augment_saturation, absolute_angle=absolute_angle)
+    else:
+        data_iterator = (DaliDataIterator if use_dali else DataIterator)(
+            path, jitter, max_size, batch_size, stride,
+            world, annotations, training=True, rotate_augment=rotate_augment, augment_brightness=augment_brightness,
+            augment_contrast=augment_contrast, augment_hue=augment_hue, augment_saturation=augment_saturation)
+    if verbose: print(data_iterator)
+
+    if verbose:
+        print('    device: {} {}'.format(
+            world, 'cpu' if not torch.cuda.is_available() else 'GPU' if world == 1 else 'GPUs'))
+        print('     batch: {}, precision: {}'.format(batch_size, 'mixed' if mixed_precision else 'full'))
+        print(' BBOX type:', 'rotated' if rotated_bbox else 'axis aligned')
+        print('Training model for {} iterations...'.format(iterations))
+
+    # Create TensorBoard writer
+    if is_master and logdir is not None:
+        from torch.utils.tensorboard import SummaryWriter
+        if verbose:
+            print('Writing TensorBoard logs to: {}'.format(logdir))
+        writer = SummaryWriter(log_dir=logdir)
+
+    scaler = GradScaler()
+    profiler = Profiler(['train', 'fw', 'bw'])
+    iteration = state.get('iteration', 0)
+    while iteration < iterations:
+        cls_losses, box_losses = [], []
+        for i, (data, target) in enumerate(data_iterator):
+            if iteration>=iterations:
+                break
+
+            # Forward pass
+            profiler.start('fw')
+
+            optimizer.zero_grad()
+            if not no_apex:
+                cls_loss, box_loss = model([data.contiguous(memory_format=torch.channels_last), target])
+            else:
+                with autocast():
+                    cls_loss, box_loss = model([data.contiguous(memory_format=torch.channels_last), target])
+            del data
+            profiler.stop('fw')
+
+            # Backward pass
+            profiler.start('bw')
+            if not no_apex:
+                with amp.scale_loss(cls_loss + box_loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
+                optimizer.step()
+            else:
+                scaler.scale(cls_loss + box_loss).backward()
+                scaler.step(optimizer)
+                scaler.update()
+
+            scheduler.step()
+
+            # Reduce all losses
+            cls_loss, box_loss = cls_loss.mean().clone(), box_loss.mean().clone()
+            if world > 1:
+                torch.distributed.all_reduce(cls_loss)
+                torch.distributed.all_reduce(box_loss)
+                cls_loss /= world
+                box_loss /= world
+            if is_master:
+                cls_losses.append(cls_loss)
+                box_losses.append(box_loss)
+
+            if is_master and not isfinite(cls_loss + box_loss):
+                raise RuntimeError('Loss is diverging!\n{}'.format(
+                    'Try lowering the learning rate.'))
+
+            del cls_loss, box_loss
+            profiler.stop('bw')
+
+            iteration += 1
+            profiler.bump('train')
+            if is_master and (profiler.totals['train'] > 60 or iteration == iterations):
+                focal_loss = torch.stack(list(cls_losses)).mean().item()
+                box_loss = torch.stack(list(box_losses)).mean().item()
+                learning_rate = optimizer.param_groups[0]['lr']
+                if verbose:
+                    msg = '[{:{len}}/{}]'.format(iteration, iterations, len=len(str(iterations)))
+                    msg += ' focal loss: {:.3f}'.format(focal_loss)
+                    msg += ', box loss: {:.3f}'.format(box_loss)
+                    msg += ', {:.3f}s/{}-batch'.format(profiler.means['train'], batch_size)
+                    msg += ' (fw: {:.3f}s, bw: {:.3f}s)'.format(profiler.means['fw'], profiler.means['bw'])
+                    msg += ', {:.1f} im/s'.format(batch_size / profiler.means['train'])
+                    msg += ', lr: {:.2g}'.format(learning_rate)
+                    print(msg, flush=True)
+
+                if is_master and logdir is not None:
+                    writer.add_scalar('focal_loss', focal_loss, iteration)
+                    writer.add_scalar('box_loss', box_loss, iteration)
+                    writer.add_scalar('learning_rate', learning_rate, iteration)
+                    del box_loss, focal_loss
+
+                if metrics_url:
+                    post_metrics(metrics_url, {
+                        'focal loss': mean(cls_losses),
+                        'box loss': mean(box_losses),
+                        'im_s': batch_size / profiler.means['train'],
+                        'lr': learning_rate
+                    })
+
+                # Save model weights
+                state.update({
+                    'iteration': iteration,
+                    'optimizer': optimizer.state_dict(),
+                    'scheduler': scheduler.state_dict(),
+                })
+                with ignore_sigint():
+                    nn_model.save(state)
+
+                profiler.reset()
+                del cls_losses[:], box_losses[:]
+
+            if val_annotations and (iteration == iterations or iteration % val_iterations == 0):
+                stats = infer(model, val_path, None, resize, max_size, batch_size, annotations=val_annotations,
+                            mixed_precision=mixed_precision, is_master=is_master, world=world, use_dali=use_dali,
+                            no_apex=no_apex, is_validation=True, verbose=False, rotated_bbox=rotated_bbox)
+                model.train()
+                if is_master and logdir is not None and stats is not None:
+                    writer.add_scalar(
+                        'Validation_Precision/mAP', stats[0], iteration)
+                    writer.add_scalar(
+                        'Validation_Precision/mAP@0.50IoU', stats[1], iteration)
+                    writer.add_scalar(
+                        'Validation_Precision/mAP@0.75IoU', stats[2], iteration)
+                    writer.add_scalar(
+                        'Validation_Precision/mAP (small)', stats[3], iteration)
+                    writer.add_scalar(
+                        'Validation_Precision/mAP (medium)', stats[4], iteration)
+                    writer.add_scalar(
+                        'Validation_Precision/mAP (large)', stats[5], iteration)
+                    writer.add_scalar(
+                        'Validation_Recall/mAR (max 1 Dets)', stats[6], iteration)
+                    writer.add_scalar(
+                        'Validation_Recall/mAR (max 10 Dets)', stats[7], iteration)
+                    writer.add_scalar(
+                        'Validation_Recall/mAR (max 100 Dets)', stats[8], iteration)
+                    writer.add_scalar(
+                        'Validation_Recall/mAR (small)', stats[9], iteration)
+                    writer.add_scalar(
+                        'Validation_Recall/mAR (medium)', stats[10], iteration)
+                    writer.add_scalar(
+                        'Validation_Recall/mAR (large)', stats[11], iteration)
+
+            if (iteration==iterations and not rotated_bbox) or (iteration>iterations and rotated_bbox):
+                break
+
+    if is_master and logdir is not None:
+        writer.close()

+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"}
+```

+192.168.0.101_20190720-11362588_Motion
+192.168.0.101_20190720-11372337_Motion
+192.168.0.101_20190720-11374138_Motion
+192.168.0.101_20190720-11452888_Motion
+192.168.0.101_20190720-12105034_Motion
+192.168.0.101_20190720-12160838_Motion
+192.168.0.101_20190720-12170737_Motion
+192.168.0.101_20190720-12175937_Motion
+192.168.0.101_20190720-12214737_Motion
+192.168.0.101_20190720-12235083_Motion
+192.168.0.101_20190720-12244383_Motion
+192.168.0.101_20190720-12283436_Motion
+192.168.0.101_20190720-13080033_Motion
+192.168.0.101_20190720-13111980_Motion
+192.168.0.101_20190720-13182131_Motion
+192.168.0.101_20190720-15191163_Motion
+192.168.0.101_20190720-15413012_Motion
+192.168.0.101_20190720-16465005_Motion
+192.168.0.101_20190720-16583605_Motion
+192.168.0.101_20190720-17234402_Motion
+192.168.0.101_20190720-17443700_Motion
+192.168.0.101_20190720-17502000_Motion
+192.168.0.101_20190720-18152897_Motion
+192.168.0.101_20190720-19041042_Motion
+192.168.0.101_20190721-06271810_Motion
+192.168.0.101_20190721-08570042_Motion
+192.168.0.101_20190721-09163489_Motion
+192.168.0.101_20190721-09345536_Motion
+192.168.0.101_20190721-09581284_Motion
+192.168.0.101_20190721-10460577_Motion
+192.168.0.101_20190721-10492125_Motion
+192.168.0.101_20190721-13324903_Motion
+192.168.0.101_20190721-14053747_Motion
+192.168.0.101_20190721-15081190_Motion
+192.168.0.101_20190721-15240738_Motion
+192.168.0.101_20190721-15280487_Motion
+192.168.0.101_20190721-15335936_Motion
+192.168.0.101_20190721-16074833_Motion
+192.168.0.101_20190721-16284528_Motion
+192.168.0.101_20190721-17213572_Motion
+192.168.0.101_20190721-17371770_Motion
+192.168.0.101_20190721-17383420_Motion
+192.168.0.101_20190721-18123715_Motion
+192.168.0.101_20190721-19404804_Motion
+192.168.0.101_20190722-06242119_Motion
+192.168.0.101_20190722-06535264_Motion
+192.168.0.101_20190722-07063613_Motion
+192.168.0.101_20190722-07160110_Motion
+192.168.0.101_20190722-07223611_Motion
+192.168.0.101_20190722-07282159_Motion
+192.168.0.101_20190722-07285409_Motion
+192.168.0.101_20190722-07290708_Motion
+192.168.0.101_20190722-07302108_Motion
+192.168.0.101_20190722-07343258_Motion
+192.168.0.101_20190722-07364109_Motion
+192.168.0.101_20190722-07382807_Motion
+192.168.0.101_20190722-07395107_Motion
+192.168.0.101_20190722-07460456_Motion
+192.168.0.101_20190722-07463858_Motion
+192.168.0.101_20190722-07483156_Motion
+192.168.0.101_20190722-07522694_Motion
+192.168.0.101_20190722-07542938_Motion
+192.168.0.101_20190722-07580427_Motion
+192.168.0.101_20190722-08042508_Motion
+192.168.0.101_20190722-08095192_Motion
+192.168.0.101_20190722-08120435_Motion
+192.168.0.101_20190722-08123134_Motion
+192.168.0.101_20190722-08125783_Motion
+192.168.0.101_20190722-08191565_Motion
+192.168.0.101_20190722-08263256_Motion
+192.168.0.101_20190722-08274306_Motion
+192.168.0.101_20190722-08514048_Motion
+192.168.0.101_20190722-08532449_Motion
+192.168.0.101_20190722-09201242_Motion
+192.168.0.101_20190722-09230544_Motion
+192.168.0.101_20190722-09310393_Motion
+192.168.0.101_20190722-09505438_Motion
+192.168.0.101_20190722-10032136_Motion
+192.168.0.101_20190722-10071237_Motion
+192.168.0.101_20190722-10253882_Motion
+192.168.0.101_20190722-10471330_Motion
+192.168.0.101_20190722-11091289_Motion
+192.168.0.101_20190722-11092887_Motion
+192.168.0.101_20190722-11132426_Motion
+192.168.0.101_20190722-11300481_Motion
+192.168.0.101_20190722-11324728_Motion
+192.168.0.101_20190722-11381481_Motion
+192.168.0.101_20190722-11383430_Motion
+192.168.0.101_20190722-11460626_Motion
+192.168.0.101_20190722-11471926_Motion
+192.168.0.101_20190722-11570325_Motion
+192.168.0.101_20190722-12124523_Motion
+192.168.0.101_20190722-12190523_Motion
+192.168.0.101_20190722-13103041_Motion
+192.168.0.101_20190722-13190967_Motion
+192.168.0.101_20190722-13544612_Motion
+192.168.0.101_20190722-13582111_Motion
+192.168.0.101_20190722-14015211_Motion
+192.168.0.101_20190722-14030761_Motion
+192.168.0.101_20190722-14292457_Motion
+192.168.0.101_20190722-14295705_Motion
+192.168.0.101_20190722-15021452_Motion
+192.168.0.101_20190722-15425543_Motion
+192.168.0.101_20190722-15430593_Motion
+192.168.0.101_20190722-15464694_Motion
+192.168.0.101_20190722-15500793_Motion
+192.168.0.101_20190722-15540538_Motion
+192.168.0.101_20190722-15562932_Motion
+192.168.0.101_20190722-15595921_Motion
+192.168.0.101_20190722-16032910_Motion
+192.168.0.101_20190722-16034659_Motion
+192.168.0.101_20190722-16280535_Motion
+192.168.0.101_20190722-16350515_Motion
+192.168.0.101_20190722-16353862_Motion
+192.168.0.101_20190722-16400558_Motion
+192.168.0.101_20190722-16564458_Motion
+192.168.0.101_20190722-16590705_Motion
+192.168.0.101_20190722-17092304_Motion
+192.168.0.101_20190722-17340899_Motion
+192.168.0.101_20190722-17512647_Motion
+192.168.0.101_20190722-17540845_Motion
+192.168.0.101_20190722-17593146_Motion
+192.168.0.101_20190722-18131094_Motion
+192.168.0.101_20190722-18310790_Motion
+192.168.0.101_20190722-18333591_Motion
+192.168.0.101_20190722-19010888_Motion
+192.168.0.101_20190722-19083933_Motion
+192.168.0.101_20190722-19090232_Motion
+192.168.0.101_20190722-19123431_Motion
+192.168.0.101_20190722-19135831_Motion
+192.168.0.101_20190723-06162619_Motion
+192.168.0.101_20190723-06505211_Motion
+192.168.0.101_20190723-07045162_Motion
+192.168.0.101_20190723-07150360_Motion
+192.168.0.101_20190723-07163007_Motion
+192.168.0.101_20190723-07203855_Motion
+192.168.0.101_20190723-07230356_Motion
+192.168.0.101_20190723-07282654_Motion
+192.168.0.101_20190723-07284053_Motion
+192.168.0.101_20190723-07340155_Motion
+192.168.0.101_20190723-07345454_Motion
+192.168.0.101_20190723-07382153_Motion
+192.168.0.101_20190723-07391451_Motion
+192.168.0.101_20190723-07441501_Motion
+192.168.0.101_20190723-07462601_Motion
+192.168.0.101_20190723-07473003_Motion
+192.168.0.101_20190723-07482502_Motion
+192.168.0.101_20190723-07495650_Motion
+192.168.0.101_20190723-07505498_Motion
+192.168.0.101_20190723-07521200_Motion
+192.168.0.101_20190723-08022747_Motion
+192.168.0.101_20190723-08173296_Motion
+192.168.0.101_20190723-08212896_Motion
+192.168.0.101_20190723-08365541_Motion
+192.168.0.101_20190723-08432990_Motion
+192.168.0.101_20190723-08501190_Motion
+192.168.0.101_20190723-08594738_Motion
+192.168.0.101_20190723-09285531_Motion
+192.168.0.101_20190723-09420829_Motion
+192.168.0.101_20190723-10123474_Motion
+192.168.0.101_20190723-10183522_Motion
+192.168.0.101_20190723-10250419_Motion
+192.168.0.101_20190723-10433566_Motion
+192.168.0.101_20190723-11212308_Motion
+192.168.0.101_20190723-11224407_Motion
+192.168.0.101_20190723-11233257_Motion
+192.168.0.101_20190723-12082848_Motion
+192.168.0.101_20190723-12305993_Motion
+192.168.0.101_20190723-12310692_Motion
+192.168.0.101_20190723-12362990_Motion
+192.168.0.101_20190723-12402990_Motion
+192.168.0.101_20190723-12502827_Motion
+192.168.0.101_20190723-13110701_Motion
+192.168.0.101_20190723-13234549_Motion
+192.168.0.101_20190723-13241598_Motion
+192.168.0.101_20190723-14555527_Motion
+192.168.0.101_20190723-14594823_Motion
+192.168.0.101_20190723-15043174_Motion
+192.168.0.101_20190723-15124570_Motion
+192.168.0.101_20190723-15312018_Motion
+192.168.0.101_20190723-15425663_Motion
+192.168.0.101_20190723-15572809_Motion
+192.168.0.101_20190723-15574909_Motion
+192.168.0.101_20190723-16143255_Motion
+192.168.0.101_20190723-16151257_Motion
+192.168.0.101_20190723-16202053_Motion
+192.168.0.101_20190723-16202802_Motion
+192.168.0.101_20190723-16395644_Motion
+192.168.0.101_20190723-16465694_Motion
+192.168.0.101_20190723-16540292_Motion
+192.168.0.101_20190723-17025889_Motion
+192.168.0.101_20190723-17043538_Motion
+192.168.0.101_20190723-17123835_Motion
+192.168.0.101_20190723-17192435_Motion
+192.168.0.101_20190723-17383430_Motion
+192.168.0.101_20190723-18144521_Motion
+192.168.0.101_20190723-18194467_Motion
+192.168.0.101_20190723-18322865_Motion
+192.168.0.101_20190723-19085206_Motion
+192.168.0.101_20190723-19414050_Motion

+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+11
+12
+13
+14
+15
+16
+17
+18
+19
+20
+21
+22
+23
+24
+25
+26
+27
+28
+29
+30
+31
+32
+33
+34
+35
+36
+37
+38
+39
+40
+41
+42
+43
+44
+45
+46
+47
+48
+49
+50
+51
+52
+53
+54
+55
+56
+57
+58
+59
+60
+61
+62
+63
+64
+65
+66
+67
+68
+69
+70
+71
+72
+73
+74
+75
+76
+77
+78
+79
+80
+81
+82
+83
+84
+85
+86
+87
+88
+89
+90
+91
+92
+93
+94
+95
+96
+97
+98
+99
+100
+101
+102
+103
+104
+105
+106
+107
+108
+109
+110
+111
+112
+113
+114
+115
+116
+117
+118
+119
+120
+121
+122
+123
+124
+125
+126
+127
+128
+129
+130
+131
+132
+133
+134
+135
+136
+137
+138
+139
+140
+141
+142
+143
+144
+145
+146
+147
+148
+149
+150
+151
+152
+153
+154
+155
+156
+157
+158
+159
+160
+161
+162
+163
+164
+165
+166
+167
+168
+169
+170
+171
+172
+173
+174
+175
+176
+177
+178
+179
+180
+181
+182
+183
+184
+185
+186
+187
+188
+189
+190
+191
+192
+193
+194
+195
+196
+197
+198
+199

+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+ba
+beo
+bo
+bu
+da
+deo
+do
+du
+eo
+ga
+geo
+go
+gu
+ha
+heo
+ho
+jeo
+jo
+ju
+ma
+meo
+mo
+mu
+na
+neo
+no
+nu
+o
+ra
+reo
+ro
+ru
+seo
+so
+su
+u
\ No newline at end of file

+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+51
+50
+49
+52
+39
+38
+37
+41
+33
+56
+55
+54
+57
+46
+46
+48
+71
+70
+74
+28
+27
+26
+30
+60
+59
+58
+63
+32
+44
+43
+42
+45
+65
+64
+68
+36
\ No newline at end of file

+import os
+import sys
+import json
+
+# path = sys.argv[1]
+
+# files = os.listdir(path)
+
+# f = open("ids.txt", 'w')
+# for filename in files:
+#     filename = filename.rstrip('.xml')
+#     f.writelines(filename+"\n")
+# f.close()
+
+f = open("ids2.txt", 'w')
+for i in range(200):
+    f.write(str(i)+'\n')
+f.close()
\ No newline at end of file

+import os
+import sys
+import json
+from xml.etree.ElementTree import parse
+
+def get_class(xml_path):
+    tree = parse(xml_path)
+    root = tree.getroot()
+    classes = root.findall("object")
+    names = [x.findtext("name") for x in classes]
+    return names
+    
+    
+
+path = sys.argv[1]
+files = os.listdir(path)
+classlist = []
+
+for file in files:
+    classes = get_class(path+'\\'+file)
+    classlist = list(set(classlist) | set(classes))
+
+classlist.sort()
+f = open("label.txt", 'w')
+for ca in classlist:
+    f.write(ca+'\n')
+f.close()
\ No newline at end of file

+import os
+import argparse
+import json
+import xml.etree.ElementTree as ET
+from typing import Dict, List
+from tqdm import tqdm
+import re
+
+
+def get_label2id(labels_path: str, labeltable_path: str) -> Dict[str, int]:
+    """id is 1 start"""
+    with open(labels_path, 'r') as f:
+        labels_str = f.read().split()
+    with open(labeltable_path, 'r') as f2:
+        table_str = f2.read().split()
+        table_int = list(map(int, table_str))
+    labels_ids = list(range(1, len(labels_str)+1))
+    for i in range(0, len(labels_str)):
+        labels_ids[i] = table_int[i]+1
+    return dict(zip(labels_str, labels_ids))
+
+
+def get_annpaths(ann_dir_path: str = None,
+                 ann_ids_path: str = None,
+                 ext: str = '',
+                 annpaths_list_path: str = None) -> List[str]:
+    # If use annotation paths list
+    if annpaths_list_path is not None:
+        with open(annpaths_list_path, 'r') as f:
+            ann_paths = f.read().split()
+        return ann_paths
+
+    # If use annotaion ids list
+    ext_with_dot = '.' + ext if ext != '' else ''
+    with open(ann_ids_path, 'r') as f:
+        ann_ids = f.read().split()
+    ann_paths = [os.path.join(ann_dir_path, aid+ext_with_dot) for aid in ann_ids]
+    return ann_paths
+
+
+def get_image_info(annotation_root, id, extract_num_from_imgid=True):
+    path = annotation_root.findtext('path')
+    if path is None:
+        filename = annotation_root.findtext('filename')
+    else:
+        filename = os.path.basename(path)
+    img_name = os.path.basename(filename)
+    # img_id = os.path.splitext(img_name)[0]
+    # if extract_num_from_imgid and isinstance(img_id, str):
+    #     img_id = int(re.findall(r'\d+', img_id)[0])
+
+    img_id = id
+    size = annotation_root.find('size')
+    width = int(size.findtext('width'))
+    height = int(size.findtext('height'))
+
+    image_info = {
+        'file_name': filename,
+        'height': height,
+        'width': width,
+        'id': img_id
+    }
+    return image_info
+
+
+def get_coco_annotation_from_obj(obj, label2id):
+    label = obj.findtext('name')
+    assert label in label2id, f"Error: {label} is not in label2id !"
+    category_id = label2id[label]
+    bndbox = obj.find('bndbox')
+    xmin = int(float(bndbox.findtext('xmin'))) - 1
+    ymin = int(float(bndbox.findtext('ymin'))) - 1
+    xmax = int(float(bndbox.findtext('xmax')))
+    ymax = int(float(bndbox.findtext('ymax')))
+    assert xmax > xmin and ymax > ymin, f"Box size error !: (xmin, ymin, xmax, ymax): {xmin, ymin, xmax, ymax}"
+    o_width = xmax - xmin
+    o_height = ymax - ymin
+    ann = {
+        'area': o_width * o_height,
+        'iscrowd': 0,
+        'bbox': [xmin, ymin, o_width, o_height],
+        'category_id': category_id,
+        'ignore': 0,
+        'segmentation': []  # This script is not for segmentation
+    }
+    return ann
+
+
+def convert_xmls_to_cocojson(annotation_paths: List[str],
+                             label2id: Dict[str, int],
+                             output_jsonpath: str,
+                             extract_num_from_imgid: bool = True):
+    output_json_dict = {
+        "images": [],
+        "type": "instances",
+        "annotations": [],
+        "categories": []
+    }
+    bnd_id = 1  # START_BOUNDING_BOX_ID, TODO input as args ?
+    print('Start converting !')
+    i = 1
+    for a_path in tqdm(annotation_paths):
+        # Read annotation xml
+        ann_tree = ET.parse(a_path)
+        ann_root = ann_tree.getroot()
+
+        img_info = get_image_info(annotation_root=ann_root, id=i,
+                                  extract_num_from_imgid=extract_num_from_imgid)
+        img_id = img_info['id']
+        output_json_dict['images'].append(img_info)
+        i+=1
+        for obj in ann_root.findall('object'):
+            ann = get_coco_annotation_from_obj(obj=obj, label2id=label2id)
+            ann.update({'image_id': img_id, 'id': bnd_id})
+            output_json_dict['annotations'].append(ann)
+            bnd_id = bnd_id + 1
+
+    for label, label_id in label2id.items():
+        category_info = {'supercategory': 'none', 'id': label_id, 'name': label}
+        output_json_dict['categories'].append(category_info)
+
+    with open(output_jsonpath, 'w') as f:
+        output_json = json.dumps(output_json_dict)
+        f.write(output_json)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='This script support converting voc format xmls to coco format json')
+    parser.add_argument('--ann_dir', type=str, default=None,
+                        help='path to annotation files directory. It is not need when use --ann_paths_list')
+    parser.add_argument('--ann_ids', type=str, default=None,
+                        help='path to annotation files ids list. It is not need when use --ann_paths_list')
+    parser.add_argument('--ann_paths_list', type=str, default=None,
+                        help='path of annotation paths list. It is not need when use --ann_dir and --ann_ids')
+    parser.add_argument('--labels', type=str, default=None,
+                        help='path to label list.')
+    parser.add_argument('--labelt', type=str, default=None,
+                        help='path to label table.')
+    parser.add_argument('--output', type=str, default='output.json', help='path to output json file')
+    parser.add_argument('--ext', type=str, default='', help='additional extension of annotation file')
+    parser.add_argument('--extract_num_from_imgid', action="store_true",
+                        help='Extract image number from the image filename')
+    args = parser.parse_args()
+    label2id = get_label2id(labels_path=args.labels, labeltable_path=args.labelt)
+    ann_paths = get_annpaths(
+        ann_dir_path=args.ann_dir,
+        ann_ids_path=args.ann_ids,
+        ext=args.ext,
+        annpaths_list_path=args.ann_paths_list
+    )
+    convert_xmls_to_cocojson(
+        annotation_paths=ann_paths,
+        label2id=label2id,
+        output_jsonpath=args.output,
+        extract_num_from_imgid=args.extract_num_from_imgid
+    )
+
+
+if __name__ == '__main__':
+    main()