Skip to content

Commit 2f5de0c

Browse files
Li-HongdaLi-Hongda
committed
update vis mask
1 parent 9f0c165 commit 2f5de0c

16 files changed

Lines changed: 194 additions & 222 deletions

‎configs/rtdetr.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
rtdetr:
22
onnx_file: "../weights/rtdetr/rtdetr_hgnetv2_l_6x_coco.onnx"
33
engine_file: "../weights/rtdetr/rtdetr_hgnetl.trt"
4-
type: "coco80"
4+
type: "coco"
55
mode: "fp32"
66
batchSize: 1
77
imageWidth: 640

‎configs/yolov5-seg.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
yolov5-seg:
22
onnx_file: "../weights/yolov5/yolov5s-seg.onnx"
33
engine_file: "../weights/yolov5/yolov5s-seg.trt"
4-
type: "coco80"
4+
type: "coco"
55
mode: "fp32"
66
batchSize: 1
77
imageWidth: 640

‎configs/yolov5.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
yolov5:
22
onnx_file: "../weights/yolov5/yolov5s.onnx"
33
engine_file: "../weights/yolov5/yolov5s.trt"
4-
type: "coco80"
4+
type: "coco"
55
mode: "fp32"
66
batchSize: 1
77
imageWidth: 640

‎configs/yolov7-p6.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
yolov7:
22
onnx_file: "../weights/yolov7/yolov7-w6.onnx"
33
engine_file: "../weights/yolov7/yolov7-w6.trt"
4-
type: "coco80"
4+
type: "coco"
55
mode: "fp32"
66
batchSize: 1
77
imageWidth: 1280

‎configs/yolov7.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
yolov7:
22
onnx_file: "../weights/yolov7/yolov7.onnx"
33
engine_file: "../weights/yolov7/yolov7.trt"
4-
type: "coco80"
4+
type: "coco"
55
mode: "fp32"
66
batchSize: 1
77
imageWidth: 640

‎configs/yolov8-seg.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
yolov8-seg:
22
onnx_file: "../weights/yolov8/yolov8s-seg.onnx"
33
engine_file: "../weights/yolov8/yolov8s-seg.trt"
4-
type: "coco80"
4+
type: "coco"
55
mode: "fp32"
66
batchSize: 1
77
imageWidth: 640

‎configs/yolov8.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
yolov8:
22
onnx_file: "../weights/yolov8/yolov8s.onnx"
33
engine_file: "../weights/yolov8/yolov8s.trt"
4-
type: "coco80"
4+
type: "coco"
55
mode: "fp32"
66
batchSize: 1
77
imageWidth: 640

‎include/basemodel.h‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,7 @@
44
#include "common.h"
55
#include "cuda_function.h"
66

7-
class Model
8-
{
7+
class Model {
98
public:
109
explicit Model(const YAML::Node &config);
1110
~Model();
@@ -27,7 +26,6 @@ class Model
2726
std::string names[10];
2827
float** cpu_buffers = new float* [10];
2928
float* gpu_buffers[10]{};
30-
// float* cpu_mask_buffer = nullptr;
3129
std::vector<int64_t> bufferSize;
3230
std::shared_ptr<nvinfer1::ICudaEngine> engine;
3331
std::unique_ptr<nvinfer1::IExecutionContext> context;

‎include/cuda_function.h‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,25 +15,26 @@ void cuda_preprocess_destroy();
1515
void cuda_postprocess_init(int num_out, int width, int height);
1616
void cuda_postprocess_destroy();
1717

18+
void process_mask_init(int num_out, int width, int height);
19+
1820
void preprocess(uint8_t* src, AffineMatrix d2s, int src_width, int src_height,
1921
float* dst, int dst_width, int dst_height,
2022
cudaStream_t stream);
2123

22-
2324
void postprocess_box(float* predict, int num_bboxes, int num_classes, int num_out, float conf_thr,
24-
float nms_thr, cudaStream_t stream, float* dst);
25+
float nms_thr, AffineMatrix mat, cudaStream_t stream, float* dst);
2526

2627
void yolov8_postprocess_box(float* predict, int num_bboxes, int num_classes, int num_out, float conf_thr,
27-
float nms_thr, cudaStream_t stream, float* dst);
28+
float nms_thr, AffineMatrix mat, cudaStream_t stream, float* dst);
2829

