@ -1,9 +1,10 @@
//
// Created by ubuntu on 1/8 /23.
// Created by ubuntu on 1/24 /23.
//
# include "config.h"
# include "utils.h"
# ifndef SEGMENT_YOLOV8_SEG_HPP
# define SEGMENT_YOLOV8_SEG_HPP
# include <fstream>
# include "common.hpp"
# include "NvInferPlugin.h"
using namespace seg ;
@ -16,23 +17,44 @@ public:
void make_pipe ( bool warmup = true ) ;
void copy_from_Mat ( const cv : : Mat & image ) ;
void copy_from_Mat ( const cv : : Mat & image , cv : : Size & size ) ;
void letterbox (
const cv : : Mat & image ,
cv : : Mat & out ,
cv : : Size & size
) ;
void infer ( ) ;
void postprocess ( std : : vector < Object > & objs ) ;
size_t in_size = 1 * 3 * INPUT_W * INPUT_H ;
float w = INPUT_W ;
float h = INPUT_H ;
float ratio = 1.0f ;
float dw = 0.f ;
float dh = 0.f ;
std : : array < std : : pair < int , int > , NUM_OUTPUT > out_sizes { } ;
std : : array < void * , NUM_OUTPUT > outputs { } ;
void postprocess (
std : : vector < Object > & objs ,
float score_thres = 0.25f ,
float iou_thres = 0.65f ,
int topk = 100 ,
int seg_channels = 32 ,
int seg_h = 160 ,
int seg_w = 160
) ;
static void draw_objects (
const cv : : Mat & image ,
cv : : Mat & res ,
const std : : vector < Object > & objs ,
const std : : vector < std : : string > & CLASS_NAMES ,
const std : : vector < std : : vector < unsigned int > > & COLORS ,
const std : : vector < std : : vector < unsigned int > > & MASK_COLORS
) ;
int num_bindings ;
int num_inputs = 0 ;
int num_outputs = 0 ;
std : : vector < Binding > input_bindings ;
std : : vector < Binding > output_bindings ;
std : : vector < void * > host_ptrs ;
std : : vector < void * > device_ptrs ;
PreParam pparam ;
private :
nvinfer1 : : ICudaEngine * engine = nullptr ;
nvinfer1 : : IRuntime * runtime = nullptr ;
nvinfer1 : : IExecutionContext * context = nullptr ;
cudaStream_t stream = nullptr ;
std : : array < void * , NUM_BINDINGS > buffs { } ;
Logger gLogger { nvinfer1 : : ILogger : : Severity : : kERROR } ;
} ;
@ -43,8 +65,6 @@ YOLOv8_seg::YOLOv8_seg(const std::string& engine_file_path)
assert ( file . good ( ) ) ;
file . seekg ( 0 , std : : ios : : end ) ;
auto size = file . tellg ( ) ;
std : : ostringstream fmt ;
file . seekg ( 0 , std : : ios : : beg ) ;
char * trtModelStream = new char [ size ] ;
assert ( trtModelStream ) ;
@ -61,6 +81,41 @@ YOLOv8_seg::YOLOv8_seg(const std::string& engine_file_path)
assert ( this - > context ! = nullptr ) ;
cudaStreamCreate ( & this - > stream ) ;
this - > num_bindings = this - > engine - > getNbBindings ( ) ;
for ( int i = 0 ; i < this - > num_bindings ; + + i )
{
Binding binding ;
nvinfer1 : : Dims dims ;
nvinfer1 : : DataType dtype = this - > engine - > getBindingDataType ( i ) ;
std : : string name = this - > engine - > getBindingName ( i ) ;
binding . name = name ;
binding . dsize = type_to_size ( dtype ) ;
bool IsInput = engine - > bindingIsInput ( i ) ;
if ( IsInput )
{
this - > num_inputs + = 1 ;
dims = this - > engine - > getProfileDimensions (
i ,
0 ,
nvinfer1 : : OptProfileSelector : : kMAX ) ;
binding . size = get_size_by_dims ( dims ) ;
binding . dims = dims ;
this - > input_bindings . push_back ( binding ) ;
// set max opt shape
this - > context - > setBindingDimensions ( i , dims ) ;
}
else
{
dims = this - > context - > getBindingDimensions ( i ) ;
binding . size = get_size_by_dims ( dims ) ;
binding . dims = dims ;
this - > output_bindings . push_back ( binding ) ;
this - > num_outputs + = 1 ;
}
}
}
@ -70,58 +125,67 @@ YOLOv8_seg::~YOLOv8_seg()
this - > engine - > destroy ( ) ;
this - > runtime - > destroy ( ) ;
cudaStreamDestroy ( this - > stream ) ;
for ( auto & ptr : this - > buff s)
for ( auto & ptr : this - > device_ptr s)
{
CHECK ( cudaFree ( ptr ) ) ;
}
for ( auto & ptr : this - > output s)
for ( auto & ptr : this - > host_ptr s)
{
CHECK ( cudaFreeHost ( ptr ) ) ;
}
}
void YOLOv8_seg : : make_pipe ( bool warmup )
{
const nvinfer1 : : Dims input_dims = this - > engine - > getBindingDimensions (
this - > engine - > getBindingIndex ( INPUT )
) ;
this - > in_size = get_size_by_dims ( input_dims ) ;
CHECK ( cudaMalloc ( & this - > buffs [ 0 ] , this - > in_size * sizeof ( float ) ) ) ;
this - > context - > setBindingDimensions ( 0 , input_dims ) ;
const int32_t output_idx = this - > engine - > getBindingIndex ( OUTPUT ) ;
const nvinfer1 : : Dims output_dims = this - > context - > getBindingDimensions ( output_idx ) ;
this - > out_sizes [ output_idx - NUM_INPUT ] . first = get_size_by_dims ( output_dims ) ;
this - > out_sizes [ output_idx - NUM_INPUT ] . second = DataTypeToSize (
this - > engine - > getBindingDataType ( output_idx ) ) ;
const int32_t proto_idx = this - > engine - > getBindingIndex ( PROTO ) ;
const nvinfer1 : : Dims proto_dims = this - > context - > getBindingDimensions ( proto_idx ) ;
this - > out_sizes [ proto_idx - NUM_INPUT ] . first = get_size_by_dims ( proto_dims ) ;
this - > out_sizes [ proto_idx - NUM_INPUT ] . second = DataTypeToSize (
this - > engine - > getBindingDataType ( proto_idx ) ) ;
for ( auto & bindings : this - > input_bindings )
{
void * d_ptr ;
CHECK ( cudaMallocAsync (
& d_ptr ,
bindings . size * bindings . dsize ,
this - > stream )
) ;
this - > device_ptrs . push_back ( d_ptr ) ;
}
for ( int i = 0 ; i < NUM_OUTPUT ; i + + )
for ( auto & bindings : this - > output_bindings )
{
const int osize = this - > out_sizes [ i ] . first * out_sizes [ i ] . second ;
CHECK ( cudaHostAlloc ( & this - > outputs [ i ] , osize , 0 ) ) ;
CHECK ( cudaMalloc ( & this - > buffs [ NUM_INPUT + i ] , osize ) ) ;
void * d_ptr , * h_ptr ;
size_t size = bindings . size * bindings . dsize ;
CHECK ( cudaMallocAsync (
& d_ptr ,
size ,
this - > stream )
) ;
CHECK ( cudaHostAlloc (
& h_ptr ,
size ,
0 )
) ;
this - > device_ptrs . push_back ( d_ptr ) ;
this - > host_ptrs . push_back ( h_ptr ) ;
}
if ( warmup )
{
for ( int i = 0 ; i < 10 ; i + + )
{
size_t isize = this - > in_size * sizeof ( float ) ;
auto * tmp = new float [ isize ] ;
CHECK ( cudaMemcpyAsync ( this - > buffs [ 0 ] ,
tmp ,
isize ,
cudaMemcpyHostToDevice ,
this - > stream ) ) ;
for ( auto & bindings : this - > input_bindings )
{
size_t size = bindings . size * bindings . dsize ;
void * h_ptr = malloc ( size ) ;
memset ( h_ptr , 0 , size ) ;
CHECK ( cudaMemcpyAsync (
this - > device_ptrs [ 0 ] ,
h_ptr ,
size ,
cudaMemcpyHostToDevice ,
this - > stream )
) ;
free ( h_ptr ) ;
}
this - > infer ( ) ;
}
printf ( " model warmup 10 times \n " ) ;
@ -129,158 +193,257 @@ void YOLOv8_seg::make_pipe(bool warmup)
}
}
void YOLOv8_seg : : copy_from_Mat ( const cv : : Mat & image )
void YOLOv8_seg : : letterbox (
const cv : : Mat & image ,
cv : : Mat & out ,
cv : : Size & size
)
{
float height = ( float ) image . rows ;
float width = ( float ) image . cols ;
float r = std : : min ( INPUT_H / height , INPUT_W / width ) ;
const float inp_h = size . height ;
const float inp_w = size . width ;
float height = image . rows ;
float width = image . cols ;
int padw = ( int ) std : : round ( width * r ) ;
int padh = ( int ) std : : round ( height * r ) ;
float r = std : : min ( inp_h / height , inp_w / width ) ;
int padw = std : : round ( width * r ) ;
int padh = std : : round ( height * r ) ;
cv : : Mat tmp ;
if ( ( int ) width ! = padw | | ( int ) height ! = padh )
{
cv : : resize ( image , tmp , cv : : Size ( padw , padh ) ) ;
cv : : resize (
image ,
tmp ,
cv : : Size ( padw , padh )
) ;
}
else
{
tmp = image . clone ( ) ;
}
float _ dw= INPUT_W - padw ;
float _ dh= INPUT_H - padh ;
float dw = inp_w - padw ;
float dh = inp_h - padh ;
_ dw/ = 2.0f ;
_ dh/ = 2.0f ;
int top = int ( std : : round ( _ dh- 0.1f ) ) ;
int bottom = int ( std : : round ( _ dh+ 0.1f ) ) ;
int left = int ( std : : round ( _ dw- 0.1f ) ) ;
int right = int ( std : : round ( _ dw+ 0.1f ) ) ;
cv : : copyMakeBorder ( tmp , tmp , top , bottom , left , right , cv : : BORDER_CONSTANT , PAD_COLOR ) ;
cv : : dnn : : blobFromImage ( tmp ,
dw / = 2.0f ;
dh / = 2.0f ;
int top = int ( std : : round ( dh - 0.1f ) ) ;
int bottom = int ( std : : round ( dh + 0.1f ) ) ;
int left = int ( std : : round ( dw - 0.1f ) ) ;
int right = int ( std : : round ( dw + 0.1f ) ) ;
cv : : copyMakeBorder (
tmp ,
tmp ,
top ,
bottom ,
left ,
right ,
cv : : BORDER_CONSTANT ,
{ 114 , 114 , 114 }
) ;
cv : : dnn : : blobFromImage ( tmp ,
out ,
1 / 255.f ,
cv : : Size ( ) ,
cv : : Scalar ( 0 , 0 , 0 ) ,
true ,
false ,
CV_32F ) ;
CHECK ( cudaMemcpyAsync ( this - > buffs [ 0 ] ,
tmp . ptr < float > ( ) ,
this - > in_size * sizeof ( float ) ,
CV_32F
) ;
this - > pparam . ratio = 1 / r ;
this - > pparam . dw = dw ;
this - > pparam . dh = dh ;
this - > pparam . height = height ;
this - > pparam . width = width ; ;
}
void YOLOv8_seg : : copy_from_Mat ( const cv : : Mat & image )
{
cv : : Mat nchw ;
auto & in_binding = this - > input_bindings [ 0 ] ;
auto width = in_binding . dims . d [ 3 ] ;
auto height = in_binding . dims . d [ 2 ] ;
cv : : Size size { width , height } ;
this - > letterbox (
image ,
nchw ,
size
) ;
this - > context - > setBindingDimensions (
0 ,
nvinfer1 : : Dims
{
4 ,
{ 1 , 3 , height , width }
}
) ;
CHECK ( cudaMemcpyAsync (
this - > device_ptrs [ 0 ] ,
nchw . ptr < float > ( ) ,
nchw . total ( ) * nchw . elemSize ( ) ,
cudaMemcpyHostToDevice ,
this - > stream ) ) ;
this - > stream )
) ;
}
this - > ratio = 1 / r ;
this - > dw = _dw ;
this - > dh = _dh ;
this - > w = width ;
this - > h = height ;
void YOLOv8_seg : : copy_from_Mat ( const cv : : Mat & image , cv : : Size & size )
{
cv : : Mat nchw ;
this - > letterbox (
image ,
nchw ,
size
) ;
this - > context - > setBindingDimensions (
0 ,
nvinfer1 : : Dims
{ 4 ,
{ 1 , 3 , size . height , size . width }
}
) ;
CHECK ( cudaMemcpyAsync (
this - > device_ptrs [ 0 ] ,
nchw . ptr < float > ( ) ,
nchw . total ( ) * nchw . elemSize ( ) ,
cudaMemcpyHostToDevice ,
this - > stream )
) ;
}
void YOLOv8_seg : : infer ( )
{
this - > context - > enqueueV2 ( buffs . data ( ) , this - > stream , nullptr ) ;
for ( int i = 0 ; i < NUM_OUTPUT ; i + + )
this - > context - > enqueueV2 (
this - > device_ptrs . data ( ) ,
this - > stream ,
nullptr
) ;
for ( int i = 0 ; i < this - > num_outputs ; i + + )
{
const int osize = this - > out_sizes [ i ] . first * out_sizes [ i ] . second ;
CHECK ( cudaMemcpyAsync ( this - > outputs [ i ] ,
this - > buffs [ NUM_INPUT + i ] ,
size_ tosize = this - > output_binding s [ i ] . size * this - > output_binding s[ i ] . d siz e;
CHECK ( cudaMemcpyAsync ( this - > host_ptr s[ i ] ,
this - > device_ptrs [ i + this - > num_inputs ] ,
osize ,
cudaMemcpyDeviceToHost ,
this - > stream ) ) ;
this - > stream )
) ;
}
cudaStreamSynchronize ( this - > stream ) ;
}
void YOLOv8_seg : : postprocess ( std : : vector < Object > & objs )
void YOLOv8_seg : : postprocess ( std : : vector < Object > & objs ,
float score_thres ,
float iou_thres ,
int topk ,
int seg_channels ,
int seg_h ,
int seg_w
)
{
objs . clear ( ) ;
auto * output = static_cast < float * > ( this - > outputs [ 0 ] ) ; // x0 y0 x1 y1 s l *32
cv : : Mat protos = cv : : Mat ( NUM_SEG_C , SEG_W * SEG_H , CV_32F ,
static_cast < float * > ( this - > outputs [ 1 ] ) ) ;
auto input_h = this - > input_bindings [ 0 ] . dims . d [ 2 ] ;
auto input_w = this - > input_bindings [ 0 ] . dims . d [ 3 ] ;
auto num_anchors = this - > output_bindings [ 0 ] . dims . d [ 1 ] ;
auto num_channels = this - > output_bindings [ 0 ] . dims . d [ 2 ] ;
auto & dw = this - > pparam . dw ;
auto & dh = this - > pparam . dh ;
auto & width = this - > pparam . width ;
auto & height = this - > pparam . height ;
auto & ratio = this - > pparam . ratio ;
auto * output = static_cast < float * > ( this - > host_ptrs [ 0 ] ) ;
cv : : Mat protos = cv : : Mat ( seg_channels , seg_h * seg_w , CV_32F ,
static_cast < float * > ( this - > host_ptrs [ 1 ] ) ) ;
std : : vector < int > labels ;
std : : vector < float > scores ;
std : : vector < cv : : Rect > bboxes ;
std : : vector < cv : : Mat > mask_confs ;
std : : vector < int > indices ;
for ( int i = 0 ; i < NUM_PROPOSAL ; i + + )
for ( int i = 0 ; i < num_anchors ; i + + )
{
float * ptr = output + i * NUM_COLS ;
float * ptr = output + i * num_channels ;
float score = * ( ptr + 4 ) ;
if ( score > CONF_THRES )
if ( score > score_thres )
{
float x0 = * ptr + + - this - > dw ;
float y0 = * ptr + + - this - > dh ;
float x1 = * ptr + + - this - > dw ;
float y1 = * ptr + + - this - > dh ;
float x0 = * ptr + + - dw ;
float y0 = * ptr + + - dh ;
float x1 = * ptr + + - dw ;
float y1 = * ptr + + - dh ;
x0 = clamp ( x0 * this - > ratio , 0.f , this - > w ) ;
y0 = clamp ( y0 * this - > ratio , 0.f , this - > h ) ;
x1 = clamp ( x1 * this - > ratio , 0.f , this - > w ) ;
y1 = clamp ( y1 * this - > ratio , 0.f , this - > h ) ;
x0 = clamp ( x0 * ratio , 0.f , width ) ;
y0 = clamp ( y0 * ratio , 0.f , height ) ;
x1 = clamp ( x1 * ratio , 0.f , width ) ;
y1 = clamp ( y1 * ratio , 0.f , height ) ;
int label = * ( + + ptr ) ;
cv : : Mat mask_conf = cv : : Mat ( 1 , NUM_SEG_C , CV_32F , + + ptr ) ;
cv : : Mat mask_conf = cv : : Mat ( 1 , seg_channels , CV_32F , + + ptr ) ;
mask_confs . push_back ( mask_conf ) ;
labels . push_back ( label ) ;
scores . push_back ( score ) ;
# if defined(BATCHED_NMS)
bboxes . push_back ( cv : : Rect_ < float > ( x0 , y0 , x1 - x0 , y1 - y0 ) ) ;
# else
bboxes . push_back ( cv : : Rect_ < float > ( x0 + label * DIS ,
y0 + label * DIS ,
x1 - x0 ,
y1 - y0 ) ) ;
# endif
}
}
std : : vector < int > indices ;
# if defined(BATCHED_NMS)
cv : : dnn : : NMSBoxesBatched ( bboxes , scores , labels , CONF_THRES , IOU_THRES , indices ) ;
cv : : dnn : : NMSBoxesBatched (
bboxes ,
scores ,
labels ,
score_thres ,
iou_thres ,
indices
) ;
# else
cv : : dnn : : NMSBoxes ( bboxes , scores , CONF_THRES , IOU_THRES , indices ) ;
cv : : dnn : : NMSBoxes (
bboxes ,
scores ,
score_thres ,
iou_thres ,
indices
) ;
# endif
cv : : Mat masks ;
int cnt = 0 ;
for ( auto & i : indices )
{
# if defined(BATCHED_NMS)
if ( cnt > = topk )
{
break ;
}
cv : : Rect tmp = bboxes [ i ] ;
# else
cv : : Rect tmp = { ( int ) ( bboxes [ i ] . x - labels [ i ] * DIS ) ,
( int ) ( bboxes [ i ] . y - labels [ i ] * DIS ) ,
bboxes [ i ] . width ,
bboxes [ i ] . height } ;
# endif
Object obj ;
obj . label = labels [ i ] ;
obj . rect = tmp ;
obj . prob = scores [ i ] ;
masks . push_back ( mask_confs [ i ] ) ;
objs . push_back ( obj ) ;
cnt + = 1 ;
}
cv : : Mat matmulRes = ( masks * protos ) . t ( ) ;
cv : : Mat maskMat = matmulRes . reshape ( indices . size ( ) , { SEG_W , SEG_H } ) ;
cv : : Mat maskMat = matmulRes . reshape ( indices . size ( ) , { seg_w , seg_h } ) ;
std : : vector < cv : : Mat > maskChannels ;
cv : : split ( maskMat , maskChannels ) ;
int scale_dw = this - > dw / INPUT_W * SEG_W ;
int scale_dh = this - > dh / INPUT_H * SEG_H ;
int scale_dw = dw / input_w * seg_w ;
int scale_dh = dh / input_h * seg_h ;
cv : : Rect roi (
scale_dw ,
scale_dh ,
SEG_W - 2 * scale_dw ,
SEG_H - 2 * scale_dh ) ;
seg_w - 2 * scale_dw ,
seg_h - 2 * scale_dh ) ;
for ( int i = 0 ; i < indices . size ( ) ; i + + )
{
@ -288,30 +451,64 @@ void YOLOv8_seg::postprocess(std::vector<Object>& objs)
cv : : exp ( - maskChannels [ i ] , dest ) ;
dest = 1.0 / ( 1.0 + dest ) ;
dest = dest ( roi ) ;
cv : : resize ( dest , mask , cv : : Size ( ( int ) this - > w , ( int ) this - > h ) , cv : : INTER_LINEAR ) ;
objs [ i ] . boxMask = mask ( objs [ i ] . rect ) > MASK_THRES ;
cv : : resize (
dest ,
mask ,
cv : : Size ( ( int ) width , ( int ) height ) ,
cv : : INTER_LINEAR
) ;
objs [ i ] . boxMask = mask ( objs [ i ] . rect ) > 0.5f ;
}
}
static void draw_objects ( const cv : : Mat & image , cv : : Mat & res , const std : : vector < Object > & objs )
void YOLOv8_seg : : draw_objects ( const cv : : Mat & image ,
cv : : Mat & res ,
const std : : vector < Object > & objs ,
const std : : vector < std : : string > & CLASS_NAMES ,
const std : : vector < std : : vector < unsigned int > > & COLORS ,
const std : : vector < std : : vector < unsigned int > > & MASK_COLORS
)
{
res = image . clone ( ) ;
cv : : Mat mask = image . clone ( ) ;
for ( auto & obj : objs )
{
int idx = obj . label ;
cv : : Scalar color = cv : : Scalar ( COLORS [ idx ] [ 0 ] , COLORS [ idx ] [ 1 ] , COLORS [ idx ] [ 2 ] ) ;
cv : : Scalar color = cv : : Scalar (
COLORS [ idx ] [ 0 ] ,
COLORS [ idx ] [ 1 ] ,
COLORS [ idx ] [ 2 ]
) ;
cv : : Scalar mask_color = cv : : Scalar (
MASK_COLORS [ idx % 20 ] [ 0 ] , MASK_COLORS [ idx % 20 ] [ 1 ] , MASK_COLORS [ idx % 20 ] [ 2 ] ) ;
cv : : rectangle ( res , obj . rect , color , 2 ) ;
MASK_COLORS [ idx % 20 ] [ 0 ] ,
MASK_COLORS [ idx % 20 ] [ 1 ] ,
MASK_COLORS [ idx % 20 ] [ 2 ]
) ;
cv : : rectangle (
res ,
obj . rect ,
color ,
2
) ;
char text [ 256 ] ;
sprintf ( text , " %s %.1f%% " , CLASS_NAMES [ idx ] , obj . prob * 100 ) ;
sprintf (
text ,
" %s %.1f%% " ,
CLASS_NAMES [ idx ] . c_str ( ) ,
obj . prob * 100
) ;
mask ( obj . rect ) . setTo ( mask_color , obj . boxMask ) ;
int baseLine = 0 ;
cv : : Size label_size = cv : : getTextSize ( text , cv : : FONT_HERSHEY_SIMPLEX , 0.4 , 1 , & baseLine ) ;
cv : : Size label_size = cv : : getTextSize (
text ,
cv : : FONT_HERSHEY_SIMPLEX ,
0.4 ,
1 ,
& baseLine
) ;
int x = ( int ) obj . rect . x ;
int y = ( int ) obj . rect . y + 1 ;
@ -319,11 +516,30 @@ static void draw_objects(const cv::Mat& image, cv::Mat& res, const std::vector<O
if ( y > res . rows )
y = res . rows ;
cv : : rectangle ( res , cv : : Rect ( x , y , label_size . width , label_size . height + baseLine ) , RECT_COLOR , - 1 ) ;
cv : : putText ( res , text , cv : : Point ( x , y + label_size . height ) ,
cv : : FONT_HERSHEY_SIMPLEX , 0.4 , TXT_COLOR , 1 ) ;
cv : : rectangle (
res ,
cv : : Rect ( x , y , label_size . width , label_size . height + baseLine ) ,
{ 0 , 0 , 255 } ,
- 1
) ;
cv : : putText (
res ,
text ,
cv : : Point ( x , y + label_size . height ) ,
cv : : FONT_HERSHEY_SIMPLEX ,
0.4 ,
{ 255 , 255 , 255 } ,
1
) ;
}
cv : : addWeighted ( res , 0.5 , mask , 0.8 , 1 , res ) ;
cv : : addWeighted (
res ,
0.5 ,
mask ,
0.8 ,
1 ,
res
) ;
}
# endif //SEGMENT_YOLOV8_SEG_HPP
triple-Mu
2 years ago