whellvalue/3part/Cyclops/include/TemplateDetector.h

#ifndef TemplateDetector_h__
#define TemplateDetector_h__

#include <string>
#include <set>
#include <opencv2/opencv.hpp>
#include <opencv2/opencv_modules.hpp>
#include "CVUtils.h"
#include <map>

#define TD_STR_CHANNEL_NAME_ORIGINAL    "Original"
#define TD_STR_ORI_IMG_FILE_NAME        "OriginImageFile"
#define TD_STR_ROI_COUNT                "ROICounts"
#define TD_STR_ROI_SCALE                "ROI_scale"
#define TD_STR_ROI_CHANNEL              "ROI_channel"
#define TD_STR_ROI_LOW_THRE             "ROI_low_thre"
#define TD_STR_ROI_HIGH_THRE            "ROI_high_thre"
#define TD_STR_ROI_AREA_THRE            "ROI_area_thre"
#define TD_STR_ROI_RECT                 "ROI_rect"
#define TD_STR_ROI_DET_RECT             "ROI_det_rect"
#define TD_STR_ROI_SUBIMG               "ROI_subImg"
#define TD_STR_ROI_METHOD               "ROI_method"
#define TD_STR_JOIN                     "_"
#define TD_STR_ROI_ID                   "ROI_id"
#define TD_STR_ROI_SMOOTH_KSIZE         "ROI_smooth_ksize"

#define TD_STR_TMPL_FILE_DOT_EXT        ".yaml"
#define TD_STR_TMPL_FILE_EXT            "yaml"
#define TD_STR_GLOBAL_PARAM_FILE_NAME   "global"

#define TD_INT_METHOD_COUNT             10
#define TD_INT_METHOD_2DFILTER          6
#define TD_INT_METHOD_2DFILTER_ABS      7
#define TD_INT_METHOD_HSV_DIS           8
#define TD_INT_METHOD_ORI               9
#define TD_STR_METHOD_2DFILTER          "2D_FILTER"
#define TD_STR_METHOD_2DFILTER_ABS      "2D_FILTER_ABS"
#define TD_STR_METHOD_HSV_DIS           "HSV_DIS"
#define TD_STR_METHOD_ORI               "ORI"

#define TD_IMG_SCALE_METHOD             INTER_AREA

using std::string;
using std::map;
using namespace cv;
using std::set;

class TemplateDescriptor
{
public:    
    TemplateDescriptor();
    TemplateDescriptor(const Mat& srcImg, float scale, char channelName, 
        const Rect& roiRect, int method, float thre, const Mat& roiImg = Mat());

    static void resizeImage(const Mat& srcImg, Mat& dstImg, double scale, int method);
    static void smoothImage(const Mat& srcImg, Mat& dstImg, int ksize);
    static string methodStr(int i);

    void extractImageChannel(const Mat& srcImg, Mat& dstImg) const;
    void convertImage(const Mat& srcImg, Mat& dstImg) const;
    // assume srcImg is bgr image
    void init(const Mat& srcImg, float scale, char channelName,
        const Rect& roiRect, int method, float thre);
    void init();

    double scale() const { return mScale; }
    void setScale(double val) { mScale = val; }
    char channelName() const { return mChannelName; }
    void setChannelName(char val) { mChannelName = val; }
    cv::Mat srcImg() const { return mSrcImg; }
    void setSrcImg(cv::Mat val) { mSrcImg = val; }
    cv::Mat roiImg() const { return mRoiImg; }
    void setRoiImg(cv::Mat val);
    double lowThre() const { return mLowThre; }
    double highThre() const { return mHighThre; }
    void setLowThre(double val);
    void setHighThre(double val);
    double areaThre() const { return mAreaThre; }
    void setAreaThre(double val) { mAreaThre = val; }
    int method() const { return mMethod; }
    void setMethod(int val);
    Rectf tmplRoiRect() const { return mTmplRoiRect; }
    void setTmplRoiRect(Rectf val) { mTmplRoiRect = val; }
    Rectf detRoiRect() const { return mDetRoiRect; }
    void setDetRoiRect(Rectf val) { mDetRoiRect = val; }
    string imgFileName() const { return mImgFileName; }
    void setImgFileName(string val) { mImgFileName = val; }
    string id() const { return mID; }
    void setID(string val) { mID = val; }
    double normScale() const { return mNormScale; }
    void setNormScale(double scale);
    int smoothKSize() const { return mSmoothKSize; }
    void setSmoothKSize(int val) { mSmoothKSize = val; }

    void operator= (const TemplateDescriptor& td);

protected:
    double mScale, mLowThre, mHighThre, mNormScale, mAreaThre;
    char mChannelName;
    int mMethod;
    Rectf mTmplRoiRect, mDetRoiRect;
    Mat mSrcImg;
    Mat mRoiImg;
    string mImgFileName;
    string mID;
    int mSmoothKSize;
};

typedef map<char, map<float, map<int, Mat> > > ChannelScaleSmoothMatMap;

class TemplateDetector
{
public:
    TemplateDetector() {};
    TemplateDetector(const string& dirPath);

    void clear();
    void clear(string imgFileName);
    void del(const set<int>& idxVec);
    void add(const TemplateDescriptor& td);
    void replace(string imgFileName, const vector<TemplateDescriptor>& vec);
    void insertAndModify(string imgFilename, const vector<TemplateDescriptor>& vec);
    void setThre(const TemplateDescriptor* pTD, double lowThre, double highThre);
    void setAll(const TemplateDescriptor* pTD, const TemplateDescriptor* pNewTD);

    void loadGlobalParams(const string& filePath);
    void loadByPathPattern(const string& dirPath);
    string saveTmplToFileStorage(string dirPath, string imgFileName);
    bool saveToDir(const string& dirPath);
    bool saveOneFile(const string& imgFilePath);

    int detect(const Mat& img, vector<Mat>& results, vector<int>& idxVec) const;
    const TemplateDescriptor* templateDescriptor(int idx) const;
    void find(string fileName, vector<const TemplateDescriptor*>& tdVec);
    void find(string fileName, vector<TemplateDescriptor*>& tdVec);

    cv::Rect globalRoiRect() const { return mGlobalRoiRect; }
    void setGlobalRoiRect(cv::Rect val) { mGlobalRoiRect = val; }
    
    static char toChannelChar(const string& str);
    static string toChannelStr(char c);
    static int channelCharIdx(char c);
    static string toTmplPathPattern(const string& dirPath);

protected:
    int initScaledImgMap(const Mat& img, ChannelScaleSmoothMatMap& imgMap) const;

    vector<TemplateDescriptor> mTmplDesVec;
    Rect mGlobalRoiRect;

private:
};

#endif // TemplateDetector_h__