#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); if (m_cp.image_show) { cv::destroyAllWindows(); imshow_wait("input_img", m_raw_img); } //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()); } if (m_cp.image_show) { cv::Mat rects_img = m_raw_img.clone(); int step_c = int(255 / (float)region_cnt); int step_cc = step_c / 2; int step_ccc = step_cc / 2; int cnt = 0; for (auto&r : drop_regions) { cv::rectangle(rects_img, r, cv::Scalar(step_cc*cnt, step_c*cnt, step_ccc*cnt), 3); cnt += 1; } imshow_wait("regions_img", rects_img); } //5 detect vector droplets_raw; int dn = detect_impl(m_raw_img, drop_regions, droplets_raw); if (dn < 2 && m_dtype == img_type::tea_grab) { //up-down flip cv::Mat flip_img; cv::flip(m_raw_img, flip_img, 0); if (m_cp.image_show) { imshow_wait("flip_img", flip_img); } vector droplets_flip; int dn_flip = detect_impl(flip_img, drop_regions, droplets_flip); for (auto&b: droplets_flip) { int y2 = flip_img.rows - b.y1; int y1 = flip_img.rows - b.y2; b.y1 = y1; b.y2 = y2; for (int i = 0; i < 5; ++i) { b.ppoint[2 * i + 1] = flip_img.rows - b.ppoint[2 * i + 1]; } } if (dn_flip > 0) { droplets_raw.insert( droplets_raw.end(), droplets_flip.begin(), droplets_flip.end()); } } /*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; if (m_dtype == img_type::tea_grab) { //grab double pre_cx, pre_cy; double min_dist_grab = m_cp.min_distance_grab; pre_cx = -min_dist_grab; pre_cy = -min_dist_grab; for (int i = 0; i < droplets.size(); ++i) { if (valid_cnt > 1) { break; } Bbox&b = droplets.at(i); double cx = 0.5*(b.x1 + b.x2); double cy = 0.5*(b.y1 + b.y2); double dist = sqrt((cx - pre_cx)*(cx - pre_cx) + (cy - pre_cy)*(cy - pre_cy)); if (dist < min_dist_grab) { continue; } double grab_x, grab_y; double angle = calalate_angle(b, grab_x, grab_y); //grab point if (valid_cnt == 0) { posinfo.tea_grab_x1 = grab_x; posinfo.tea_grab_y1 = grab_y; posinfo.tea_grab_angle1 = angle; } else { posinfo.tea_grab_x2 = grab_x; posinfo.tea_grab_y2 = grab_y; posinfo.tea_grab_angle2 = angle; } pre_cx = cx; pre_cy = cy; b.status = 1; valid_cnt += 1; } } else { //cut for (int i = 0; i < droplets.size();++i) { if (i > 1) { break; } Bbox&b = droplets.at(i); b.status = 1; // selected double grab_x, grab_y; double angle = calalate_angle(b, grab_x, grab_y); valid_cnt += 1; if (i == 0) { // 切割点是3、4的中间的点 posinfo.tea_cut_x1 = 0.5 * (b.ppoint[4] + b.ppoint[6]); posinfo.tea_cut_y1 = 0.5 * (b.ppoint[5] + b.ppoint[7]); posinfo.tea_cut_angle1 = angle; } else { // 切割点是3、4的中间的点 posinfo.tea_cut_x2 = 0.5 * (b.ppoint[4] + b.ppoint[6]); posinfo.tea_cut_y2 = 0.5 * (b.ppoint[5] + 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(); for (auto& b : droplets) { //rectangle cv::Rect r = cv::Rect(cv::Point2i(b.x1, b.y1), cv::Point2i(b.x2, b.y2)); if (b.status > 0) { cv::rectangle(img_rst, r, cv::Scalar(0, 0, 255),2); } else { cv::rectangle(img_rst, r, cv::Scalar(0, 255, 0),2); } //score char name[256]; cv::Scalar color(120, 120, 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); //points cv::circle(img_rst, cv::Point(int(b.ppoint[0]), int(b.ppoint[1])), 4, cv::Scalar(255, 0, 255), -1, 3, 0); cv::circle(img_rst, cv::Point(int(b.ppoint[2]), int(b.ppoint[3])), 4, cv::Scalar(0, 255, 255), -1, 3, 0); cv::circle(img_rst, cv::Point(int(b.ppoint[4]), int(b.ppoint[5])), 4, cv::Scalar(255, 0, 0), -1, 3, 0); cv::circle(img_rst, cv::Point(int(b.ppoint[6]), int(b.ppoint[7])), 4, cv::Scalar(0, 255, 0), -1, 3, 0); cv::circle(img_rst, cv::Point(int(b.ppoint[8]), int(b.ppoint[9])), 4, cv::Scalar(0, 0, 255), -1, 3, 0); //grab points if (m_dtype == img_type::tea_grab) { double grab_x, grab_y; calalate_angle(b, grab_x, grab_y); //cv::circle(img_rst, cv::Point(int(grab_x), int(grab_y)), 4, cv::Scalar(0, 215, 255), -1, 3, 0); //lines, p4-p5, p5-grab cv::line(img_rst, cv::Point(int(b.ppoint[6]), int(b.ppoint[7])), cv::Point(int(b.ppoint[8]), int(b.ppoint[9])), cv::Scalar(0, 215, 255), 2); cv::line(img_rst, cv::Point(int(b.ppoint[8]), int(b.ppoint[9])), cv::Point(int(grab_x), int(grab_y)), cv::Scalar(0, 215, 255), 2); //line x int radius = 20; int cx = int(grab_x); int cy = int(grab_y); cv::line(img_rst, cv::Point(cx - radius, cy - radius), cv::Point(cx + radius, cy + radius), cv::Scalar(0, 215, 255), 2); cv::line(img_rst, cv::Point(cx - radius, cy + radius), cv::Point(cx + radius, cy - radius), cv::Scalar(0, 215, 255), 2); } //cut points if (m_dtype == img_type::tea_cut) { //lines, p3-p4 cv::line(img_rst, cv::Point(int(b.ppoint[4]), int(b.ppoint[5])), cv::Point(int(b.ppoint[6]), int(b.ppoint[7])), cv::Scalar(0, 215, 255), 2); //line x int cx = int(0.5 * (b.ppoint[4] + b.ppoint[6])); int cy = int(0.5 * (b.ppoint[5] + b.ppoint[7])); int radius = 20; cv::line(img_rst, cv::Point(cx - radius, cy - radius), cv::Point(cx + radius, cy + radius), cv::Scalar(0, 215, 255),2); cv::line(img_rst, cv::Point(cx - radius, cy + radius), cv::Point(cx + radius, cy - radius), cv::Scalar(0, 215, 255),2); } } if (m_cp.image_show) { imshow_wait("result_img", img_rst); } 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; } int CTeaSort::detect_impl( cv::Mat& raw_img, //input, image vector&drop_regions, //input, detect regions vector &droplets_raw //output, detect result ) { //return number of detect result droplets_raw.clear(); for (auto rect : drop_regions) { Mat roi = 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 << "," << rect.y << "," << rect.width << "," << rect.height << "]," << " roi image detect over. tea number is " << head_droplets.size(); m_pLogger->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()); } } return droplets_raw.size(); } double CTeaSort::calalate_angle( Bbox&b, //input double& grab_x, //output double& grab_y //output ) { grab_x = grab_y = 0.0; 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]; if (m_dtype == img_type::tea_grab) { //calculate line of p4 ans p5 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); } //计算抓取点 double pr = (double)m_cp.offset_grab; double dx = pr * sin(angle); double dy = pr * cos(angle); grab_x = x5 + dx; grab_y = y5 + dy; } else { //tea cut, calculate line of p3 ans p4 angle = atan2(x3 - x4, y3 - 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) || (imginfo->channel!=1 && imginfo->channel!=3)) { if (m_pLogger) { m_pLogger->ERRORINFO(m_imgId + m_dtype_str + "input image invalid."); } throw_msg(m_imgId + " invalid input image"); } if (imginfo->channel == 1) { cv::Mat tmp_img = imginfo2mat(imginfo); vector channels; for (size_t i = 0; i < 3; ++i) { channels.push_back(tmp_img); } cv::merge(channels, m_raw_img); } else { 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(); } }