tea_cv/tea_detect.h

#pragma once

#include <opencv.hpp>
#include "logger.h"
#include "data_def.h"
#include "inference.h"


namespace graft_cv {

	class RetinaDrop {
		struct DropBox {
			float x1;
			float y1;
			float x2;
			float y2;
		};

		struct DropRes {
			float confidence;
			DropBox drop_box;
			std::vector<cv::Point2f> keypoints;			
		};

	public:
		explicit RetinaDrop(CGcvLogger* pLogger=0, float obj_th=0.6, float nms_th=0.4);
		~RetinaDrop();
		bool LoadModel(std::string onnx_path); 
		std::vector<Bbox> RunModel(cv::Mat& img, CGcvLogger* pInstanceLogger=0);
		bool IsModelLoaded();
		float GetNmsThreshold();
		void SetThreshold(float object_threshold, float nms_threshold);

	private:
		void generate_anchors();		
		int post_process(
			cv::Mat &vec_Mat, 
			std::vector<cv::Mat> &result_matrix, 
			std::vector<RetinaDrop::DropRes>& valid_result);

		void nms_detect(
			std::vector<DropRes>& detections,
			std::vector<int>& keep);

		static float iou_calculate(
			const DropBox& det_a, 
			const DropBox& det_b);

		int BATCH_SIZE; //default 1
		int INPUT_CHANNEL; //default 3
		int IMAGE_WIDTH; //default 640
		int IMAGE_HEIGHT; //default 640
		float m_obj_threshold; // default 0.5
		float m_nms_threshold; // default 0.45		
		

		cv::Mat m_refer_matrix;
		int m_anchor_num;
		int m_bbox_head;
		
		std::vector<int> m_feature_sizes;
		std::vector<int> m_feature_steps; 
		std::vector<int> m_feature_maps;
		std::vector<std::vector<int>>m_anchor_sizes; 
		int m_sum_of_feature;

		cv::dnn::Net m_model;
		float m_variance[2];
		cv::Scalar m_img_mean;
		cv::Size m_size_detection;
		bool m_model_loaded;
		CGcvLogger* m_pLogger;
	};


	class YoloDrop {
		struct DropBox {
			float x1;
			float y1;
			float x2;
			float y2;
		};

		struct DropRes {
			float confidence;
			DropBox drop_box;
			std::vector<cv::Point2f> keypoints;
		};

	public:
		explicit YoloDrop(CGcvLogger* pLogger = 0, float obj_th = 0.6, float nms_th = 0.4);
		~YoloDrop();
		bool LoadModel(std::string onnx_path);
		std::vector<Bbox> RunModel(cv::Mat& img, CGcvLogger* pInstanceLogger = 0);
		bool IsModelLoaded();
		float GetNmsThreshold();
		void SetThreshold(float object_threshold, float nms_threshold);

	private:
		Inference m_infer;
		void generate_anchors();
		int post_process(
			cv::Mat &vec_Mat,
			std::vector<cv::Mat> &result_matrix,
			std::vector<YoloDrop::DropRes>& valid_result);

		void nms_detect(
			std::vector<DropRes>& detections,
			std::vector<int>& keep);

		static float iou_calculate(
			const DropBox& det_a,
			const DropBox& det_b);

		int BATCH_SIZE; //default 1
		int INPUT_CHANNEL; //default 3
		int IMAGE_WIDTH; //default 640
		int IMAGE_HEIGHT; //default 640
		float m_obj_threshold; // default 0.5
		float m_nms_threshold; // default 0.45		


		cv::Mat m_refer_matrix;
		int m_anchor_num;
		int m_bbox_head;

		std::vector<int> m_feature_sizes;
		std::vector<int> m_feature_steps;
		std::vector<int> m_feature_maps;
		std::vector<std::vector<int>>m_anchor_sizes;
		int m_sum_of_feature;

		cv::dnn::Net m_model;
		float m_variance[2];
		cv::Scalar m_img_mean;
		cv::Size m_size_detection;
		bool m_model_loaded;
		CGcvLogger* m_pLogger;
	};
}