python爬虫菜鸟教程，CenterNet 模型后处理 (C++和python代码)

    float hm_chn0[Height][Width] = {0};float hm_chn1[Height][Width] = {0};float reg_chn0[Height][Width] = {0};float reg_chn1[Height][Width] = {0};float wh_chn0[Height][Width] = {0};float wh_chn1[Height][Width] = {0};// if (desc->astMI_OutputTensorDescs[0].eElmFormat == MI_IPU_FORMAT_FP32)// {//     //MI_FLOAT* data = (MI_FLOAT*)(OutputTensorVector->astArrayTensors[0].ptTensorData[0]);//     //memcpy(hmData, OutputTensorVector->astArrayTensors[0].ptTensorData[0],  sizeof(MI_FLOAT)*s32ClassCount );//     //memcpy(whData, OutputTensorVector->astArrayTensors[1].ptTensorData[0],  sizeof(MI_FLOAT)*s32ClassCount );//     //memcpy(regData, OutputTensorVector->astArrayTensors[2].ptTensorData[0], sizeof(MI_FLOAT)*s32ClassCount);// }MI_FLOAT* phmdata = (MI_FLOAT*)(OutputTensorVector->astArrayTensors[0].ptTensorData[0]);MI_FLOAT* pwhdata = (MI_FLOAT*)(OutputTensorVector->astArrayTensors[1].ptTensorData[0]);MI_FLOAT* pregdata = (MI_FLOAT*)(OutputTensorVector->astArrayTensors[2].ptTensorData[0]);for(unsigned int h = 0; h < Height; h++){for(unsigned int w = 0; w < Width; w++)   {//number+=1;if(s32ClassCount%8==0){//取模型推理出的数据，因为模型的输出做了数据对齐，这个要乘以8hm_chn0[h][w] =*(phmdata+((h*Height+w)*8));hm_chn1[h][w] =*(phmdata+((h*Height+w)*8+1));//heatmap_chn0[h][w] = (unsigned char)hm_chn0[h][w]*255;//heatmap_chn1[h][w] = (unsigned char)hm_chn1[h][w]*255;wh_chn0[h][w] =*(pwhdata+((h*Height+w)*8));wh_chn1[h][w] =*(pwhdata+((h*Height+w)*8+1));reg_chn0[h][w] =*(pregdata+((h*Height+w)*8));reg_chn1[h][w] =*(pregdata+((h*Height+w)*8+1));}}}cout<<"hm[0] " << hm_chn0[0][0]<<" " << hm_chn0[0][1]<<endl;//cout<<"hm[0] " << heatmap_chn0[0][0]<<" " << heatmap_chn0[0][1]<<endl;// ofs << std::endl << "}" << std::endl << std::endl;// ofs.close();// Mat heatmap_image(Height, Width, CV_8UC1);// heatmap_image.data = heatmap[0];// imwrite("heat.jpg",heatmap_image);// heatmap_image.data = heatmap[1];// imwrite("heat_1.jpg",heatmap_image);Mat img_hm_0(Height, Width, CV_8UC1);Mat img_hm_1(Height, Width, CV_8UC1);Mat heatmap(Height, Width, CV_8UC1);Mat heatmap_1(Height, Width, CV_8UC1);Mat src(Height, Width, CV_32FC1);float a =0;for (int row = 0;row < Height;row++){for (int col = 0;col < Width;col++){//a = data[(200*row+col)*8] *255 ;a = hm_chn0[row][col]  *255;src.at<float>(row, col)= a ;}}src.convertTo(heatmap,CV_8U,10,0);imwrite("heatmap_0.jpg",heatmap);//src.convertTo(img_hm_0,CV_8U);src.convertTo(img_hm_1,CV_8U);unsigned char c  =  img_hm_0.at<uchar>(0, 0);unsigned char b  =  img_hm_0.at<uchar>(0, 1);cout<<" hm 0 unchar is" << (int)c <<' '<< (int)b << endl;//  对heatmap做maxpool,我用膨胀来代替了cv::Mat mat;//Mat HmImg1(Height, Width, CV_8UC1);Mat Hmax1(Height, Width, CV_8UC1);// Mat HmImg2(Height, Width, CV_8UC1);Mat Hmax2(Height, Width, CV_8UC1);
