tea_sorter.cpp 8.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343
  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. //4 generate_detect_windows(vector<Rect>&boxes)
  48. vector<Rect> drop_regions;
  49. int region_cnt = generate_detect_windows(drop_regions);
  50. if (region_cnt == 0) {
  51. stringstream buff_;
  52. buff_ << m_imgId << m_dtype_str << "tea detect image regions' size == 0";
  53. m_pLogger->ERRORINFO(buff_.str());
  54. return 1;
  55. }
  56. else {
  57. stringstream bufftmp;
  58. bufftmp << m_imgId << m_dtype_str << "tea detect image regions' size = "<<region_cnt;
  59. m_pLogger->INFO(bufftmp.str());
  60. }
  61. //5 detect
  62. vector<Bbox> droplets_raw;
  63. for (auto rect : drop_regions) {
  64. Mat roi = m_raw_img(rect);
  65. vector<Bbox> head_droplets = m_drop_detector.RunModel(roi, m_pLogger);
  66. if (m_pLogger) {
  67. stringstream buff_;
  68. buff_ << m_imgId << m_dtype_str << "-------crop_rect["<< rect.x<<","<<rect.y<<","<<rect.width
  69. <<","<<rect.height<<"],"
  70. <<" roi image detect over. tea number is " << head_droplets.size();
  71. m_pLogger->INFO(buff_.str());
  72. }
  73. for (Bbox& b : head_droplets) {
  74. b.x1 += rect.x;
  75. b.x2 += rect.x;
  76. b.y1 += rect.y;
  77. b.y2 += rect.y;
  78. for (int i = 0; i < 5; ++i) {
  79. b.ppoint[2 * i] += rect.x;
  80. b.ppoint[2 * i + 1] += rect.y;
  81. }
  82. }
  83. if (head_droplets.size()) {
  84. droplets_raw.insert(
  85. droplets_raw.end(),
  86. head_droplets.begin(),
  87. head_droplets.end());
  88. }
  89. }
  90. if (m_pLogger) {
  91. stringstream buff_;
  92. buff_ << m_imgId<<m_dtype_str << "image detect over. tea number is " << droplets_raw.size();
  93. m_pLogger->INFO(buff_.str());
  94. }
  95. //6 nms, width(height) filt and area calculation
  96. vector<Bbox> droplets;
  97. vector<int> keep;
  98. nms_bbox(droplets_raw, m_drop_detector.GetNmsThreshold(), keep);
  99. //width(height) filter
  100. for (int i : keep) {
  101. droplets.push_back(droplets_raw[i]);
  102. }
  103. if (m_pLogger) {
  104. stringstream buff_;
  105. buff_ << m_imgId << m_dtype_str << "after nms, keep tea number is " << droplets.size();
  106. m_pLogger->INFO(buff_.str());
  107. }
  108. int valid_cnt = 0;
  109. for (int i = 0; i < droplets.size();++i) {
  110. if (i > 1) { break; }
  111. Bbox&b = droplets.at(i);
  112. double angle = calalate_angle(b);
  113. valid_cnt += 1;
  114. //grab point
  115. if (i == 0) {
  116. if (m_dtype == img_type::tea_grab) {
  117. posinfo.tea_grab_x1 = b.ppoint[8];
  118. posinfo.tea_grab_y1 = b.ppoint[9];
  119. posinfo.tea_grab_angle1 = angle;
  120. }
  121. else {
  122. posinfo.tea_cut_x1 = b.ppoint[6];
  123. posinfo.tea_cut_y1 = b.ppoint[7];
  124. posinfo.tea_cut_angle1 = angle;
  125. }
  126. }
  127. else {
  128. if (m_dtype == img_type::tea_grab) {
  129. posinfo.tea_grab_x2 = b.ppoint[8];
  130. posinfo.tea_grab_y2 = b.ppoint[9];
  131. posinfo.tea_grab_angle2 = angle;
  132. }
  133. else {
  134. posinfo.tea_cut_x2 = b.ppoint[6];
  135. posinfo.tea_cut_y2 = b.ppoint[7];
  136. posinfo.tea_cut_angle2 = angle;
  137. }
  138. }
  139. }
  140. //6 draw
  141. if (m_cp.image_return) {
  142. this->clear_imginfo();
  143. cv::Mat img_rst = m_raw_img.clone();
  144. int cnt = 0;
  145. for (auto& b : droplets) {
  146. char name[256];
  147. cv::Scalar color(20, 0, 0);//bgr
  148. sprintf_s(name, "%.2f", b.score);
  149. cv::putText(img_rst, name,
  150. cv::Point(b.x1, b.y1),
  151. cv::FONT_HERSHEY_COMPLEX, 0.7, color, 2);
  152. cv::Rect r = cv::Rect(cv::Point2i(b.x1, b.y1), cv::Point2i(b.x2, b.y2));
  153. if (cnt < 2) {
  154. cv::rectangle(img_rst, r, cv::Scalar(0, 0, 255));
  155. }
  156. else {
  157. cv::rectangle(img_rst, r, cv::Scalar(0, 255, 0));
  158. }
  159. cv::rectangle(img_rst, r, cv::Scalar(0, 0, 255));
  160. cv::circle(img_rst, cv::Point(int(b.ppoint[0]), int(b.ppoint[1])), 4, cv::Scalar(0, 0, 255), -1, 8, 0);
  161. cv::circle(img_rst, cv::Point(int(b.ppoint[2]), int(b.ppoint[3])), 4, cv::Scalar(0, 255, 255), -1, 8, 0);
  162. cv::circle(img_rst, cv::Point(int(b.ppoint[4]), int(b.ppoint[5])), 4, cv::Scalar(255, 0, 255), -1, 8, 0);
  163. cv::circle(img_rst, cv::Point(int(b.ppoint[6]), int(b.ppoint[7])), 4, cv::Scalar(0, 255, 0), -1, 8, 0);
  164. cv::circle(img_rst, cv::Point(int(b.ppoint[8]), int(b.ppoint[9])), 4, cv::Scalar(255, 0, 0), -1, 8, 0);
  165. cnt += 1;
  166. }
  167. m_pImginfoRaw = mat2imginfo(m_raw_img);
  168. m_pImginfoDetected = mat2imginfo(img_rst);
  169. posinfo.pp_images[0] = m_pImginfoRaw;
  170. posinfo.pp_images[1] = m_pImginfoDetected;
  171. if (m_ppImgSaver && *m_ppImgSaver) {
  172. (*m_ppImgSaver)->saveImage(img_rst, m_imgId + "_rst_0");
  173. }
  174. }
  175. //拍照无苗, 返回识别结果-1
  176. if (valid_cnt == 0) { return -1; }
  177. return 0;
  178. }
  179. double CTeaSort::calalate_angle(Bbox&b) {
  180. double angle = 0.0;
  181. float x3,y3,x4,y4,x5,y5;
  182. x3 = b.ppoint[4];
  183. y3 = b.ppoint[5];
  184. x4 = b.ppoint[6];
  185. y4 = b.ppoint[7];
  186. x5 = b.ppoint[8];
  187. y5 = b.ppoint[9];
  188. double r45 = sqrt((x4 - x5)*(x4 - x5) + (y4 - y5)*(y4 - y5));
  189. if (r45 < 15.0) {
  190. angle = atan2(x5 - x3, y5 - y3);
  191. }
  192. else {
  193. angle = atan2(x5 - x4, y5 - y4);
  194. }
  195. angle *= (180.0 / 3.1415926);
  196. return angle;
  197. }
  198. int CTeaSort::load_data(
  199. ImgInfo*imginfo,
  200. const char* fn/* = 0*/)
  201. {
  202. //数据加载功能实现,并生成imageid,保存原始数据到文件
  203. int rst = 0;
  204. //generate image id
  205. if (m_dtype == img_type::tea_grab) {
  206. m_imgId = getImgId(img_type::tea_grab);
  207. m_dtype_str = string(" tea_grab ");
  208. }
  209. else {
  210. m_imgId = getImgId(img_type::tea_cut);
  211. m_dtype_str = string(" tea_cut ");
  212. }
  213. if (imginfo) {
  214. if (m_pLogger) {
  215. stringstream buff;
  216. buff << m_imgId << m_dtype_str << "image, width=" << imginfo->width
  217. << "\theight=" << imginfo->height;
  218. m_pLogger->INFO(buff.str());
  219. }
  220. if (!isvalid(imginfo)) {
  221. if (m_pLogger) {
  222. m_pLogger->ERRORINFO(m_imgId + m_dtype_str + "input image invalid.");
  223. }
  224. throw_msg(m_imgId + " invalid input image");
  225. }
  226. m_raw_img = imginfo2mat(imginfo);
  227. }
  228. else {
  229. cv::Mat img = imread(fn, cv::IMREAD_COLOR);
  230. if (img.empty()) {
  231. if (m_pLogger) {
  232. m_pLogger->ERRORINFO(m_imgId + m_dtype_str + "input image invalid:" + string(fn));
  233. }
  234. throw_msg(m_imgId + m_dtype_str + "invalid input image: " + string(fn));
  235. }
  236. if (m_pLogger) {
  237. stringstream buff;
  238. buff << m_imgId << m_dtype_str << "image, width=" << img.cols
  239. << "\theight=" << img.rows;
  240. m_pLogger->INFO(buff.str());
  241. }
  242. m_raw_img = img.clone();
  243. }
  244. //image saver
  245. if (m_ppImgSaver && *m_ppImgSaver) {
  246. (*m_ppImgSaver)->saveImage(m_raw_img, m_imgId);
  247. }
  248. return rst;
  249. }
  250. int CTeaSort::load_model()
  251. {
  252. bool b = false;
  253. if (!m_drop_detector.IsModelLoaded()) {
  254. if (m_dtype == img_type::tea_grab) {
  255. b = m_drop_detector.LoadModel(m_cp.model_path_grab);
  256. }
  257. else {
  258. b = m_drop_detector.LoadModel(m_cp.model_path_cut);
  259. }
  260. }
  261. else {
  262. b = true;
  263. }
  264. return b ? 0 : 1;
  265. }
  266. void CTeaSort::clear_imginfo() {
  267. if (m_pImginfoDetected) {
  268. imginfo_release(&m_pImginfoDetected);
  269. m_pImginfoDetected = 0;
  270. }
  271. if (m_pImginfoRaw) {
  272. imginfo_release(&m_pImginfoRaw);
  273. m_pImginfoRaw = 0;
  274. }
  275. }
  276. int CTeaSort::generate_detect_windows(vector<Rect>&boxes)
  277. {
  278. boxes.clear();
  279. int grid_row = m_cp.grid_row_cut;
  280. int grid_col = m_cp.grid_col_cut;
  281. int grid_padding = m_cp.grid_padding_cut;
  282. if (m_dtype == img_type::tea_grab) {
  283. grid_row = m_cp.grid_row_grab;
  284. grid_col = m_cp.grid_col_grab;
  285. grid_padding = m_cp.grid_padding_grab;
  286. }
  287. if (grid_row < 1) { grid_row = 1; }
  288. if (grid_col < 1) { grid_col = 1; }
  289. if (grid_padding < 0) { grid_padding = 0; }
  290. int block_height = int(m_raw_img.rows / (float)grid_row + 0.5);
  291. int block_width = int(m_raw_img.cols / (float)grid_col + 0.5);
  292. for (int r = 0; r < grid_row; ++r) {
  293. for (int c = 0; c < grid_col; ++c) {
  294. int x0 = c*block_width - grid_padding;
  295. int y0 = r*block_height - grid_padding;
  296. int x1 = (c+1)*block_width + grid_padding;
  297. int y1 = (r+1)*block_height + grid_padding;
  298. if (x0 < 0) { x0 = 0; }
  299. if (y0 < 0) { y0 = 0; }
  300. if (x1 > m_raw_img.cols) { x1 = m_raw_img.cols; }
  301. if (y1 > m_raw_img.rows) { y1 = m_raw_img.rows; }
  302. Rect r(x0, y0, x1-x0, y1-y0);
  303. boxes.push_back(r);
  304. }
  305. }
  306. return boxes.size();
  307. }
  308. }