123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523 |
- #include <opencv.hpp>
- #include <math.h>
- #include <io.h>
- #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<Rect>&boxes)
- vector<Rect> 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 = "<<region_cnt;
- m_pLogger->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<Bbox> 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<Bbox> 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<Bbox> 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());
- }
- }*/
- if (m_pLogger) {
- stringstream buff_;
- buff_ << m_imgId<<m_dtype_str << "image detect over. tea number is " << droplets_raw.size();
- m_pLogger->INFO(buff_.str());
- }
- //6 nms, width(height) filt and area calculation
- vector<Bbox> droplets;
- vector<int> 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<Rect>&drop_regions, //input, detect regions
- vector<Bbox> &droplets_raw //output, detect result
- )
- {
- //return number of detect result
- droplets_raw.clear();
- for (auto rect : drop_regions) {
- Mat roi = raw_img(rect);
- vector<Bbox> 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<Mat> 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<Rect>&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();
- }
- }
|