#include #include #include #include "tea_sorter.h" #include "utils.h" using namespace cv; namespace graft_cv{ CTeaSort::CTeaSort( ConfigParam& cp, img_type dtpye, CGcvLogger*pLog) : m_cp(cp), m_dtype(dtpye), m_pLogger(pLog), m_ppImgSaver(0), m_pImginfoRaw(0), m_pImginfoDetected(0) { m_drop_detector = RetinaDrop(m_pLogger, 0.5, 0.5); } CTeaSort::~CTeaSort() { clear_imginfo(); } int CTeaSort::detect( ImgInfo*imginfo, PositionInfo& posinfo, const char* fn ) { //m_head_droplets.clear(); ////0 文件目录有效性检测 //int at = _access(m_temporary_dir.c_str(),0); //int ai = _access(m_image_dir.c_str(),0); //if(at==-1){ // m_pLogger->ERRORINFO( // string("invalid temporary folder: ")+m_temporary_dir); // return 1; //} //if(ai==-1){ // m_pLogger->ERRORINFO( // string("invalid image folder: ")+m_image_dir); // return 1; //} //1 model status if (!m_drop_detector.IsModelLoaded()) { m_pLogger->ERRORINFO( string("drople detect model NOT loaded")); return 1; } //2 update recognize threshold if (m_dtype == img_type::tea_grab) { m_drop_detector.SetThreshold(m_cp.object_threshold_grab, m_cp.nms_threshold_grab); } else { m_drop_detector.SetThreshold(m_cp.object_threshold_cut, m_cp.nms_threshold_cut); } //3 load data load_data(imginfo, fn); //4 detect vector droplets_raw = m_drop_detector.RunModel(m_raw_img,m_pLogger); if (m_pLogger) { stringstream buff_; buff_ << m_imgId<INFO(buff_.str()); } //5 nms, width(height) filt and area calculation vector droplets; vector keep; nms_bbox(droplets_raw, m_drop_detector.GetNmsThreshold(), keep); //width(height) filter for (int i : keep) { droplets.push_back(droplets_raw[i]); } m_pLogger->INFO(string("nms droplets")); int valid_cnt = 0; for (int i = 0; i < droplets.size();++i) { if (i > 1) { break; } Bbox&b = droplets.at(i); double angle = calalate_angle(b); valid_cnt += 1; //grab point if (i == 0) { if (m_dtype == img_type::tea_grab) { posinfo.tea_grab_x1 = b.ppoint[8]; posinfo.tea_grab_y1 = b.ppoint[9]; posinfo.tea_grab_angle1 = angle; } else { posinfo.tea_cut_x1 = b.ppoint[6]; posinfo.tea_cut_y1 = b.ppoint[7]; posinfo.tea_cut_angle1 = angle; } } else { if (m_dtype == img_type::tea_grab) { posinfo.tea_grab_x2 = b.ppoint[8]; posinfo.tea_grab_y2 = b.ppoint[9]; posinfo.tea_grab_angle2 = angle; } else { posinfo.tea_cut_x2 = b.ppoint[6]; posinfo.tea_cut_y2 = b.ppoint[7]; posinfo.tea_cut_angle2 = angle; } } } //6 draw if (m_cp.image_return) { this->clear_imginfo(); cv::Mat img_rst = m_raw_img.clone(); int cnt = 0; for (auto& b : droplets) { char name[256]; cv::Scalar color(20, 0, 0);//bgr sprintf_s(name, "%.2f", b.score); cv::putText(img_rst, name, cv::Point(b.x1, b.y1), cv::FONT_HERSHEY_COMPLEX, 0.7, color, 2); cv::Rect r = cv::Rect(cv::Point2i(b.x1, b.y1), cv::Point2i(b.x2, b.y2)); if (cnt < 2) { cv::rectangle(img_rst, r, cv::Scalar(0, 0, 255)); } else { cv::rectangle(img_rst, r, cv::Scalar(0, 255, 0)); } cv::rectangle(img_rst, r, cv::Scalar(0, 0, 255)); cv::circle(img_rst, cv::Point(int(b.ppoint[0]), int(b.ppoint[1])), 4, cv::Scalar(0, 0, 255), -1, 8, 0); cv::circle(img_rst, cv::Point(int(b.ppoint[2]), int(b.ppoint[3])), 4, cv::Scalar(0, 255, 255), -1, 8, 0); cv::circle(img_rst, cv::Point(int(b.ppoint[4]), int(b.ppoint[5])), 4, cv::Scalar(255, 0, 255), -1, 8, 0); cv::circle(img_rst, cv::Point(int(b.ppoint[6]), int(b.ppoint[7])), 4, cv::Scalar(0, 255, 0), -1, 8, 0); cv::circle(img_rst, cv::Point(int(b.ppoint[8]), int(b.ppoint[9])), 4, cv::Scalar(255, 0, 0), -1, 8, 0); cnt += 1; } m_pImginfoRaw = mat2imginfo(m_raw_img); m_pImginfoDetected = mat2imginfo(img_rst); posinfo.pp_images[0] = m_pImginfoRaw; posinfo.pp_images[1] = m_pImginfoDetected; if (m_ppImgSaver && *m_ppImgSaver) { (*m_ppImgSaver)->saveImage(img_rst, m_imgId + "_rst_0"); } } //拍照无苗, 返回识别结果-1 if (valid_cnt == 0) { return -1; } return 0; } double CTeaSort::calalate_angle(Bbox&b) { double angle = 0.0; float x3,y3,x4,y4,x5,y5; x3 = b.ppoint[4]; y3 = b.ppoint[5]; x4 = b.ppoint[6]; y4 = b.ppoint[7]; x5 = b.ppoint[8]; y5 = b.ppoint[9]; double r45 = sqrt((x4 - x5)*(x4 - x5) + (y4 - y5)*(y4 - y5)); if (r45 < 15.0) { angle = atan2(x5 - x3, y5 - y3); } else { angle = atan2(x5 - x4, y5 - y4); } angle *= (180.0 / 3.1415926); return angle; } int CTeaSort::load_data( ImgInfo*imginfo, const char* fn/* = 0*/) { //数据加载功能实现,并生成imageid,保存原始数据到文件 int rst = 0; //generate image id if (m_dtype == img_type::tea_grab) { m_imgId = getImgId(img_type::tea_grab); m_dtype_str = string(" tea_grab "); } else { m_imgId = getImgId(img_type::tea_cut); m_dtype_str = string(" tea_cut "); } if (imginfo) { if (m_pLogger) { stringstream buff; buff << m_imgId << m_dtype_str << "image, width=" << imginfo->width << "\theight=" << imginfo->height; m_pLogger->INFO(buff.str()); } if (!isvalid(imginfo)) { if (m_pLogger) { m_pLogger->ERRORINFO(m_imgId + m_dtype_str + "input image invalid."); } throw_msg(m_imgId + " invalid input image"); } m_raw_img = imginfo2mat(imginfo); } else { cv::Mat img = imread(fn, cv::IMREAD_COLOR); if (img.empty()) { if (m_pLogger) { m_pLogger->ERRORINFO(m_imgId + m_dtype_str + "input image invalid:" + string(fn)); } throw_msg(m_imgId + m_dtype_str + "invalid input image: " + string(fn)); } if (m_pLogger) { stringstream buff; buff << m_imgId << m_dtype_str << "image, width=" << img.cols << "\theight=" << img.rows; m_pLogger->INFO(buff.str()); } m_raw_img = img.clone(); } //image saver if (m_ppImgSaver && *m_ppImgSaver) { (*m_ppImgSaver)->saveImage(m_raw_img, m_imgId); } return rst; } int CTeaSort::load_model() { bool b = false; if (!m_drop_detector.IsModelLoaded()) { if (m_dtype == img_type::tea_grab) { b = m_drop_detector.LoadModel(m_cp.model_path_grab); } else { b = m_drop_detector.LoadModel(m_cp.model_path_cut); } } else { b = true; } return b ? 0 : 1; } void CTeaSort::clear_imginfo() { if (m_pImginfoDetected) { imginfo_release(&m_pImginfoDetected); m_pImginfoDetected = 0; } if (m_pImginfoRaw) { imginfo_release(&m_pImginfoRaw); m_pImginfoRaw = 0; } } }