2930
void rtdetr_postprocess_box(float* predict_box, float* predict_cls, int num_bboxes, int num_classes, int num_out,
30-
float conf_thr, int imageWidth, int imageHeight, cudaStream_t stream, float* dst);
31+
float conf_thr, int imageWidth, int imageHeight, AffineMatrix mat, cudaStream_t stream, float* dst);
3132

3233
void postprocess_box_mask(float* predict, int num_bboxes, int num_classes, int num_out,
33-
float conf_thr, float nms_thr, cudaStream_t stream, float* dst);
34+
float conf_thr, float nms_thr, AffineMatrix mat, cudaStream_t stream, float* dst);
3435

3536
void yolov8_postprocess_box_mask(float* predict, int num_bboxes, int num_classes, int num_out,
36-
float conf_thr, float nms_thr, cudaStream_t stream, float* dst);
37+
float conf_thr, float nms_thr, AffineMatrix mat, cudaStream_t stream, float* dst);
3738

3839
void process_mask(float* out, float* proto, uint8_t* dst , int num_out, int dst_width,
3940
int dst_height, int out_w, int proto_size,

‎include/detection.h‎

Lines changed: 14 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33

44
#include "basemodel.h"
55

6-
namespace Category{
7-
const std::vector<std::string> coco80 = {
6+
namespace Category {
7+
const std::vector<std::string> coco = {
88
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
99
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
1010
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
@@ -15,19 +15,14 @@ namespace Category{
1515
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
1616
"hair drier", "toothbrush"
1717
};
18-
const std::vector<std::string> coco91 = {
19-
"person", "bicycle","car","motorcycle","airplane","bus","train","truck","boat","traffic light",
20-
"fire hydrant","street sign","stop sign","parking meter","bench","bird","cat","dog","horse","sheep","cow","elephant","bear","zebra","giraffe",
21-
"hat","backpack","umbrella","shoe","eye glasses","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite","baseball bat",
22-
"baseball glove","skateboard","surfboard","tennis racket","bottle","plate","wine glass","cup","fork","knife","spoon","bowl","banana","apple",
23-
"sandwich","orange","broccoli","carrot","hot dog","pizza","donut","cake","chair","couch","potted plant","bed","mirror","dining table","window",
24-
"desk","toilet","door","tv","laptop","mouse","remote","keyboard","cell phone","microwave","oven","toaster","sink","refrigerator","blender",
25-
"book","clock","vase","scissors","teddy bear","hair drier","toothbrush","hair brush"
26-
};
18+
const std::vector<std::string> voc = {
19+
"aeroplane","bicycle","bird","boat","bottle","bus","car","cat","chair","cow","diningtable",
20+
"dog","horse","motorbike","person","pottedplant","sheep","sofa","train","tvmonitor"
21+
};
2722
}
2823

29-
namespace Color{
30-
const std::vector<cv::Scalar> coco80{
24+
namespace Color {
25+
const std::vector<cv::Scalar> coco {
3126
cv::Scalar(128, 77, 207),cv::Scalar(65, 32, 208),cv::Scalar(0, 224, 45),cv::Scalar(3, 141, 219),cv::Scalar(80, 239, 253),cv::Scalar(239, 184, 12),
3227
cv::Scalar(7, 144, 145),cv::Scalar(161, 88, 57),cv::Scalar(0, 166, 46),cv::Scalar(218, 113, 53),cv::Scalar(193, 33, 128),cv::Scalar(190, 94, 113),
3328
cv::Scalar(113, 123, 232),cv::Scalar(69, 205, 80),cv::Scalar(18, 170, 49),cv::Scalar(89, 51, 241),cv::Scalar(153, 191, 154),cv::Scalar(27, 26, 69),
@@ -43,24 +38,12 @@ namespace Color{
4338
cv::Scalar(102, 195, 55),cv::Scalar(160, 26, 91),cv::Scalar(60, 94, 66),cv::Scalar(204, 169, 193),cv::Scalar(126, 4, 181),cv::Scalar(229, 209, 196),
4439
cv::Scalar(195, 170, 186),cv::Scalar(155, 207, 148)
4540
};
46-
const std::vector<cv::Scalar> coco91{
47-
cv::Scalar(148, 99, 164),cv::Scalar(65, 172, 90),cv::Scalar(18, 117, 190),cv::Scalar(173, 208, 229),cv::Scalar(37, 162, 147),cv::Scalar(121, 99, 42),
48-
cv::Scalar(218, 173, 104),cv::Scalar(193, 213, 138),cv::Scalar(142, 168, 45),cv::Scalar(107, 143, 94),cv::Scalar(242, 89, 7),cv::Scalar(87, 218, 248),
49-
cv::Scalar(126, 168, 9),cv::Scalar(86, 152, 105),cv::Scalar(155, 135, 251),cv::Scalar(73, 234, 44),cv::Scalar(177, 37, 42),cv::Scalar(219, 215, 54),
50-
cv::Scalar(124, 207, 143),cv::Scalar(7, 81, 209),cv::Scalar(254, 18, 130),cv::Scalar(71, 54, 73),cv::Scalar(172, 198, 63),cv::Scalar(64, 217, 224),
51-
cv::Scalar(105, 224, 25),cv::Scalar(41, 52, 130),cv::Scalar(220, 27, 193),cv::Scalar(65, 222, 86),cv::Scalar(250, 150, 201),cv::Scalar(201, 150, 105),
52-
cv::Scalar(104, 96, 142),cv::Scalar(111, 230, 54),cv::Scalar(105, 24, 22),cv::Scalar(42, 226, 101),cv::Scalar(67, 26, 144),cv::Scalar(155, 113, 106),
53-
cv::Scalar(152, 196, 216),cv::Scalar(58, 68, 152),cv::Scalar(68, 230, 213),cv::Scalar(169, 143, 129),cv::Scalar(191, 102, 41),cv::Scalar(5, 73, 170),
54-
cv::Scalar(15, 73, 233),cv::Scalar(95, 13, 71),cv::Scalar(25, 92, 218),cv::Scalar(85, 173, 16),cv::Scalar(247, 158, 17),cv::Scalar(36, 28, 8),
55-
cv::Scalar(31, 100, 134),cv::Scalar(131, 71, 45),cv::Scalar(158, 190, 91),cv::Scalar(90, 207, 220),cv::Scalar(125, 77, 228),cv::Scalar(40, 156, 67),
56-
cv::Scalar(35, 250, 69),cv::Scalar(229, 61, 245),cv::Scalar(210, 201, 106),cv::Scalar(184, 35, 131),cv::Scalar(47, 124, 120),cv::Scalar(1, 114, 23),
57-
cv::Scalar(99, 181, 17),cv::Scalar(77, 141, 151),cv::Scalar(79, 33, 95),cv::Scalar(194, 111, 146),cv::Scalar(187, 199, 138),cv::Scalar(129, 215, 40),
58-
cv::Scalar(160, 209, 144),cv::Scalar(139, 121, 58),cv::Scalar(97, 208, 197),cv::Scalar(185, 105, 171),cv::Scalar(160, 96, 136),cv::Scalar(232, 26, 26),
59-
cv::Scalar(34, 165, 109),cv::Scalar(19, 86, 215),cv::Scalar(205, 209, 199),cv::Scalar(131, 91, 25),cv::Scalar(51, 201, 16),cv::Scalar(64, 35, 128),
60-
cv::Scalar(120, 161, 247),cv::Scalar(123, 164, 190),cv::Scalar(15, 191, 40),cv::Scalar(11, 44, 117),cv::Scalar(198, 136, 70),cv::Scalar(14, 224, 240),
61-
cv::Scalar(60, 186, 193),cv::Scalar(253, 190, 129),cv::Scalar(134, 228, 173),cv::Scalar(219, 156, 214),cv::Scalar(137, 67, 254),cv::Scalar(178, 223, 250),
62-
cv::Scalar(219, 199, 139)
63-
};
41+
const std::vector<cv::Scalar> voc {
42+
cv::Scalar(128, 77, 207),cv::Scalar(65, 32, 208),cv::Scalar(0, 224, 45),cv::Scalar(3, 141, 219),cv::Scalar(80, 239, 253),cv::Scalar(239, 184, 12),
43+
cv::Scalar(7, 144, 145),cv::Scalar(161, 88, 57),cv::Scalar(0, 166, 46),cv::Scalar(218, 113, 53),cv::Scalar(193, 33, 128),cv::Scalar(190, 94, 113),
44+
cv::Scalar(113, 123, 232),cv::Scalar(69, 205, 80),cv::Scalar(18, 170, 49),cv::Scalar(89, 51, 241),cv::Scalar(153, 191, 154),cv::Scalar(27, 26, 69),
45+
cv::Scalar(20, 186, 194),cv::Scalar(210, 202, 167),cv::Scalar(196, 113, 204),cv::Scalar(9, 81, 88),cv::Scalar(191, 162, 67),cv::Scalar(227, 73, 120)
46+
};
6447
};
6548

6649
struct Box {

0 commit comments

Comments
 (0)