tea_cv/tea_sorter.cpp

#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();
}

}