Skip to content

Commit 32c217a

Browse files
GeekAlexisGeekAlexis
committed
Add support for yolov4-p5 and yolov4-p6
1 parent 9840fbc commit 32c217a

File tree

7 files changed

+1014
-24
lines changed

7 files changed

+1014
-24
lines changed

‎README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ Only required for SSD (not supported on Ubuntu 20.04)
101101
```
102102
- Image sequence:
103103
```bash
104-
python3 app.py --input_uri img_%06d.jpg --mot
104+
python3 app.py --input_uri %06d.jpg --mot
105105
```
106106
- Video file:
107107
```bash

‎fastmot/models/label.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""
22
90-class COCO labels
3-
`unlabled` (id = 0) is replaced with `head` to work with CrowdHuman
3+
`unlabled` (id = 0) is replaced with `head` for CrowdHuman
4+
These are different from the default 80-class COCO labels used by YOLO
45
"""
56

67
LABEL_MAP = (

‎fastmot/models/yolo.py

+111-6
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,6 @@ class YOLO:
2222

2323
@classmethod
2424
def add_plugin(cls, network):
25-
"""
26-
Adapted from https://github.com/jkjung-avt/tensorrt_demos
27-
"""
2825
def get_plugin_creator(plugin_name):
2926
plugin_creators = trt.get_plugin_registry().plugin_creator_list
3027
for plugin_creator in plugin_creators:
@@ -107,6 +104,114 @@ class YOLOv4(YOLO):
107104
INPUT_SHAPE = (3, 512, 512)
108105
LAYER_FACTORS = [8, 16, 32]
109106
SCALES = [1.2, 1.1, 1.05]
110-
ANCHORS = [[11, 22, 24, 60, 37, 116],
111-
[54, 186, 69, 268, 89, 369],
112-
[126, 491, 194, 314, 278, 520]]
107+
ANCHORS = [[11,22, 24,60, 37,116],
108+
[54,186, 69,268, 89,369],
109+
[126,491, 194,314, 278,520]]
110+
111+
112+
"""
113+
The following models are supported but not provided.
114+
Modify paths, # classes, input shape, and anchors according to your Darknet cfg for custom model.
115+
"""
116+
117+
class YOLOv4CSP(YOLO):
118+
ENGINE_PATH = Path(__file__).parent / 'yolov4-csp.trt'
119+
MODEL_PATH = Path(__file__).parent / 'yolov4-csp.onnx'
120+
NUM_CLASSES = 1
121+
LETTERBOX = True
122+
NEW_COORDS = True
123+
INPUT_SHAPE = (3, 512, 512)
124+
LAYER_FACTORS = [8, 16, 32]
125+
SCALES = [2.0, 2.0, 2.0]
126+
ANCHORS = [[12,16, 19,36, 40,28],
127+
[36,75, 76,55, 72,146],
128+
[142,110, 192,243, 459,401]]
129+
130+
131+
class YOLOv4xMish(YOLO):
132+
ENGINE_PATH = Path(__file__).parent / 'yolov4x-mish.trt'
133+
MODEL_PATH = Path(__file__).parent / 'yolov4x-mish.onnx'
134+
NUM_CLASSES = 1
135+
LETTERBOX = True
136+
NEW_COORDS = True
137+
INPUT_SHAPE = (3, 640, 640)
138+
LAYER_FACTORS = [8, 16, 32]
139+
SCALES = [2.0, 2.0, 2.0]
140+
ANCHORS = [[12,16, 19,36, 40,28],
141+
[36,75, 76,55, 72,146],
142+
[142,110, 192,243, 459,401]]
143+
144+
145+
class YOLOv4P5(YOLO):
146+
ENGINE_PATH = Path(__file__).parent / 'yolov4-p5.trt'
147+
MODEL_PATH = Path(__file__).parent / 'yolov4-p5.onnx'
148+
NUM_CLASSES = 1
149+
LETTERBOX = True
150+
NEW_COORDS = True
151+
INPUT_SHAPE = (3, 896, 896)
152+
LAYER_FACTORS = [8, 16, 32]
153+
SCALES = [2.0, 2.0, 2.0]
154+
ANCHORS = [[13,17, 31,25, 24,51, 61,45],
155+
[48,102, 119,96, 97,189, 217,184],
156+
[171,384, 324,451, 616,618, 800,800]]
157+
158+
159+
class YOLOv4P6(YOLO):
160+
ENGINE_PATH = Path(__file__).parent / 'yolov4-p6.trt'
161+
MODEL_PATH = Path(__file__).parent / 'yolov4-p6.onnx'
162+
NUM_CLASSES = 1
163+
LETTERBOX = True
164+
NEW_COORDS = True
165+
INPUT_SHAPE = (3, 1280, 1280)
166+
LAYER_FACTORS = [8, 16, 32, 64]
167+
SCALES = [2.0, 2.0, 2.0, 2.0]
168+
ANCHORS = [[13,17, 31,25, 24,51, 61,45],
169+
[61,45, 48,102, 119,96, 97,189],
170+
[97,189, 217,184, 171,384, 324,451],
171+
[324,451, 545,357, 616,618, 1024,1024]]
172+
173+
174+
class YOLOv4Tiny(YOLO):
175+
ENGINE_PATH = Path(__file__).parent / 'yolov4-tiny.trt'
176+
MODEL_PATH = Path(__file__).parent / 'yolov4-tiny.onnx'
177+
NUM_CLASSES = 1
178+
INPUT_SHAPE = (3, 416, 416)
179+
LAYER_FACTORS = [32, 16]
180+
SCALES = [1.05, 1.05]
181+
ANCHORS = [[81,82, 135,169, 344,319],
182+
[23,27, 37,58, 81,82]]
183+
184+
185+
class YOLOv3(YOLO):
186+
ENGINE_PATH = Path(__file__).parent / 'yolov3.trt'
187+
MODEL_PATH = Path(__file__).parent / 'yolov3.onnx'
188+
NUM_CLASSES = 1
189+
INPUT_SHAPE = (3, 416, 416)
190+
LAYER_FACTORS = [32, 16, 8]
191+
SCALES = [1., 1.]
192+
ANCHORS = [[116,90, 156,198, 373,326],
193+
[30,61, 62,45, 59,119],
194+
[10,13, 16,30, 33,23]]
195+
196+
197+
class YOLOv3SPP(YOLO):
198+
ENGINE_PATH = Path(__file__).parent / 'yolov3-spp.trt'
199+
MODEL_PATH = Path(__file__).parent / 'yolov3-spp.onnx'
200+
NUM_CLASSES = 1
201+
INPUT_SHAPE = (3, 608, 608)
202+
LAYER_FACTORS = [32, 16, 8]
203+
SCALES = [1., 1.]
204+
ANCHORS = [[116,90, 156,198, 373,326],
205+
[30,61, 62,45, 59,119],
206+
[10,13, 16,30, 33,23]]
207+
208+
209+
class YOLOv3Tiny(YOLO):
210+
ENGINE_PATH = Path(__file__).parent / 'yolov3-tiny.trt'
211+
MODEL_PATH = Path(__file__).parent / 'yolov3-tiny.onnx'
212+
NUM_CLASSES = 1
213+
INPUT_SHAPE = (3, 416, 416)
214+
LAYER_FACTORS = [32, 16]
215+
SCALES = [1., 1.]
216+
ANCHORS = [[81,82, 135,169, 344,319],
217+
[10,14, 23,27, 37,58]]

‎fastmot/plugins/Makefile

+1-1
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ TENSORRT_LIBS=-L"/usr/lib/x86_64-linux-gnu"
2121

2222
# INCS and LIBS
2323
INCS=-I"/usr/local/cuda/include" $(TENSORRT_INCS) -I"/usr/local/include" -I"plugin"
24-
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
24+
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
2525

2626
.PHONY: all clean
2727

‎fastmot/plugins/README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -1 +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).
1+
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).

‎fastmot/plugins/yolo_layer.cu

+1-14
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,3 @@
1-
/*
2-
* yolo_layer.cu
3-
*
4-
* This code was originally written by wang-xinyu under MIT license.
5-
* I took it from:
6-
*
7-
* https://github.com/wang-xinyu/tensorrtx/tree/master/yolov4
8-
*
9-
* and made necessary modifications.
10-
*
11-
* - JK Jung
12-
*/
13-
141
#include "yolo_layer.h"
152

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

416403
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);

0 commit comments

Comments
 (0)