////HmImg1.data = heatmap[0];Mat hmimg1 = img_hm_0.clone();//HmImg2.data = heatmap[1];Mat hmimg2 = img_hm_0.clone();cv::Mat element = getStructuringElement(MORPH_RECT, Size(3, 3));dilate(hmimg1, Hmax1, element);dilate(hmimg2, Hmax2, element);/////  for (unsigned int i = 0; i < onebuf; i++)// {//     if (hmimg1.data[i] != Hmax1.data[i])//     {//         hm[0][i] = 0;//     }//     if (hmimg2.data[i] != Hmax2.data[i])//     {//         hm[1][i] = 0;//     }// }for (int row = 0;row < Height;row++){for (int col = 0;col < Width;col++){if(hmimg1.at<uchar>(row,col) != Hmax1.at<uchar>(row,col)){hm_chn0[row][col] = 0;}if(hmimg2.at<uchar>(row,col) != Hmax2.at<uchar>(row,col)){hm_chn1[row][col] = 0;}}}// savefor (int row = 0;row < Height;row++){for (int col = 0;col < Width;col++){//a = data[(200*row+col)*8] *255 ;a = hm_chn0[row][col]  *255;src.at<float>(row, col)= a ;}}src.convertTo(heatmap_1,CV_8U,10,0);imwrite("heatmap_1.jpg",heatmap);cout<<"ssssssssss"<<endl;float topk_scores[Chn][topN];int topk_inds[Chn][topN];int topk_ys[Chn][topN];int topk_xs[Chn][topN];// get topN
//   //在两张热度图中分别取最大的10个点//topk(hm[0], onebuf, topN, topk_scores[0], topk_inds[0]);//topk(hm[1], onebuf, topN, topk_scores[1], topk_inds[1]);//在二维数据中 &hm_chn0[1] 表示首行地址 hm_chn0+1 代表第一行地址，  hm_chn0[0] 和 *hm_chn0 表示首行元素地址，*(hm_chn1+1)第一行元素地址， *hm_chn0+1 表示首行下一个元素地址topk(hm_chn0[0], onebuf, topN, topk_scores[0], topk_inds[0]);topk(hm_chn1[0], onebuf, topN, topk_scores[1], topk_inds[1]);
//
//float scores[Chn * topN];int num = 0;//求这20个点的坐标for (unsigned int cl = 0; cl < Chn; cl++){for (int n = 0; n < topN; n++){   //坐标 = y*200+x.//存放的是每张图的位置 0-40000topk_inds[cl][n] = topk_inds[cl][n] % (onebuf);//y 0-200topk_ys[cl][n] = (int)(topk_inds[cl][n] / Height);topk_xs[cl][n] = (int)(topk_inds[cl][n] % Height);//取值范围0-1scores[num] = topk_scores[cl][n];num++;//cout <<"scores is  " << scores[num]<<endl;}}float topk_score[topN]; //输出数值最大的10个点的数值//取值范围 0-20 存放的是得分高的位置int topk_ind[topN];//    //在20个点里取10个点topk(scores, Chn * topN, topN, topk_score, topk_ind);cout <<"Top scores is  " << topk_score[0]<<endl;int topk_clses[topN];int topk_y[topN];   int topk_x[topN];   int ind[topN];     for (unsigned int cl = 0; cl < topN; cl++){//取值范围是0-1，就是这最大十个点，所属的类别topk_clses[cl] = (int)(topk_ind[cl] / topN);//10个值的 x,y坐标（中心坐标）相对于200*200来说topk_y[cl] = topk_ys[topk_clses[cl]][topk_ind[cl] % topN];topk_x[cl] = topk_xs[topk_clses[cl]][topk_ind[cl] % topN];//取值范围是0-H*w(40000)ind[cl] = topk_inds[topk_clses[cl]][topk_ind[cl] % topN];}//以上代码求出了在【112,112,2】上的10个最大值坐标和对应的值//对wh，reg处理，引入回归量float feat_reg[onebuf][Chn] = {0};float feat_wh[onebuf][Chn] = {0};float reg_view[topN][Chn];float wh_view[topN][Chn];float x[topN];float y[topN];float dets[topN][6];//这个循环如果并入maxpool的循环里，处理时间会多2ms，所以依然放在这边/// for (unsigned int i = 0; i < onebuf; i++)// {//     feat_reg[i][0] = reg[0][i];//     feat_reg[i][1] = reg[1][i];//     feat_wh[i][0] = wh[0][i];//     feat_wh[i][1] = wh[1][i];// }int inc = 0;for (int row = 0;row < Height;row++){for (int col = 0;col < Width;col++){feat_reg[inc][0] = reg_chn0[row][col];feat_reg[inc][1] = reg_chn1[row][col];feat_wh[inc][0] = wh_chn0[row][col];feat_wh[inc][1] = wh_chn1[row][col];inc+=1;}}/std::vector<int> ids;std::vector<cv::Rect> boxes;std::vector<float> confidences;cout<<"vvvvvvv"<<endl;for (int num = 0; num < topN; num++){//reg 存放的是中心点的的偏移量reg_view[num][0] = feat_reg[ind[num]][0];reg_view[num][1] = feat_reg[ind[num]][1];//坐标增加回归量x[num] = topk_x[num] + reg_view[num][0];y[num] = topk_y[num] + reg_view[num][1];//读取前10个索引对应的whwh_view[num][0] = feat_wh[ind[num]][0];wh_view[num][1] = feat_wh[ind[num]][1];//输出[10,6]的检测结果，其中10是置信top10，6是4(bboxes)+1(scores)+1(clses)dets[num][0] = (x[num] - (wh_view[num][0] / 2)) * 4;if(dets[num][0] < 0){dets[num][0] = 0;}dets[num][1] = (y[num] - (wh_view[num][1] / 2)) * 4;if(dets[num][1] < 0){dets[num][1] = 0;}dets[num][2] = (x[num] + (wh_view[num][0] / 2)) * 4;if(dets[num][2] < 0){dets[num][2] = 0;}dets[num][3] = (y[num] + (wh_view[num][1] / 2)) * 4;if(dets[num][3] < 0){dets[num][3] = 0;}
//dets[num][4] = topk_score[num];dets[num][5] = topk_clses[num];ids.push_back(dets[num][5]);confidences.push_back(dets[num][4]);boxes.emplace_back((int)dets[num][0], (int)dets[num][1], (int)(wh_view[num][0] * 4), (int)(wh_view[num][1] * 4));}//nmscout<<"gggggggggggggggg"<<endl;std::vector<int> indices;float score_threshold = 0.1;float nms_threshold = 0.1;NMSBoxes(boxes, confidences, score_threshold, nms_threshold, indices);//string filename=(string)(pstPreProcessedData->pImagePath);//cv 默认格式bgr,hwccv::Mat img = cv::imread(image_path, -1);if (img.empty()) {std::cout << " error!  image don't exist!" << std::endl;exit(1);}
//    //网络大小int net_w, net_h;net_w = Width * 4;net_h = Height * 4;std::vector<float> Result_str;cout<<"indec size is  "<<(int)(indices.size()) <<endl;for (size_t i = 0; i < indices.size(); ++i){//这4个点都是对于448*448图片来说的int idx = indices[i];cv::Rect box = boxes[idx];float xmin = static_cast<float>(box.x);float ymin = static_cast<float>(box.y);float xmax = xmin + static_cast<float>(box.width);float ymax = ymin + static_cast<float>(box.height);cout << xmin << " " << ymin << " " << xmax << " " << ymax << " " << confidences[idx] << " " << ids[idx] << " ";cout << endl;Result_str.push_back(xmin);Result_str.push_back(ymin);Result_str.push_back(xmax);Result_str.push_back(ymax);Result_str.push_back(confidences[idx]);Result_str.push_back(ids[idx]);//我们把这些点映射回原图if (img.cols > img.rows) //宽大于高{xmin = xmin * img.cols / net_w;xmax = xmax * img.cols / net_w;ymin = (ymin * img.cols / net_w) - ((img.cols - img.rows) / 2);ymax = (ymax * img.cols / net_w) - ((img.cols - img.rows) / 2);if(ymin < 0){ymin = 0;}if(ymax > img.rows){ymax = img.rows;}}else //高大于宽{ymin = ymin * img.rows / net_h;ymax = ymax * img.rows / net_h;xmin = (xmin * img.rows / net_h) - ((img.rows - img.cols) / 2);xmax = (xmax * img.rows / net_h) - ((img.rows - img.cols) / 2);if(xmin < 0){xmin = 0;}if(xmax > img.cols){xmax = img.cols;}}cout << xmin << " " << ymin << " " << xmax << " " << ymax << " " << confidences[idx] << " " << ids[idx] << " ";cout << endl;cv::rectangle(img, Point((int)(xmin),(int)(ymin)), Point((int)(xmax),(int)(ymax)), Scalar(0, 0, 255), 2);std::string save_out ="EEEEEE.jpg";imwrite(save_out, img);}

