2021-1-capstone-design1 / HCS_Project2

Go to a project
Toggle navigation Toggle navigation pinning
정우진
d1114a3d 1 parent 04050124

nvidia retinanet container

Inline Side-by-side
Showing 65 changed files with 7799 additions and 0 deletions
# 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&deg;, 90&deg;, 180&deg; or 270&deg;.
* `--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
File mode changed
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()