Skip to content

Commit

Permalink
Add support for yolov4-p5 and yolov4-p6
Browse files Browse the repository at this point in the history
  • Loading branch information
GeekAlexis committed Jul 5, 2021
1 parent 9840fbc commit 32c217a
Showing 7 changed files with 1,014 additions and 24 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -101,7 +101,7 @@ Only required for SSD (not supported on Ubuntu 20.04)
```
- Image sequence:
```bash
python3 app.py --input_uri img_%06d.jpg --mot
python3 app.py --input_uri %06d.jpg --mot
```
- Video file:
```bash
3 changes: 2 additions & 1 deletion fastmot/models/label.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""
90-class COCO labels
`unlabled` (id = 0) is replaced with `head` to work with CrowdHuman
`unlabled` (id = 0) is replaced with `head` for CrowdHuman
These are different from the default 80-class COCO labels used by YOLO
"""

LABEL_MAP = (
117 changes: 111 additions & 6 deletions fastmot/models/yolo.py
Original file line number Diff line number Diff line change
@@ -22,9 +22,6 @@ class YOLO:

@classmethod
def add_plugin(cls, network):
"""
Adapted from https://github.com/jkjung-avt/tensorrt_demos
"""
def get_plugin_creator(plugin_name):
plugin_creators = trt.get_plugin_registry().plugin_creator_list
for plugin_creator in plugin_creators:
@@ -107,6 +104,114 @@ class YOLOv4(YOLO):
INPUT_SHAPE = (3, 512, 512)
LAYER_FACTORS = [8, 16, 32]
SCALES = [1.2, 1.1, 1.05]
ANCHORS = [[11, 22, 24, 60, 37, 116],
[54, 186, 69, 268, 89, 369],
[126, 491, 194, 314, 278, 520]]
ANCHORS = [[11,22, 24,60, 37,116],
[54,186, 69,268, 89,369],
[126,491, 194,314, 278,520]]


"""
The following models are supported but not provided.
Modify paths, # classes, input shape, and anchors according to your Darknet cfg for custom model.
"""

class YOLOv4CSP(YOLO):
ENGINE_PATH = Path(__file__).parent / 'yolov4-csp.trt'
MODEL_PATH = Path(__file__).parent / 'yolov4-csp.onnx'
NUM_CLASSES = 1
LETTERBOX = True
NEW_COORDS = True
INPUT_SHAPE = (3, 512, 512)
LAYER_FACTORS = [8, 16, 32]
SCALES = [2.0, 2.0, 2.0]
ANCHORS = [[12,16, 19,36, 40,28],
[36,75, 76,55, 72,146],
[142,110, 192,243, 459,401]]


class YOLOv4xMish(YOLO):
ENGINE_PATH = Path(__file__).parent / 'yolov4x-mish.trt'
MODEL_PATH = Path(__file__).parent / 'yolov4x-mish.onnx'
NUM_CLASSES = 1
LETTERBOX = True
NEW_COORDS = True
INPUT_SHAPE = (3, 640, 640)
LAYER_FACTORS = [8, 16, 32]
SCALES = [2.0, 2.0, 2.0]
ANCHORS = [[12,16, 19,36, 40,28],
[36,75, 76,55, 72,146],
[142,110, 192,243, 459,401]]


class YOLOv4P5(YOLO):
ENGINE_PATH = Path(__file__).parent / 'yolov4-p5.trt'
MODEL_PATH = Path(__file__).parent / 'yolov4-p5.onnx'
NUM_CLASSES = 1
LETTERBOX = True
NEW_COORDS = True
INPUT_SHAPE = (3, 896, 896)
LAYER_FACTORS = [8, 16, 32]
SCALES = [2.0, 2.0, 2.0]
ANCHORS = [[13,17, 31,25, 24,51, 61,45],
[48,102, 119,96, 97,189, 217,184],
[171,384, 324,451, 616,618, 800,800]]


class YOLOv4P6(YOLO):
ENGINE_PATH = Path(__file__).parent / 'yolov4-p6.trt'
MODEL_PATH = Path(__file__).parent / 'yolov4-p6.onnx'
NUM_CLASSES = 1
LETTERBOX = True
NEW_COORDS = True
INPUT_SHAPE = (3, 1280, 1280)
LAYER_FACTORS = [8, 16, 32, 64]
SCALES = [2.0, 2.0, 2.0, 2.0]
ANCHORS = [[13,17, 31,25, 24,51, 61,45],
[61,45, 48,102, 119,96, 97,189],
[97,189, 217,184, 171,384, 324,451],
[324,451, 545,357, 616,618, 1024,1024]]


class YOLOv4Tiny(YOLO):
ENGINE_PATH = Path(__file__).parent / 'yolov4-tiny.trt'
MODEL_PATH = Path(__file__).parent / 'yolov4-tiny.onnx'
NUM_CLASSES = 1
INPUT_SHAPE = (3, 416, 416)
LAYER_FACTORS = [32, 16]
SCALES = [1.05, 1.05]
ANCHORS = [[81,82, 135,169, 344,319],
[23,27, 37,58, 81,82]]


class YOLOv3(YOLO):
ENGINE_PATH = Path(__file__).parent / 'yolov3.trt'
MODEL_PATH = Path(__file__).parent / 'yolov3.onnx'
NUM_CLASSES = 1
INPUT_SHAPE = (3, 416, 416)
LAYER_FACTORS = [32, 16, 8]
SCALES = [1., 1.]
ANCHORS = [[116,90, 156,198, 373,326],
[30,61, 62,45, 59,119],
[10,13, 16,30, 33,23]]


class YOLOv3SPP(YOLO):
ENGINE_PATH = Path(__file__).parent / 'yolov3-spp.trt'
MODEL_PATH = Path(__file__).parent / 'yolov3-spp.onnx'
NUM_CLASSES = 1
INPUT_SHAPE = (3, 608, 608)
LAYER_FACTORS = [32, 16, 8]
SCALES = [1., 1.]
ANCHORS = [[116,90, 156,198, 373,326],
[30,61, 62,45, 59,119],
[10,13, 16,30, 33,23]]


class YOLOv3Tiny(YOLO):
ENGINE_PATH = Path(__file__).parent / 'yolov3-tiny.trt'
MODEL_PATH = Path(__file__).parent / 'yolov3-tiny.onnx'
NUM_CLASSES = 1
INPUT_SHAPE = (3, 416, 416)
LAYER_FACTORS = [32, 16]
SCALES = [1., 1.]
ANCHORS = [[81,82, 135,169, 344,319],
[10,14, 23,27, 37,58]]
2 changes: 1 addition & 1 deletion fastmot/plugins/Makefile
Original file line number Diff line number Diff line change
@@ -21,7 +21,7 @@ TENSORRT_LIBS=-L"/usr/lib/x86_64-linux-gnu"

# INCS and LIBS
INCS=-I"/usr/local/cuda/include" $(TENSORRT_INCS) -I"/usr/local/include" -I"plugin"
LIBS=-L"/usr/local/cuda/lib64" $(TENSORRT_LIBS) -L"/usr/local/lib" -Wl,--start-group -lnvinfer -lnvparsers -lnvinfer_plugin -lcudnn -lcublas -lcudart_static -lnvToolsExt -lcudart -lrt -ldl -lpthread -Wl,--end-group
LIBS=-L"/usr/local/cuda/lib64" $(TENSORRT_LIBS) -L"/usr/local/lib" -Wl,--start-group -lnvinfer -lnvparsers -lnvinfer_plugin -lcudnn -lcublas -lnvToolsExt -lcudart -lrt -ldl -lpthread -Wl,--end-group

.PHONY: all clean

2 changes: 1 addition & 1 deletion fastmot/plugins/README.md
Original file line number Diff line number Diff line change
@@ -1 +1 @@
"yolo_layer.h" and "yolo_layer.cu" are taken from [tensorrt_demos](https://github.com/jkjung-avt/tensorrt_demos). The original code is under [MIT License](https://github.com/jkjung-avt/tensorrt_demos/blob/master/LICENSE).
YOLO layer plugins are adapted from [tensorrt_demos](https://github.com/jkjung-avt/tensorrt_demos). The original code is under [MIT License](https://github.com/jkjung-avt/tensorrt_demos/blob/master/LICENSE).
15 changes: 1 addition & 14 deletions fastmot/plugins/yolo_layer.cu
Original file line number Diff line number Diff line change
@@ -1,16 +1,3 @@
/*
* yolo_layer.cu
*
* This code was originally written by wang-xinyu under MIT license.
* I took it from:
*
* https://github.com/wang-xinyu/tensorrtx/tree/master/yolov4
*
* and made necessary modifications.
*
* - JK Jung
*/

#include "yolo_layer.h"

using namespace Yolo;
@@ -410,7 +397,7 @@ namespace nvinfer1
assert(yolo_width > 0 && yolo_height > 0);
assert(anchors[0] > 0.0f && anchors[1] > 0.0f);
assert(num_classes > 0);
assert(input_multiplier == 8 || input_multiplier == 16 || input_multiplier == 32);
assert(input_multiplier == 8 || input_multiplier == 16 || input_multiplier == 32 || input_multiplier == 64 || input_multiplier == 128);
assert(scale_x_y >= 1.0);

YoloLayerPlugin* obj = new YoloLayerPlugin(yolo_width, yolo_height, num_anchors, anchors, num_classes, yolo_width * input_multiplier, yolo_height * input_multiplier, scale_x_y, new_coords);
Loading

0 comments on commit 32c217a

Please sign in to comment.