tea_sorter.cpp 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502
  1. #include <opencv.hpp>
  2. #include <math.h>
  3. #include <io.h>
  4. #include "tea_sorter.h"
  5. #include "utils.h"
  6. using namespace cv;
  7. namespace graft_cv{
  8. CTeaSort::CTeaSort(
  9. ConfigParam& cp,
  10. img_type dtpye,
  11. CGcvLogger*pLog)
  12. :
  13. m_cp(cp),
  14. m_dtype(dtpye),
  15. m_pLogger(pLog),
  16. m_ppImgSaver(0),
  17. m_pImginfoRaw(0),
  18. m_pImginfoDetected(0)
  19. {
  20. m_drop_detector = RetinaDrop(m_pLogger, 0.5, 0.5);
  21. }
  22. CTeaSort::~CTeaSort()
  23. {
  24. clear_imginfo();
  25. }
  26. int CTeaSort::detect(
  27. ImgInfo*imginfo,
  28. PositionInfo& posinfo,
  29. const char* fn
  30. )
  31. {
  32. //1 model status
  33. if (!m_drop_detector.IsModelLoaded()) {
  34. m_pLogger->ERRORINFO(
  35. string("tea detect model NOT loaded"));
  36. return 1;
  37. }
  38. //2 update recognize threshold
  39. if (m_dtype == img_type::tea_grab) {
  40. m_drop_detector.SetThreshold(m_cp.object_threshold_grab, m_cp.nms_threshold_grab);
  41. }
  42. else {
  43. m_drop_detector.SetThreshold(m_cp.object_threshold_cut, m_cp.nms_threshold_cut);
  44. }
  45. //3 load data
  46. load_data(imginfo, fn);
  47. if (m_cp.image_show) {
  48. cv::destroyAllWindows();
  49. imshow_wait("input_img", m_raw_img);
  50. }
  51. //4 generate_detect_windows(vector<Rect>&boxes)
  52. vector<Rect> drop_regions;
  53. int region_cnt = generate_detect_windows(drop_regions);
  54. if (region_cnt == 0) {
  55. stringstream buff_;
  56. buff_ << m_imgId << m_dtype_str << "tea detect image regions' size == 0";
  57. m_pLogger->ERRORINFO(buff_.str());
  58. return 1;
  59. }
  60. else {
  61. stringstream bufftmp;
  62. bufftmp << m_imgId << m_dtype_str << "tea detect image regions' size = "<<region_cnt;
  63. m_pLogger->INFO(bufftmp.str());
  64. }
  65. if (m_cp.image_show) {
  66. cv::Mat rects_img = m_raw_img.clone();
  67. int step_c = int(255 / (float)region_cnt);
  68. int step_cc = step_c / 2;
  69. int step_ccc = step_cc / 2;
  70. int cnt = 0;
  71. for (auto&r : drop_regions) {
  72. cv::rectangle(rects_img, r, cv::Scalar(step_cc*cnt, step_c*cnt, step_ccc*cnt), 3);
  73. cnt += 1;
  74. }
  75. imshow_wait("regions_img", rects_img);
  76. }
  77. //5 detect
  78. vector<Bbox> droplets_raw;
  79. int dn = detect_impl(m_raw_img, drop_regions, droplets_raw);
  80. //if (dn < 2 && m_dtype == img_type::tea_grab) {
  81. // //up-down flip
  82. // cv::Mat flip_img;
  83. // cv::flip(m_raw_img, flip_img, 0);
  84. // if (m_cp.image_show) {
  85. // imshow_wait("flip_img", flip_img);
  86. // }
  87. // vector<Bbox> droplets_flip;
  88. // int dn_flip = detect_impl(flip_img, drop_regions, droplets_flip);
  89. // for (auto&b: droplets_flip) {
  90. // int y2 = flip_img.rows - b.y1;
  91. // int y1 = flip_img.rows - b.y2;
  92. // b.y1 = y1;
  93. // b.y2 = y2;
  94. //
  95. // for (int i = 0; i < 5; ++i) {
  96. // b.ppoint[2 * i + 1] = flip_img.rows - b.ppoint[2 * i + 1];
  97. // }
  98. // }
  99. // if (dn_flip > 0) {
  100. // droplets_raw.insert(
  101. // droplets_raw.end(),
  102. // droplets_flip.begin(),
  103. // droplets_flip.end());
  104. // }
  105. //}
  106. /*for (auto rect : drop_regions) {
  107. Mat roi = m_raw_img(rect);
  108. vector<Bbox> head_droplets = m_drop_detector.RunModel(roi, m_pLogger);
  109. if (m_pLogger) {
  110. stringstream buff_;
  111. buff_ << m_imgId << m_dtype_str << "-------crop_rect["<< rect.x<<","<<rect.y<<","<<rect.width
  112. <<","<<rect.height<<"],"
  113. <<" roi image detect over. tea number is " << head_droplets.size();
  114. m_pLogger->INFO(buff_.str());
  115. }
  116. for (Bbox& b : head_droplets) {
  117. b.x1 += rect.x;
  118. b.x2 += rect.x;
  119. b.y1 += rect.y;
  120. b.y2 += rect.y;
  121. for (int i = 0; i < 5; ++i) {
  122. b.ppoint[2 * i] += rect.x;
  123. b.ppoint[2 * i + 1] += rect.y;
  124. }
  125. }
  126. if (head_droplets.size()) {
  127. droplets_raw.insert(
  128. droplets_raw.end(),
  129. head_droplets.begin(),
  130. head_droplets.end());
  131. }
  132. }*/
  133. if (m_pLogger) {
  134. stringstream buff_;
  135. buff_ << m_imgId<<m_dtype_str << "image detect over. tea number is " << droplets_raw.size();
  136. m_pLogger->INFO(buff_.str());
  137. }
  138. //6 nms, width(height) filt and area calculation
  139. vector<Bbox> droplets;
  140. vector<int> keep;
  141. nms_bbox(droplets_raw, m_drop_detector.GetNmsThreshold(), keep);
  142. //width(height) filter
  143. for (int i : keep) {
  144. droplets.push_back(droplets_raw[i]);
  145. }
  146. if (m_pLogger) {
  147. stringstream buff_;
  148. buff_ << m_imgId << m_dtype_str << "after nms, keep tea number is " << droplets.size();
  149. m_pLogger->INFO(buff_.str());
  150. }
  151. int valid_cnt = 0;
  152. for (int i = 0; i < droplets.size();++i) {
  153. if (i > 1) { break; }
  154. Bbox&b = droplets.at(i);
  155. double grab_x, grab_y;
  156. double angle = calalate_angle(b, grab_x, grab_y);
  157. valid_cnt += 1;
  158. //grab point
  159. if (i == 0) {
  160. if (m_dtype == img_type::tea_grab) {
  161. posinfo.tea_grab_x1 = grab_x;
  162. posinfo.tea_grab_y1 = grab_y;
  163. posinfo.tea_grab_angle1 = angle;
  164. }
  165. else {
  166. // 切割点是3、4的中间的点
  167. posinfo.tea_cut_x1 = 0.5 * (b.ppoint[4] + b.ppoint[6]);
  168. posinfo.tea_cut_y1 = 0.5 * (b.ppoint[5] + b.ppoint[7]);
  169. posinfo.tea_cut_angle1 = angle;
  170. }
  171. }
  172. else {
  173. if (m_dtype == img_type::tea_grab) {
  174. posinfo.tea_grab_x2 = grab_x;
  175. posinfo.tea_grab_y2 = grab_y;
  176. posinfo.tea_grab_angle2 = angle;
  177. }
  178. else {
  179. // 切割点是3、4的中间的点
  180. posinfo.tea_cut_x2 = 0.5 * (b.ppoint[4] + b.ppoint[6]);
  181. posinfo.tea_cut_y2 = 0.5 * (b.ppoint[5] + b.ppoint[7]);
  182. posinfo.tea_cut_angle2 = angle;
  183. }
  184. }
  185. }
  186. //6 draw
  187. if (m_cp.image_return) {
  188. this->clear_imginfo();
  189. cv::Mat img_rst = m_raw_img.clone();
  190. int cnt = 0;
  191. for (auto& b : droplets) {
  192. //rectangle
  193. cv::Rect r = cv::Rect(cv::Point2i(b.x1, b.y1), cv::Point2i(b.x2, b.y2));
  194. if (cnt < 2) {
  195. cv::rectangle(img_rst, r, cv::Scalar(0, 0, 255),2);
  196. }
  197. else {
  198. cv::rectangle(img_rst, r, cv::Scalar(0, 255, 0),2);
  199. }
  200. //score
  201. char name[256];
  202. cv::Scalar color(120, 120, 0);//bgr
  203. sprintf_s(name, "%.2f", b.score);
  204. cv::putText(img_rst, name,
  205. cv::Point(b.x1, b.y1),
  206. cv::FONT_HERSHEY_COMPLEX, 0.7, color, 2);
  207. //points
  208. cv::circle(img_rst, cv::Point(int(b.ppoint[0]), int(b.ppoint[1])), 4, cv::Scalar(255, 0, 255), -1, 3, 0);
  209. cv::circle(img_rst, cv::Point(int(b.ppoint[2]), int(b.ppoint[3])), 4, cv::Scalar(0, 255, 255), -1, 3, 0);
  210. cv::circle(img_rst, cv::Point(int(b.ppoint[4]), int(b.ppoint[5])), 4, cv::Scalar(255, 0, 0), -1, 3, 0);
  211. cv::circle(img_rst, cv::Point(int(b.ppoint[6]), int(b.ppoint[7])), 4, cv::Scalar(0, 255, 0), -1, 3, 0);
  212. cv::circle(img_rst, cv::Point(int(b.ppoint[8]), int(b.ppoint[9])), 4, cv::Scalar(0, 0, 255), -1, 3, 0);
  213. //grab points
  214. if (m_dtype == img_type::tea_grab) {
  215. double grab_x, grab_y;
  216. calalate_angle(b, grab_x, grab_y);
  217. //cv::circle(img_rst, cv::Point(int(grab_x), int(grab_y)), 4, cv::Scalar(0, 215, 255), -1, 3, 0);
  218. //lines, p4-p5, p5-grab
  219. cv::line(img_rst,
  220. cv::Point(int(b.ppoint[6]), int(b.ppoint[7])),
  221. cv::Point(int(b.ppoint[8]), int(b.ppoint[9])),
  222. cv::Scalar(0, 215, 255), 2);
  223. cv::line(img_rst,
  224. cv::Point(int(b.ppoint[8]), int(b.ppoint[9])),
  225. cv::Point(int(grab_x), int(grab_y)),
  226. cv::Scalar(0, 215, 255), 2);
  227. //line x
  228. int radius = 20;
  229. int cx = int(grab_x);
  230. int cy = int(grab_y);
  231. cv::line(img_rst, cv::Point(cx - radius, cy - radius), cv::Point(cx + radius, cy + radius), cv::Scalar(0, 215, 255), 2);
  232. cv::line(img_rst, cv::Point(cx - radius, cy + radius), cv::Point(cx + radius, cy - radius), cv::Scalar(0, 215, 255), 2);
  233. }
  234. //cut points
  235. if (m_dtype == img_type::tea_cut) {
  236. //lines, p3-p4
  237. cv::line(img_rst,
  238. cv::Point(int(b.ppoint[4]), int(b.ppoint[5])),
  239. cv::Point(int(b.ppoint[6]), int(b.ppoint[7])),
  240. cv::Scalar(0, 215, 255), 2);
  241. //line x
  242. int cx = int(0.5 * (b.ppoint[4] + b.ppoint[6]));
  243. int cy = int(0.5 * (b.ppoint[5] + b.ppoint[7]));
  244. int radius = 20;
  245. cv::line(img_rst, cv::Point(cx - radius, cy - radius), cv::Point(cx + radius, cy + radius), cv::Scalar(0, 215, 255),2);
  246. cv::line(img_rst, cv::Point(cx - radius, cy + radius), cv::Point(cx + radius, cy - radius), cv::Scalar(0, 215, 255),2);
  247. }
  248. cnt += 1;
  249. }
  250. if (m_cp.image_show) {
  251. imshow_wait("result_img", img_rst);
  252. }
  253. m_pImginfoRaw = mat2imginfo(m_raw_img);
  254. m_pImginfoDetected = mat2imginfo(img_rst);
  255. posinfo.pp_images[0] = m_pImginfoRaw;
  256. posinfo.pp_images[1] = m_pImginfoDetected;
  257. if (m_ppImgSaver && *m_ppImgSaver) {
  258. (*m_ppImgSaver)->saveImage(img_rst, m_imgId + "_rst_0");
  259. }
  260. }
  261. //拍照无苗, 返回识别结果-1
  262. if (valid_cnt == 0) { return -1; }
  263. return 0;
  264. }
  265. int CTeaSort::detect_impl(
  266. cv::Mat& raw_img, //input, image
  267. vector<Rect>&drop_regions, //input, detect regions
  268. vector<Bbox> &droplets_raw //output, detect result
  269. )
  270. {
  271. //return number of detect result
  272. droplets_raw.clear();
  273. for (auto rect : drop_regions) {
  274. Mat roi = raw_img(rect);
  275. vector<Bbox> head_droplets = m_drop_detector.RunModel(roi, m_pLogger);
  276. if (m_pLogger) {
  277. stringstream buff_;
  278. buff_ << m_imgId << m_dtype_str << "-------crop_rect[" << rect.x << "," << rect.y << "," << rect.width
  279. << "," << rect.height << "],"
  280. << " roi image detect over. tea number is " << head_droplets.size();
  281. m_pLogger->INFO(buff_.str());
  282. }
  283. for (Bbox& b : head_droplets) {
  284. b.x1 += rect.x;
  285. b.x2 += rect.x;
  286. b.y1 += rect.y;
  287. b.y2 += rect.y;
  288. for (int i = 0; i < 5; ++i) {
  289. b.ppoint[2 * i] += rect.x;
  290. b.ppoint[2 * i + 1] += rect.y;
  291. }
  292. }
  293. if (head_droplets.size()) {
  294. droplets_raw.insert(
  295. droplets_raw.end(),
  296. head_droplets.begin(),
  297. head_droplets.end());
  298. }
  299. }
  300. return droplets_raw.size();
  301. }
  302. double CTeaSort::calalate_angle(
  303. Bbox&b, //input
  304. double& grab_x, //output
  305. double& grab_y //output
  306. )
  307. {
  308. grab_x = grab_y = 0.0;
  309. double angle = 0.0;
  310. float x3,y3,x4,y4,x5,y5;
  311. x3 = b.ppoint[4];
  312. y3 = b.ppoint[5];
  313. x4 = b.ppoint[6];
  314. y4 = b.ppoint[7];
  315. x5 = b.ppoint[8];
  316. y5 = b.ppoint[9];
  317. if (m_dtype == img_type::tea_grab) {
  318. //calculate line of p4 ans p5
  319. double r45 = sqrt((x4 - x5)*(x4 - x5) + (y4 - y5)*(y4 - y5));
  320. if (r45 < 15.0) {
  321. angle = atan2(x5 - x3, y5 - y3);
  322. }
  323. else {
  324. angle = atan2(x5 - x4, y5 - y4);
  325. }
  326. //计算抓取点
  327. double pr = (double)m_cp.offset_grab;
  328. double dx = pr * sin(angle);
  329. double dy = pr * cos(angle);
  330. grab_x = x5 + dx;
  331. grab_y = y5 + dy;
  332. }
  333. else {
  334. //tea cut, calculate line of p3 ans p4
  335. angle = atan2(x3 - x4, y3 - y4);
  336. }
  337. angle *= (180.0 / 3.1415926);
  338. return angle;
  339. }
  340. int CTeaSort::load_data(
  341. ImgInfo*imginfo,
  342. const char* fn/* = 0*/)
  343. {
  344. //数据加载功能实现,并生成imageid,保存原始数据到文件
  345. int rst = 0;
  346. //generate image id
  347. if (m_dtype == img_type::tea_grab) {
  348. m_imgId = getImgId(img_type::tea_grab);
  349. m_dtype_str = string(" tea_grab ");
  350. }
  351. else {
  352. m_imgId = getImgId(img_type::tea_cut);
  353. m_dtype_str = string(" tea_cut ");
  354. }
  355. if (imginfo) {
  356. if (m_pLogger) {
  357. stringstream buff;
  358. buff << m_imgId << m_dtype_str << "image, width=" << imginfo->width
  359. << "\theight=" << imginfo->height;
  360. m_pLogger->INFO(buff.str());
  361. }
  362. if (!isvalid(imginfo) || (imginfo->channel!=1 && imginfo->channel!=3)) {
  363. if (m_pLogger) {
  364. m_pLogger->ERRORINFO(m_imgId + m_dtype_str + "input image invalid.");
  365. }
  366. throw_msg(m_imgId + " invalid input image");
  367. }
  368. if (imginfo->channel == 1) {
  369. cv::Mat tmp_img = imginfo2mat(imginfo);
  370. vector<Mat> channels;
  371. for (size_t i = 0; i < 3; ++i) { channels.push_back(tmp_img); }
  372. cv::merge(channels, m_raw_img);
  373. }
  374. else {
  375. m_raw_img = imginfo2mat(imginfo);
  376. }
  377. }
  378. else {
  379. cv::Mat img = imread(fn, cv::IMREAD_COLOR);
  380. if (img.empty()) {
  381. if (m_pLogger) {
  382. m_pLogger->ERRORINFO(m_imgId + m_dtype_str + "input image invalid:" + string(fn));
  383. }
  384. throw_msg(m_imgId + m_dtype_str + "invalid input image: " + string(fn));
  385. }
  386. if (m_pLogger) {
  387. stringstream buff;
  388. buff << m_imgId << m_dtype_str << "image, width=" << img.cols
  389. << "\theight=" << img.rows;
  390. m_pLogger->INFO(buff.str());
  391. }
  392. m_raw_img = img.clone();
  393. }
  394. //image saver
  395. if (m_ppImgSaver && *m_ppImgSaver) {
  396. (*m_ppImgSaver)->saveImage(m_raw_img, m_imgId);
  397. }
  398. return rst;
  399. }
  400. int CTeaSort::load_model()
  401. {
  402. bool b = false;
  403. if (!m_drop_detector.IsModelLoaded()) {
  404. if (m_dtype == img_type::tea_grab) {
  405. b = m_drop_detector.LoadModel(m_cp.model_path_grab);
  406. }
  407. else {
  408. b = m_drop_detector.LoadModel(m_cp.model_path_cut);
  409. }
  410. }
  411. else {
  412. b = true;
  413. }
  414. return b ? 0 : 1;
  415. }
  416. void CTeaSort::clear_imginfo() {
  417. if (m_pImginfoDetected) {
  418. imginfo_release(&m_pImginfoDetected);
  419. m_pImginfoDetected = 0;
  420. }
  421. if (m_pImginfoRaw) {
  422. imginfo_release(&m_pImginfoRaw);
  423. m_pImginfoRaw = 0;
  424. }
  425. }
  426. int CTeaSort::generate_detect_windows(vector<Rect>&boxes)
  427. {
  428. boxes.clear();
  429. int grid_row = m_cp.grid_row_cut;
  430. int grid_col = m_cp.grid_col_cut;
  431. int grid_padding = m_cp.grid_padding_cut;
  432. if (m_dtype == img_type::tea_grab) {
  433. grid_row = m_cp.grid_row_grab;
  434. grid_col = m_cp.grid_col_grab;
  435. grid_padding = m_cp.grid_padding_grab;
  436. }
  437. if (grid_row < 1) { grid_row = 1; }
  438. if (grid_col < 1) { grid_col = 1; }
  439. if (grid_padding < 0) { grid_padding = 0; }
  440. int block_height = int(m_raw_img.rows / (float)grid_row + 0.5);
  441. int block_width = int(m_raw_img.cols / (float)grid_col + 0.5);
  442. for (int r = 0; r < grid_row; ++r) {
  443. for (int c = 0; c < grid_col; ++c) {
  444. int x0 = c*block_width - grid_padding;
  445. int y0 = r*block_height - grid_padding;
  446. int x1 = (c+1)*block_width + grid_padding;
  447. int y1 = (r+1)*block_height + grid_padding;
  448. if (x0 < 0) { x0 = 0; }
  449. if (y0 < 0) { y0 = 0; }
  450. if (x1 > m_raw_img.cols) { x1 = m_raw_img.cols; }
  451. if (y1 > m_raw_img.rows) { y1 = m_raw_img.rows; }
  452. Rect r(x0, y0, x1-x0, y1-y0);
  453. boxes.push_back(r);
  454. }
  455. }
  456. return boxes.size();
  457. }
  458. }