开发笔记:手撕coreML之yolov2objectdetection物体检测（含源代码）

execute download.sh to download the pre-trained model % sh download.sh

- (nullable yoloModelOutput *)predictionFromInput__0:(CVPixelBufferRef)input__0 error:(NSError * _Nullable * _Nullable)error;

- (void)viewDidLoad {
[super viewDidLoad];
// Do any additional setup after loading the view, typically from a nib.

//load image
NSString* imagePath_=[[NSBundle mainBundle] pathForResource:@"dog416" ofType:@"jpg"];
std::string imgPath = std::string([imagePath_ UTF8String]);
cv::Mat image = cv::imread(imgPath);
cv::cvtColor(image, image, CV_BGR2RGBA);


//set classtxt path
NSString* classtxtPath_ = [ [NSBundle mainBundle] pathForResource:@"classtxt" ofType:@"txt"];
std::string classtxtPath = std::string([classtxtPath_ UTF8String]);

//init Detection
bool useCpuOny= false;
MLComputeUnits computeUnit = MLComputeUnitsAll;
cv::Size scaleSize(608, 608);
CDetectObject objectDetection;
objectDetection.init(useCpuOny, computeUnit, classtxtPath, scaleSize);
//run detection
std::vector detectionResults;
objectDetection.implDetection(image, detectionResults);

//draw rectangles
cv::Mat showImage;
cv::resize(image, showImage, scaleSize);
for (int i=0; i)
{
cv::rectangle(showImage,detectionResults[i].box, cv::Scalar(255, 0,0), 3);
}

//show in iphone
cv::cvtColor(showImage, showImage, cv::COLOR_RGBA2BGRA);
[self showUIImage:showImage];
}

struct DetectionInfo {
std::string name;
float confidence;
cv::Rect2d box;
};

//init model
int CDetectObject::init(const BOOL useCpuOnly, const MLComputeUnits computeUnit, const std::string& classtxtPath, const cv::Size& scaleSize){

//init configuration
option = [[MLPredictionOptions alloc] init];
option.usesCPUOnly = useCpuOnly;

config = [ [MLModelConfiguration alloc] init];
config.computeUnits = computeUnit;

NSError* err;
Model = [[yoloModel alloc] initWithConfiguration:config error:&err];

//init paramss
inputSize = scaleSize;
maxBoundingBoxes = 10;
confidenceThreshold = 0.5;
nmsThreshold = 0.6;
// anchor boxes
anchors = {0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828};

//load labels
int ret = loadClasstxt(classtxtPath, classes);

return ret;
}

int CDetectObject::implDetection(const cv::Mat& image, std::vector& detectionResults){

if(image.empty()){
NSLog(@"Error! image is empty!");
return -1;
}

//preprocessing
cv::Mat inputImage;
preprocessImage(image, inputImage);

//prediction
MLMultiArray* outFeature = predictImageScene(inputImage);

//analyze the output
std::vector<int> idxList;
std::vector<float> confidenceList;
std::vector boxesList;
parseFeature(outFeature, idxList, confidenceList, boxesList);

//nms box
std::vector<int> indices;
cv::dnn::NMSBoxes(boxesList, confidenceList, confidenceThreshold, nmsThreshold, indices);

//get result
for (int i=0; i){
int idx = indices[i];
DetectionInfo objectInfo;
objectInfo.name = classes[idxList[idx]];
objectInfo.confidence = confidenceList[idx];
objectInfo.box = boxesList[idx];
detectionResults.push_back(objectInfo);
}

return 0;
}

MLMultiArray* CDetectObject::predictImageScene(const cv::Mat& imgTensor) {
//preprocess image

//convert to cvPixelbuffer
ins::PixelBufferPool mat2pixelbuffer;
CVPixelBufferRef buffer = mat2pixelbuffer.GetPixelBuffer(imgTensor);

//predict from image
NSError *error;
yoloModelInput *input = [[yoloModelInput alloc] initWithInput__0:buffer];

yoloModelOutput *output = [Model predictionFromFeatures:input options:option error:&error];

return output.output__0;
}

void CDetectObject::parseFeature(MLMultiArray* feature, std::vector<int>& ids, std::vector<float>& confidences, std::vector& boxes){

NSArray* featureShape = feature.shape;
int d0 = [[featureShape objectAtIndex:0] intValue];
int d1 = [[featureShape objectAtIndex:1] intValue];
int d2 = [[featureShape objectAtIndex:2] intValue];

int stride0 = [feature.strides[0] intValue];
int stride1 = [feature.strides[1] intValue];
int stride2 = [feature.strides[2] intValue];

int blockSize = 32;
int gridHeight = d1;
int gridWidth = d2;
int boxesPerCell = 5;//Int(anchors.count/5)
int numClasses = (int)classes.size();

double* pdata = (double*)feature.dataPointer;

for (int cy =0; cy){
for (int cx =0; cx){
for (int b=0; b){
int channel = b*(numClasses + 5);

int laterId= cx*stride2+cy*stride1;
float tx = (float)pdata[channel*stride0 + laterId];
float ty = (float)pdata[(channel+1)*stride0 + laterId];
float tw = (float)pdata[(channel+2)*stride0 + laterId];
float th = (float)pdata[(channel+3)*stride0 + laterId];
float tc = (float)pdata[(channel+4)*stride0 + laterId];

// The predicted tx and ty coordinates are relative to the location
// of the grid cell; we use the logistic sigmoid to constrain these
// coordinates to the range 0 - 1. Then we add the cell coordinates
// (0-12) and multiply by the number of pixels per grid cell (32).
// Now x and y represent center of the bounding box in the original
// 608x608 image space.
float x = (float(cx) + sigmoid(tx)) * blockSize;
float y = (float(cy) + sigmoid(ty)) * blockSize;

// The size of the bounding box, tw and th, is predicted relative to
// the size of an "anchor" box. Here we also transform the width and
// height into the original 608x608 image space.
float w = exp(tw) * anchors[2*b] * blockSize;
float h = exp(th) * anchors[2*b + 1] * blockSize;

// The confidence value for the bounding box is given by tc. We use
// the logistic sigmoid to turn this into a percentage.
float cOnfidence= sigmoid(tc);
std::vector<float> classesProb(numClasses);
for (int i = 0; i i) {
int offset = (channel+5+i)*stride0 + laterId;
classesProb[i] = (float)pdata[offset];
}
softmax(classesProb);

// Find the index of the class with the largest score.
auto max_itr = std::max_element(classesProb.begin(), classesProb.end());
int index = int(max_itr - classesProb.begin());
// Combine the confidence score for the bounding box, which tells us
// how likely it is that there is an object in this box (but not what
// kind of object it is), with the largest class prediction, which
// tells us what kind of object it detected (but not where).
float cOnfidenceInClass= classesProb[index] * confidence;
if(confidence>confidenceThreshold){
// Since we compute 19x19x5 = 1805 bounding boxes, we only want to
// keep the ones whose combined score is over a certain threshold.
//if (confidenceInClass > confidenceThreshold){
cv::Rect2d rect =cv::Rect2d(float(x-w*0.5), float(y-h*0.5), float(w), float(h));
ids.push_back(index);
confidences.push_back(confidenceInClass);
boxes.push_back(rect);
}
}
}
}
}