Micaela Verucchi
412 lines
15 KiB
C++
412 lines
15 KiB
C++
#include "MobilenetDetection.h"
|
|
|
|
bool boxProbCmp(const tk::dnn::box &a, const tk::dnn::box &b)
|
|
{
|
|
return (a.prob > b.prob);
|
|
}
|
|
|
|
namespace tk
|
|
{
|
|
namespace dnn
|
|
{
|
|
|
|
void MobilenetDetection::generate_ssd_priors(const SSDSpec *specs, const int n_specs, bool clamp)
|
|
{
|
|
n_priors = 0;
|
|
for (int i = 0; i < n_specs; i++)
|
|
{
|
|
n_priors += specs[i].featureSize * specs[i].featureSize * 6;
|
|
}
|
|
|
|
// std::cout<<"n priors: "<<n_priors<<std::endl;
|
|
// std::cout<<"n priors: "<<n_specs<<std::endl;
|
|
|
|
priors = (float *)malloc(N_COORDS * n_priors * sizeof(float));
|
|
|
|
int i_prio = 0;
|
|
float scale, x_center, y_center, h, w, size, ratio;
|
|
int min, max;
|
|
for (int i = 0; i < n_specs; i++)
|
|
{
|
|
scale = (float)imageSize / (float)specs[i].shrinkage;
|
|
min = specs[i].boxHeight > specs[i].boxWidth ? specs[i].boxWidth : specs[i].boxHeight;
|
|
max = specs[i].boxHeight < specs[i].boxWidth ? specs[i].boxWidth : specs[i].boxHeight;
|
|
for (int j = 0; j < specs[i].featureSize; j++)
|
|
{
|
|
for (int k = 0; k < specs[i].featureSize; k++)
|
|
{
|
|
//small sized square box
|
|
size = min;
|
|
x_center = (k + 0.5f) / scale;
|
|
y_center = (j + 0.5f) / scale;
|
|
h = w = (float)size / (float)imageSize;
|
|
|
|
priors[i_prio * N_COORDS + 0] = x_center;
|
|
priors[i_prio * N_COORDS + 1] = y_center;
|
|
priors[i_prio * N_COORDS + 2] = w;
|
|
priors[i_prio * N_COORDS + 3] = h;
|
|
++i_prio;
|
|
|
|
//big sized square box
|
|
size = sqrt(max * min);
|
|
h = w = (float)size / (float)imageSize;
|
|
|
|
priors[i_prio * N_COORDS + 0] = x_center;
|
|
priors[i_prio * N_COORDS + 1] = y_center;
|
|
priors[i_prio * N_COORDS + 2] = w;
|
|
priors[i_prio * N_COORDS + 3] = h;
|
|
++i_prio;
|
|
|
|
//change h/w ratio of the small sized box
|
|
size = min;
|
|
h = w = size / (float)imageSize;
|
|
ratio = sqrt(specs[i].ratio1);
|
|
priors[i_prio * N_COORDS + 0] = x_center;
|
|
priors[i_prio * N_COORDS + 1] = y_center;
|
|
priors[i_prio * N_COORDS + 2] = w * ratio;
|
|
priors[i_prio * N_COORDS + 3] = h / ratio;
|
|
++i_prio;
|
|
|
|
priors[i_prio * N_COORDS + 0] = x_center;
|
|
priors[i_prio * N_COORDS + 1] = y_center;
|
|
priors[i_prio * N_COORDS + 2] = w / ratio;
|
|
priors[i_prio * N_COORDS + 3] = h * ratio;
|
|
++i_prio;
|
|
|
|
ratio = sqrt(specs[i].ratio2);
|
|
priors[i_prio * N_COORDS + 0] = x_center;
|
|
priors[i_prio * N_COORDS + 1] = y_center;
|
|
priors[i_prio * N_COORDS + 2] = w * ratio;
|
|
priors[i_prio * N_COORDS + 3] = h / ratio;
|
|
++i_prio;
|
|
|
|
priors[i_prio * N_COORDS + 0] = x_center;
|
|
priors[i_prio * N_COORDS + 1] = y_center;
|
|
priors[i_prio * N_COORDS + 2] = w / ratio;
|
|
priors[i_prio * N_COORDS + 3] = h * ratio;
|
|
++i_prio;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (clamp)
|
|
{
|
|
for (int i = 0; i < n_priors * N_COORDS; i++)
|
|
{
|
|
priors[i] = priors[i] > 1.0f ? 1.0f : priors[i];
|
|
priors[i] = priors[i] < 0.0f ? 0.0f : priors[i];
|
|
|
|
// std::cout<<priors[i]<<" ";
|
|
// if((i+1)%4 == 0)
|
|
// std::cout<< i/4 <<" " <<std::endl;
|
|
}
|
|
}
|
|
}
|
|
|
|
void MobilenetDetection::convert_locatios_to_boxes_and_center(float *priors, const int n_priors, float *locations, const float centerVariance, const float sizeVariance)
|
|
{
|
|
float cur_x, cur_y;
|
|
for (int i = 0; i < n_priors; i++)
|
|
{
|
|
locations[i * N_COORDS + 0] = locations[i * N_COORDS + 0] * centerVariance * priors[i * N_COORDS + 2] + priors[i * N_COORDS + 0];
|
|
locations[i * N_COORDS + 1] = locations[i * N_COORDS + 1] * centerVariance * priors[i * N_COORDS + 3] + priors[i * N_COORDS + 1];
|
|
locations[i * N_COORDS + 2] = exp(locations[i * N_COORDS + 2] * sizeVariance) * priors[i * N_COORDS + 2];
|
|
locations[i * N_COORDS + 3] = exp(locations[i * N_COORDS + 3] * sizeVariance) * priors[i * N_COORDS + 3];
|
|
|
|
cur_x = locations[i * N_COORDS + 0];
|
|
cur_y = locations[i * N_COORDS + 1];
|
|
|
|
locations[i * N_COORDS + 0] = cur_x - locations[i * N_COORDS + 2] / 2;
|
|
locations[i * N_COORDS + 1] = cur_y - locations[i * N_COORDS + 3] / 2;
|
|
locations[i * N_COORDS + 2] = cur_x + locations[i * N_COORDS + 2] / 2;
|
|
locations[i * N_COORDS + 3] = cur_y + locations[i * N_COORDS + 3] / 2;
|
|
|
|
// std::cout<<locations[i*N_COORDS + 0]<<" "<<locations[i*N_COORDS + 1]<<" "<<locations[i*N_COORDS + 2]<<" "<<locations[i*N_COORDS + 3]<<" "<<std::endl;
|
|
}
|
|
}
|
|
|
|
float MobilenetDetection::iou(const tk::dnn::box &a, const tk::dnn::box &b)
|
|
{
|
|
float max_x = a.x > b.x ? a.x : b.x;
|
|
float max_y = a.y > b.y ? a.y : b.y;
|
|
float min_w = a.w < b.w ? a.w : b.w;
|
|
float min_h = a.h < b.h ? a.h : b.h;
|
|
|
|
float ao_w = min_w - max_x > 0 ? min_w - max_x : 0;
|
|
float ao_h = min_h - max_y > 0 ? min_h - max_y : 0;
|
|
|
|
// std::cout<<" ao w: "<<ao_w<<" ao h: "<<ao_h<<std::endl;
|
|
|
|
float area_overlap = ao_w * ao_h;
|
|
float area_0_w = a.w - a.x > 0 ? a.w - a.x : 0;
|
|
float area_0_h = a.h - a.y > 0 ? a.h - a.y : 0;
|
|
|
|
float area_1_w = b.w - b.x > 0 ? b.w - b.x : 0;
|
|
float area_1_h = b.h - b.y > 0 ? b.h - b.y : 0;
|
|
|
|
float area_0 = area_0_h * area_0_w;
|
|
float area_1 = area_1_h * area_1_w;
|
|
|
|
// std::cout<<" area_overlap : "<<area_overlap<<" area_0: "<<area_0<<" area_1: "<<area_1<<std::endl;
|
|
|
|
float iou = area_overlap / (area_0 + area_1 - area_overlap + 1e-5);
|
|
return iou;
|
|
}
|
|
|
|
std::vector<tk::dnn::box> MobilenetDetection::postprocess(float *locations, float *confidences, const int n_values, const float threshold, const int n_classes, const float iou_thresh, const int width, const int height)
|
|
{
|
|
float *conf_per_class;
|
|
std::vector<tk::dnn::box> detections;
|
|
for (int i = 1; i < n_classes; i++)
|
|
{
|
|
conf_per_class = &confidences[i * n_values];
|
|
std::vector<tk::dnn::box> boxes;
|
|
for (int j = 0; j < n_values; j++)
|
|
{
|
|
|
|
if (conf_per_class[j] > threshold)
|
|
{
|
|
tk::dnn::box b;
|
|
b.cl = i;
|
|
b.prob = conf_per_class[j];
|
|
b.x = locations[j * N_COORDS + 0];
|
|
b.y = locations[j * N_COORDS + 1];
|
|
b.w = locations[j * N_COORDS + 2];
|
|
b.h = locations[j * N_COORDS + 3];
|
|
|
|
boxes.push_back(b);
|
|
}
|
|
}
|
|
std::sort(boxes.begin(), boxes.end(), boxProbCmp);
|
|
// for(auto b:boxes)
|
|
// b.print();
|
|
std::vector<tk::dnn::box> remaining;
|
|
while (boxes.size() > 0)
|
|
{
|
|
remaining.clear();
|
|
|
|
tk::dnn::box b;
|
|
b.cl = boxes[0].cl;
|
|
b.prob = boxes[0].prob;
|
|
b.x = boxes[0].x * width;
|
|
b.y = boxes[0].y * height;
|
|
b.w = boxes[0].w * width;
|
|
b.h = boxes[0].h * height;
|
|
detections.push_back(b);
|
|
for (size_t j = 1; j < boxes.size(); j++)
|
|
{
|
|
if (iou(boxes[0], boxes[j]) <= iou_thresh)
|
|
{
|
|
remaining.push_back(boxes[j]);
|
|
}
|
|
}
|
|
boxes = remaining;
|
|
}
|
|
}
|
|
// std::cout<<"picked"<<std::endl;
|
|
// for(auto b:detections)
|
|
// b.print();
|
|
|
|
return detections;
|
|
}
|
|
|
|
float MobilenetDetection::get_color2(int c, int x, int max)
|
|
{
|
|
float ratio = ((float)x / max) * 5;
|
|
int i = floor(ratio);
|
|
int j = ceil(ratio);
|
|
ratio -= i;
|
|
float r = (1 - ratio) * __colors[i % 6][c % 3] + ratio * __colors[j % 6][c % 3];
|
|
//printf("%f\n", r);
|
|
return r;
|
|
}
|
|
|
|
void MobilenetDetection::init(std::string tensor_path, int input_size, int n_classes)
|
|
{
|
|
this->imageSize = input_size;
|
|
this->classes = n_classes;
|
|
|
|
const int n_SSDSpec = 6;
|
|
SSDSpec specs[6];
|
|
|
|
|
|
if(input_size == 300)
|
|
{
|
|
specs[0].setAll(19, 16, 60, 105, 2, 3);
|
|
specs[1].setAll(10, 32, 105, 150, 2, 3);
|
|
specs[2].setAll(5, 64, 150, 195, 2, 3);
|
|
specs[3].setAll(3, 100, 195, 240, 2, 3);
|
|
specs[4].setAll(2, 150, 240, 285, 2, 3);
|
|
specs[5].setAll(1, 300, 285, 330, 2, 3);
|
|
}
|
|
else if(input_size == 512)
|
|
{
|
|
specs[0].setAll(32, 16, 60, 105, 2, 3);
|
|
specs[1].setAll(16, 32, 105, 150, 2, 3);
|
|
specs[2].setAll(8, 64, 150, 195, 2, 3);
|
|
specs[3].setAll(4, 100, 195, 240, 2, 3);
|
|
specs[4].setAll(2, 150, 240, 285, 2, 3);
|
|
specs[5].setAll(1, 300, 285, 330, 2, 3);
|
|
}
|
|
else
|
|
{
|
|
FatalError("Input size for mobilenet not supported");
|
|
}
|
|
|
|
|
|
generate_ssd_priors(specs, n_SSDSpec);
|
|
|
|
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
|
|
|
|
checkCuda(cudaMallocHost(&input, sizeof(dnnType) * netRT->input_dim.tot()));
|
|
checkCuda(cudaMalloc(&input_d, sizeof(dnnType) * netRT->input_dim.tot()));
|
|
|
|
locations_h = (float *)malloc(N_COORDS * n_priors * sizeof(float));
|
|
confidences_h = (float *)malloc(n_priors * classes * sizeof(float));
|
|
|
|
dim = tk::dnn::dataDim_t(1, 3, imageSize, imageSize, 1);
|
|
|
|
for (int c = 0; c < classes; c++)
|
|
{
|
|
int offset = c * 123457 % classes;
|
|
float r = get_color2(2, offset, classes);
|
|
float g = get_color2(1, offset, classes);
|
|
float b = get_color2(0, offset, classes);
|
|
colors[c] = cv::Scalar(int(255.0 * b), int(255.0 * g), int(255.0 * r));
|
|
}
|
|
|
|
if(classes == 21)
|
|
{
|
|
const char *classes_names_[] = {
|
|
"BACKGROUND", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
|
|
"car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike",
|
|
"person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"};
|
|
classesNames = std::vector<std::string>(classes_names_, std::end(classes_names_));
|
|
|
|
}
|
|
else if (classes == 81)
|
|
{
|
|
const char *classes_names_[] = {
|
|
"BACKGROUND", "person" , "bicycle" , "car" , "motorbike" , "aeroplane" , "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" , "sofa" , "pottedplant" , "bed" , "diningtable" ,
|
|
"toilet" , "tvmonitor" , "laptop" , "mouse" , "remote" , "keyboard" ,
|
|
"cell phone" , "microwave" , "oven" , "toaster" , "sink" , "refrigerator" ,
|
|
"book" , "clock" , "vase" , "scissors" , "teddy bear" , "hair drier" , "toothbrush"};
|
|
classesNames = std::vector<std::string>(classes_names_, std::end(classes_names_));
|
|
|
|
}
|
|
else
|
|
{
|
|
FatalError("Number of classes not supported for mobilenet");
|
|
}
|
|
|
|
|
|
}
|
|
|
|
cv::Mat MobilenetDetection::draw()
|
|
{
|
|
tk::dnn::box b;
|
|
for (size_t i = 0; i < detected.size(); i++)
|
|
{
|
|
b = detected[i];
|
|
std::string det_class = classesNames[b.cl];
|
|
cv::rectangle(origImg, cv::Point(b.x, b.y), cv::Point(b.w, b.h), colors[b.cl], 2);
|
|
// draw label
|
|
cv::Size textSize = getTextSize(det_class, cv::FONT_HERSHEY_SIMPLEX, fontScale, thickness, &baseline);
|
|
cv::rectangle(origImg, cv::Point(b.x, b.y), cv::Point((b.x + textSize.width - 2), (b.y - textSize.height - 2)), colors[b.cl], -1);
|
|
cv::putText(origImg, det_class, cv::Point(b.x, (b.y - (baseline / 2))), cv::FONT_HERSHEY_SIMPLEX, fontScale, cv::Scalar(255, 255, 255), thickness);
|
|
}
|
|
return origImg;
|
|
}
|
|
|
|
|
|
void MobilenetDetection::preprocess(const bool gpu)
|
|
{
|
|
std::cout<<"preprocess"<<std::endl;
|
|
if(gpu){
|
|
cv::cuda::GpuMat im_Orig;
|
|
cv::cuda::GpuMat frame_resize, frame_nomean, frame_scaled;
|
|
|
|
im_Orig = cv::cuda::GpuMat(origImg);
|
|
cv::cuda::resize (im_Orig, frame_resize, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
|
|
|
|
// resize(origImg, frame_resize, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
|
|
frame_resize.convertTo(frame_nomean, CV_32FC3, 1, -127);
|
|
frame_nomean.convertTo(frame_scaled, CV_32FC3, 1 / 128.0, 0);
|
|
|
|
//copy image into tensor and copy it into GPU
|
|
cv::cuda::GpuMat bgr[3];
|
|
cv::cuda::split(frame_scaled, bgr);
|
|
|
|
for(int i=0; i < netRT->input_dim.c; i++){
|
|
int idx = i * frame_scaled.rows * frame_scaled.cols;
|
|
checkCuda( cudaMemcpy((void *)&input_d[idx], (void *)bgr[i].data, frame_scaled.rows * frame_scaled.cols* sizeof(float), cudaMemcpyDeviceToDevice) );
|
|
}
|
|
}
|
|
else{
|
|
//resize image, remove mean, divide by std
|
|
cv::Mat frame_resize, frame_nomean, frame_scaled;
|
|
resize(origImg, frame_resize, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
|
|
frame_resize.convertTo(frame_nomean, CV_32FC3, 1, -127);
|
|
frame_nomean.convertTo(frame_scaled, CV_32FC3, 1 / 128.0, 0);
|
|
|
|
//copy image into tensor and copy it into GPU
|
|
cv::Mat bgr[3];
|
|
cv::split(frame_scaled, bgr);
|
|
for (int i = 0; i < netRT->input_dim.c; i++){
|
|
int idx = i * frame_scaled.rows * frame_scaled.cols;
|
|
memcpy((void *)&input[idx], (void *)bgr[i].data, frame_scaled.rows * frame_scaled.cols * sizeof(dnnType));
|
|
}
|
|
checkCuda(cudaMemcpyAsync(input_d, input, netRT->input_dim.tot() * sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
|
|
}
|
|
|
|
|
|
}
|
|
|
|
void MobilenetDetection::update(cv::Mat &img)
|
|
{
|
|
TIMER_START
|
|
detected.clear();
|
|
|
|
//save origin image
|
|
origImg = img;
|
|
cv::Size sz = origImg.size();
|
|
|
|
//preprocess
|
|
preprocess();
|
|
|
|
//do inference
|
|
tk::dnn::dataDim_t dim2 = dim;
|
|
printCenteredTitle(" TENSORRT inference ", '=', 30);
|
|
{
|
|
dim2.print();
|
|
TIMER_START
|
|
netRT->infer(dim2, input_d);
|
|
TIMER_STOP
|
|
dim2.print();
|
|
}
|
|
|
|
//get confidences and locations
|
|
conf = (dnnType *)netRT->buffersRT[3];
|
|
loc = (dnnType *)netRT->buffersRT[4];
|
|
|
|
checkCuda(cudaMemcpy(confidences_h, conf, n_priors * classes * sizeof(float), cudaMemcpyDeviceToHost));
|
|
checkCuda(cudaMemcpy(locations_h, loc, N_COORDS * n_priors * sizeof(float), cudaMemcpyDeviceToHost));
|
|
|
|
//postprocess
|
|
convert_locatios_to_boxes_and_center(priors, n_priors, locations_h, centerVariance, sizeVariance);
|
|
detected = postprocess(locations_h, confidences_h, n_priors, confThresh, classes, iouThreshold, sz.width, sz.height);
|
|
|
|
TIMER_STOP
|
|
stats.push_back(t_ns);
|
|
}
|
|
|
|
} // namespace dnn
|
|
} // namespace tk
|