#include "yolo.h"
|
|
#include "yolo.hpp"
|
|
#include <fstream>
|
#include <vector>
|
#include <stdexcept>
|
|
#include "struct.h"
|
|
namespace cppyolo
|
{
|
sdkyolo::sdkyolo(const char *cfg, const char *weights, const char *name, const int gpu)
|
:det_(NULL)
|
,infos_(NULL)
|
,obj_cnt_(0)
|
,image_(NULL)
|
,width_(0)
|
,height_(0)
|
,channel_(0)
|
{
|
names_.clear();
|
if (!init(cfg, weights, name, gpu)){
|
throw std::runtime_error("init yolo error");
|
}
|
}
|
|
sdkyolo::~sdkyolo()
|
{
|
if (det_) delete det_;
|
}
|
|
static void objects_names_from_file(std::string filename, std::vector<std::string> &names){
|
std::ifstream file(filename);
|
|
if (!file.is_open()){
|
printf("open %s file error\n", filename.c_str());
|
return;
|
}
|
for(std::string line; getline(file, line);) names.push_back(line);
|
|
printf("names count %d\n", names.size());
|
}
|
|
const char* sdkyolo::obj_name_by_type(const int typ)const{
|
if(names_.empty() || typ < 0 || typ >= names_.size()) return NULL;
|
return names_.at(typ).c_str();
|
}
|
|
bool sdkyolo::init(const char *cfg, const char *weights, const char *name, const int gpu){
|
if (det_) return true;
|
|
if(!cfg || !weights || !name){
|
printf("init Detector error\n");
|
return false;
|
}
|
|
if(names_.empty())
|
objects_names_from_file(name, names_);
|
|
det_ = new Detector(cfg, weights, gpu);
|
return true;
|
}
|
|
int sdkyolo::buffer_to_image(const unsigned char *data, const int w, const int h, const int color_channel){
|
|
int size = w*h;
|
int size2 = size*2;
|
|
int c = color_channel;
|
if (w != width_ || h != height_ || color_channel != channel_){
|
if (image_){
|
free(image_->data);
|
delete image_;
|
}
|
image_ = new image_t;
|
image_->h = h;
|
image_->w = w;
|
image_->c = c;
|
image_->data = (float*)calloc(h*w*c, sizeof(float));
|
}
|
|
// image im = make_image(w, h, c);
|
const unsigned char *srcData = data;
|
|
int count = 0;
|
switch(c){
|
case 1:{
|
for (; count < size; ){
|
image_->data[count] =
|
image_->data[w*h + count] =
|
image_->data[w*h*2 + count] =
|
(float)(srcData[count])/255.0;
|
|
++count;
|
}
|
break;
|
}
|
case 3:{
|
float* desData = image_->data;
|
|
for(int i = 0;i<size;i++){
|
*(desData) = *(srcData + 2) /255.0f;
|
*(desData+size) = *(srcData + 1) /255.0f;
|
*(desData+size2) = *(srcData) /255.0f;
|
|
desData++;
|
srcData+=3;
|
}
|
break;
|
}
|
|
default:
|
printf("Channel number not supported.\n");
|
break;
|
}
|
|
return 0;
|
}
|
|
int sdkyolo::detect(const cIMAGE *img, const float thrsh, const bool use_mean, void **objs, int *objCount){
|
if (!det_) return -1;
|
|
const int color_channel = img->channel;
|
buffer_to_image(img->data, img->width, img->height, color_channel);
|
|
|
std::vector<bbox_t> result_vec = det_->detect(*image_, thrsh, use_mean);
|
// det->free_image(*im);
|
// delete im;
|
|
if (obj_cnt_ < result_vec.size()){
|
free(infos_);
|
|
obj_cnt_ = result_vec.size();
|
infos_ = (cObjInfo*)malloc(obj_cnt_ * sizeof(cObjInfo));
|
}
|
|
int count = 0;
|
for(auto &i : result_vec){
|
|
cObjInfo info;
|
info.typ = i.obj_id;
|
info.prob = i.prob;
|
info.rcObj.left = i.x;
|
info.rcObj.top = i.y;
|
info.rcObj.right = i.x+i.w;
|
info.rcObj.bottom = i.y+i.h;
|
|
infos_[count++] = info;
|
}
|
|
*objCount = count;
|
*objs = infos_;
|
|
return count;
|
}
|
|
} // namespace cppyolo
|