python版本

import cv2
import numpy as np
import torch
import torch.nn.functional as Ffrom data import CenterAffinedef gather_feature(fmap, index, mask=None, use_transform=False):if use_transform:# change a (N, C, H, W) tenor to (N, HxW, C) shapebatch, channel = fmap.shape[:2]fmap = fmap.view(batch, channel, -1).permute((0, 2, 1)).contiguous()dim = fmap.size(-1)index = index.unsqueeze(len(index.shape)).expand(*index.shape, dim)fmap = fmap.gather(dim=1, index=index)if mask is not None:mask = mask.unsqueeze(2).expand_as(fmap)fmap = fmap[mask]fmap = fmap.reshape(-1, dim)return fmapclass CenterNetDecoder(object):@staticmethoddef decode(fmap, wh, reg=None, cat_spec_wh=False, K=100):r"""decode output feature map to detection resultsArgs:fmap(Tensor): output feature mapwh(Tensor): tensor that represents predicted width-heightreg(Tensor): tensor that represens regression of center pointscat_spec_wh(bool): whether apply gather on tensor `wh` or notK(int): topk value"""batch, channel, height, width = fmap.shapefmap = CenterNetDecoder.pseudo_nms(fmap)scores, index, clses, ys, xs = CenterNetDecoder.topk_score(fmap, K=K)if reg is not None:reg = gather_feature(reg, index, use_transform=True)reg = reg.reshape(batch, K, 2)xs = xs.view(batch, K, 1) + reg[:, :, 0:1]ys = ys.view(batch, K, 1) + reg[:, :, 1:2]else:xs = xs.view(batch, K, 1) + 0.5ys = ys.view(batch, K, 1) + 0.5wh = gather_feature(wh, index, use_transform=True)if cat_spec_wh:wh = wh.view(batch, K, channel, 2)clses_ind = clses.view(batch, K, 1, 1).expand(batch, K, 1, 2).long()wh = wh.gather(2, clses_ind).reshape(batch, K, 2)else:wh = wh.reshape(batch, K, 2)clses = clses.reshape(batch, K, 1).float()scores = scores.reshape(batch, K, 1)half_w, half_h = wh[..., 0:1] / 2, wh[..., 1:2] / 2bboxes = torch.cat([xs - half_w, ys - half_h, xs + half_w, ys + half_h], dim=2)detections = (bboxes, scores, clses)return detections@staticmethoddef transform_boxes(boxes, img_info, scale=1):r"""transform predicted boxes to target boxesArgs:boxes(Tensor): torch Tensor with (Batch, N, 4) shapeimg_info(dict): dict contains all information of original imagescale(float): used for multiscale testing"""boxes = boxes.cpu().numpy().reshape(-1, 4)center = img_info["center"]size = img_info["size"]output_size = (img_info["width"], img_info["height"])src, dst = CenterAffine.generate_src_and_dst(center, size, output_size)trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))coords = boxes.reshape(-1, 2)aug_coords = np.column_stack((coords, np.ones(coords.shape[0])))target_boxes = np.dot(aug_coords, trans.T).reshape(-1, 4)return target_boxes@staticmethoddef pseudo_nms(fmap, pool_size=3):r"""apply max pooling to get the same effect of nmsArgs:fmap(Tensor): output tensor of previous steppool_size(int): size of max-pooling"""pad = (pool_size - 1) // 2fmap_max = F.max_pool2d(fmap, pool_size, stride=1, padding=pad)keep = (fmap_max == fmap).float()return fmap * keep@staticmethoddef topk_score(scores, K=40):"""get top K point in score map"""batch, channel, height, width = scores.shape# get topk score and its index in every H x W(channel dim) feature maptopk_scores, topk_inds = torch.topk(scores.reshape(batch, channel, -1), K)topk_inds = topk_inds % (height * width)topk_ys = (topk_inds / width).int().float()topk_xs = (topk_inds % width).int().float()# get all topk in in a batchtopk_score, index = torch.topk(topk_scores.reshape(batch, -1), K)# div by K because index is grouped by K(C x K shape)topk_clses = (index / K).int()topk_inds = gather_feature(topk_inds.view(batch, -1, 1), index).reshape(batch, K)topk_ys = gather_feature(topk_ys.reshape(batch, -1, 1), index).reshape(batch, K)topk_xs = gather_feature(topk_xs.reshape(batch, -1, 1), index).reshape(batch, K)return topk_score, topk_inds, topk_clses, topk_ys, topk_xs

python爬虫菜鸟教程。