#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 ) { //1 model status if (!m_drop_detector.IsModelLoaded()) { m_pLogger->ERRORINFO( string("tea 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 generate_detect_windows(vector&boxes) vector drop_regions; int region_cnt = generate_detect_windows(drop_regions); if (region_cnt == 0) { stringstream buff_; buff_ << m_imgId << m_dtype_str << "tea detect image regions' size == 0"; m_pLogger->ERRORINFO(buff_.str()); return 1; } else { stringstream bufftmp; bufftmp << m_imgId << m_dtype_str << "tea detect image regions' size = "<INFO(bufftmp.str()); } //5 detect vector droplets_raw; for (auto rect : drop_regions) { Mat roi = m_raw_img(rect); vector head_droplets = m_drop_detector.RunModel(roi, m_pLogger); if (m_pLogger) { stringstream buff_; buff_ << m_imgId << m_dtype_str << "-------crop_rect["<< rect.x<<","<INFO(buff_.str()); } for (Bbox& b : head_droplets) { b.x1 += rect.x; b.x2 += rect.x; b.y1 += rect.y; b.y2 += rect.y; for (int i = 0; i < 5; ++i) { b.ppoint[2 * i] += rect.x; b.ppoint[2 * i + 1] += rect.y; } } if (head_droplets.size()) { droplets_raw.insert( droplets_raw.end(), head_droplets.begin(), head_droplets.end()); } } if (m_pLogger) { stringstream buff_; buff_ << m_imgId<INFO(buff_.str()); } //6 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]); } if (m_pLogger) { stringstream buff_; buff_ << m_imgId << m_dtype_str << "after nms, keep tea number is " << droplets.size(); m_pLogger->INFO(buff_.str()); } 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; } } int CTeaSort::generate_detect_windows(vector&boxes) { boxes.clear(); int grid_row = m_cp.grid_row_cut; int grid_col = m_cp.grid_col_cut; int grid_padding = m_cp.grid_padding_cut; if (m_dtype == img_type::tea_grab) { grid_row = m_cp.grid_row_grab; grid_col = m_cp.grid_col_grab; grid_padding = m_cp.grid_padding_grab; } if (grid_row < 1) { grid_row = 1; } if (grid_col < 1) { grid_col = 1; } if (grid_padding < 0) { grid_padding = 0; } int block_height = int(m_raw_img.rows / (float)grid_row + 0.5); int block_width = int(m_raw_img.cols / (float)grid_col + 0.5); for (int r = 0; r < grid_row; ++r) { for (int c = 0; c < grid_col; ++c) { int x0 = c*block_width - grid_padding; int y0 = r*block_height - grid_padding; int x1 = (c+1)*block_width + grid_padding; int y1 = (r+1)*block_height + grid_padding; if (x0 < 0) { x0 = 0; } if (y0 < 0) { y0 = 0; } if (x1 > m_raw_img.cols) { x1 = m_raw_img.cols; } if (y1 > m_raw_img.rows) { y1 = m_raw_img.rows; } Rect r(x0, y0, x1-x0, y1-y0); boxes.push_back(r); } } return boxes.size(); } }