├── readme.md
├── README.md
├── .gitignore
├── CMakeLists.txt
├── base
    ├── common.hpp
    ├── classification.hpp
    ├── classification.cpp
    └── common.cpp
└── main.cpp


/readme.md:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | 说明：本代码用于计算机视觉任务中，分类模型数据的预打标签和数据清洗
2 | 
3 | 具体原理详见：https://zhuanlan.zhihu.com/p/359960152
4 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
 1 | # Prerequisites
 2 | *.d
 3 | 
 4 | # Compiled Object files
 5 | *.slo
 6 | *.lo
 7 | *.o
 8 | *.obj
 9 | 
10 | # Precompiled Headers
11 | *.gch
12 | *.pch
13 | 
14 | # Compiled Dynamic libraries
15 | *.so
16 | *.dylib
17 | *.dll
18 | 
19 | # Fortran module files
20 | *.mod
21 | *.smod
22 | 
23 | # Compiled Static libraries
24 | *.lai
25 | *.la
26 | *.a
27 | *.lib
28 | 
29 | # Executables
30 | *.exe
31 | *.out
32 | *.app
33 | 


--------------------------------------------------------------------------------
/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | cmake_minimum_required(VERSION 3.15)
 2 | project(main)
 3 | 
 4 | set(CMAKE_CXX_STANDARD 11)
 5 | 
 6 | # 指定libTorch位置
 7 | include_directories(${CMAKE_SOURCE_DIR}/3rdParty/libtorch/include)
 8 | link_directories(${CMAKE_SOURCE_DIR}/3rdParty/libtorch/lib)
 9 | 
10 | # 指定opencv位置
11 | include_directories(${CMAKE_SOURCE_DIR}/3rdParty/opencv/include)
12 | link_directories(${CMAKE_SOURCE_DIR}/3rdParty/opencv/lib)
13 | 
14 | # 指定cuda位置
15 | include_directories(/usr/local/cuda-10.1/targets/x86_64-linux/include)
16 | link_directories(/usr/local/cuda-10.1/targets/x86_64-linux/lib)
17 | link_directories(/usr/local/cuda-10.1/targets/x86_64-linux/lib/stubs)
18 | 
19 | 
20 | list(APPEND LIBS PUBLIC c10)
21 | list(APPEND LIBS PUBLIC torch)
22 | list(APPEND LIBS PUBLIC nvrtc)
23 | list(APPEND LIBS PUBLIC cuda)
24 | list(APPEND LIBS PUBLIC caffe2)
25 | list(APPEND LIBS PUBLIC caffe2_gpu)
26 | 
27 | list(APPEND LIBS PUBLIC opencv_core)
28 | list(APPEND LIBS PUBLIC opencv_imgproc)
29 | list(APPEND LIBS PUBLIC opencv_imgcodecs)
30 | list(APPEND LIBS PUBLIC opencv_highgui)
31 | list(APPEND LIBS PUBLIC opencv_ml)
32 | list(APPEND LIBS PUBLIC opencv_video)
33 | list(APPEND LIBS PUBLIC opencv_videoio)
34 | list(APPEND LIBS PUBLIC opencv_calib3d)
35 | 
36 | include_directories(${CMAKE_SOURCE_DIR}/base)
37 | aux_source_directory(${CMAKE_SOURCE_DIR}/base dirSrc)
38 | 
39 | add_executable(main main.cpp ${dirSrc})
40 | target_link_libraries(main ${LIBS})
41 | 


--------------------------------------------------------------------------------
/base/common.hpp:
--------------------------------------------------------------------------------
 1 | 
 2 | #ifndef MAIN_COMMON_HPP
 3 | #define MAIN_COMMON_HPP
 4 | 
 5 | #include <iostream>
 6 | #include <vector>
 7 | #include "opencv2/opencv.hpp"
 8 | 
 9 | enum eReturnStatus{
10 |     ProcessSuccess = 0,
11 |     ProcessFailure = 1
12 | };
13 | 
14 | ///////////////////////////////////
15 | //函数功能：获取文件夹下所有文件的名称
16 | //输入：folderPath 文件夹地址
17 | //输出：fileNames 文件名称
18 | ///////////////////////////////////
19 | int getFileNameFromFolder(const std::string &folderPath, std::vector<std::string> &fileNames);
20 | 
21 | ///////////////////////////////////
22 | //函数功能：分割字符串
23 | //输入：str 待分割字符串, div 分割符
24 | //输出：outStr 分割结果
25 | ///////////////////////////////////
26 | void splitString(const std::string &str, std::vector<std::string> &outStr, const char div = ' ');
27 | 
28 | ///////////////////////////////////
29 | //函数功能：softmax计算
30 | //输入：values
31 | //输出：values
32 | ///////////////////////////////////
33 | void softmax(std::vector<float>& values);
34 | 
35 | ///////////////////////////////////
36 | //函数功能：图像扩增方法, ten crop
37 | //输入：img 输入图像,  ratio crop的比例，默认0.9
38 | //输出：augImgs 扩增后图像
39 | ///////////////////////////////////
40 | int tenCrop(cv::Mat &img, std::vector<cv::Mat> &augImgs, const float &ratio = 0.9);
41 | 
42 | ///////////////////////////////////
43 | //函数功能：创建文件夹
44 | //输入：path 文件夹地址,支持同时创建多级目录
45 | ///////////////////////////////////
46 | void makeDir(const std::string &path);
47 | 
48 | 
49 | #endif //MAIN_COMMON_HPP
50 | 


--------------------------------------------------------------------------------
/base/classification.hpp:
--------------------------------------------------------------------------------
 1 | 
 2 | #ifndef MAIN_CLASSIFICATION_HPP
 3 | #define MAIN_CLASSIFICATION_HPP
 4 | 
 5 | #include "torch/script.h"
 6 | #include "opencv2/opencv.hpp"
 7 | 
 8 | class classification {
 9 | public:
10 |     typedef struct classificationResult{
11 |         std::vector<float> softmaxRes; // 分类结果, softmax后概率
12 |         std::vector<float> beforeSoftmaxRes; // 分类结果, softmax前特征
13 |         int inferenceRes = -1; // 类别结果
14 |         float score = 0.0; // 当前类别对应的得分
15 |     }clsOutput;
16 | public:
17 |     classification(){;}
18 |     ~classification(){;}
19 | 
20 |     ///////////////////////////////////
21 |     //函数功能：模型初始化
22 |     //输入：modelPath　模型地址
23 |     ///////////////////////////////////
24 |     int init(const std::string &modelPath);
25 | 
26 |     ///////////////////////////////////
27 |     //函数功能：模型推理，支持单张图
28 |     //输入：img 输入图片
29 |     //输出：out 推理结果
30 |     ///////////////////////////////////
31 |     int inference(const cv::Mat &img, clsOutput &out);
32 | 
33 |     ///////////////////////////////////
34 |     //函数功能：模型推理，支持多张图
35 |     //输入：imgs 输入图片组
36 |     //输出：outs 推理结果
37 |     ///////////////////////////////////
38 |     int inferenceVector(const std::vector<cv::Mat> &imgs, std::vector<clsOutput> &outs);
39 | 
40 |     ///////////////////////////////////
41 |     //函数功能：设置模型均值参数
42 |     //输入：str 均值参数
43 |     ///////////////////////////////////
44 |     void setMean(const std::string &str);
45 | 
46 |     ///////////////////////////////////
47 |     //函数功能：设置模型方差参数
48 |     //输入：str 方差参数
49 |     ///////////////////////////////////
50 |     void setStd(const std::string &str);
51 | 
52 | private:
53 |     ///////////////////////////////////
54 |     //函数功能：图片预处理
55 |     //输入：img 输入图片
56 |     //输出：outTensor 预处理后的结果
57 |     ///////////////////////////////////
58 |     int preProcess(const cv::Mat &img, torch::Tensor &outTensor);
59 | 
60 |     ///////////////////////////////////
61 |     //函数功能：模型推理结果后处理
62 |     //输入：inputTensor 输入推理后的特征
63 |     //输出：out 后处理最终结果
64 |     ///////////////////////////////////
65 |     int postProcess(const torch::Tensor &inputTensor, clsOutput &out);
66 | 
67 | private:
68 |     std::shared_ptr<torch::jit::script::Module> mNet;
69 |     std::vector<float> mMean;
70 |     std::vector<float> mStd;
71 | };
72 | 
73 | 
74 | #endif //MAIN_CLASSIFICATION_HPP
75 | 


--------------------------------------------------------------------------------
/base/classification.cpp:
--------------------------------------------------------------------------------
 1 | 
 2 | #include "classification.hpp"
 3 | #include "common.hpp"
 4 | 
 5 | int classification::init(const std::string &modelPath){
 6 |     mNet = torch::jit::load(modelPath, torch::kCUDA);
 7 |     return eReturnStatus::ProcessSuccess;
 8 | }
 9 | 
10 | void classification::setMean(const std::string &str){
11 |     std::vector<std::string> sMean;
12 |     splitString(str, sMean, ',');
13 |     for(auto m : sMean){
14 |         mMean.push_back(std::atof(m.c_str()));
15 |     }
16 | }
17 | 
18 | void classification::setStd(const std::string &str){
19 |     std::vector<std::string> sStd;
20 |     splitString(str, sStd, ',');
21 |     for(auto s : sStd){
22 |         mStd.push_back(std::atof(s.c_str()));
23 |     }
24 | }
25 | 
26 | int classification::inference(const cv::Mat &img, clsOutput &out){
27 |     if(img.empty()){
28 |         return eReturnStatus::ProcessFailure;
29 |     }
30 |     torch::Tensor outTensor;
31 |     preProcess(img, outTensor);
32 |     torch::Tensor res = mNet->forward({outTensor}).toTensor();
33 |     postProcess(res, out);
34 |     return eReturnStatus::ProcessSuccess;
35 | }
36 | 
37 | int classification::inferenceVector(const std::vector<cv::Mat> &imgs, std::vector<clsOutput> &outs){
38 |     for(int i = 0; i < imgs.size(); i++){
39 |         clsOutput out;
40 |         inference(imgs[i], out);
41 |         outs.push_back(out);
42 |     }
43 |     return eReturnStatus::ProcessSuccess;
44 | }
45 | 
46 | int classification::preProcess(const cv::Mat &img, torch::Tensor &outTensor){
47 |     //根据训练方式设置
48 |     cv::Mat img1, imgFloat;
49 |     cvtColor(img, img1, cv::COLOR_BGR2RGB);
50 |     img1.convertTo(imgFloat, CV_32F, 1.0 / 255);
51 |     if(3 == mMean.size()){
52 |         cv::Scalar mean(mMean[0], mMean[1], mMean[2]);
53 |         imgFloat -= mean;
54 |     }
55 |     if(3 == mStd.size()){
56 |         std::vector<cv::Mat> vImgFloat;
57 |         cv::split(imgFloat, vImgFloat);
58 |         vImgFloat[0] /= mStd[0];
59 |         vImgFloat[1] /= mStd[1];
60 |         vImgFloat[2] /= mStd[2];
61 |         cv::merge(vImgFloat, imgFloat);
62 |     }
63 |     torch::Tensor tImg = torch::CPU(torch::kFloat32).tensorFromBlob(imgFloat.data, { 1, imgFloat.rows, imgFloat.cols, 3 });
64 |     tImg = tImg.permute({ 0,3,1,2 });
65 |     outTensor = torch::autograd::make_variable(tImg, false).to(at::kCUDA);
66 |     return eReturnStatus::ProcessSuccess;
67 | }
68 | 
69 | int classification::postProcess(const torch::Tensor &inputTensor, clsOutput &out){
70 |     torch::Tensor resTmp = inputTensor.softmax(1).to(torch::kCPU);
71 |     for(int i = 0; i < resTmp.size(1); i++){
72 |         out.softmaxRes.push_back(resTmp[0][i].item().toFloat());
73 |     }
74 |     out.inferenceRes = std::get<1>(resTmp.max(1, true)).item().toInt();
75 |     out.score = out.softmaxRes[out.inferenceRes];
76 |     inputTensor.to(torch::kCPU);
77 |     for(int i = 0; i < inputTensor.size(1); i++){
78 |         out.beforeSoftmaxRes.push_back(inputTensor[0][i].item().toFloat());
79 |     }
80 |     return eReturnStatus::ProcessSuccess;
81 | }


--------------------------------------------------------------------------------
/base/common.cpp:
--------------------------------------------------------------------------------
 1 | 
 2 | #include "common.hpp"
 3 | #include <sstream>
 4 | #include <dirent.h>
 5 | #include <string>
 6 | #include "torch/script.h"
 7 | #include <unistd.h>
 8 | 
 9 | int getFileNameFromFolder(const std::string &folderPath, std::vector<std::string> &fileNames){
10 |     DIR *pDir;
11 |     struct dirent* pDirent;
12 |     pDir = opendir(folderPath.c_str());
13 |     if(NULL == pDir){
14 |         std::cout << "dir is null..." << std::endl;
15 |         return 0;
16 |     }
17 |     while(NULL != (pDirent = readdir(pDir))){
18 |         if(0 == strcmp(pDirent->d_name,".") || 0 == strcmp(pDirent->d_name,"..")){
19 |             continue;
20 |         }
21 |         else if (4 == pDirent->d_type){
22 |             std::string folderPathNew = folderPath + "/" + pDirent->d_name;
23 |             getFileNameFromFolder(folderPathNew, fileNames);
24 |         }
25 |         else{
26 |             std::string fileName = folderPath + "/" + pDirent->d_name;
27 |             fileNames.push_back(fileName);
28 |         }
29 |     }
30 |     closedir(pDir);
31 |     return 1;
32 | }
33 | 
34 | void splitString(const std::string &str, std::vector<std::string> &outStr, const char div){
35 |     outStr.clear();
36 |     std::istringstream iss(str);
37 |     std::string tmp;
38 |     while (std::getline(iss, tmp, div)) {
39 |         if (tmp != "") {
40 |             outStr.emplace_back(std::move(tmp));
41 |         }
42 |     }
43 | }
44 | 
45 | void softmax(std::vector<float>& values){
46 |     torch::Tensor tmpValues = torch::tensor(values);
47 |     tmpValues = tmpValues.softmax(0);
48 |     for(int i = 0; i < values.size(); i++){
49 |         values[i] = tmpValues[i].item().toFloat();
50 |     }
51 | }
52 | 
53 | int tenCrop(cv::Mat &img, std::vector<cv::Mat> &augImgs, const float &ratio){
54 |     if(img.empty()){
55 |         return 0;
56 |     }
57 |     int newWidth = int(ratio*img.cols);
58 |     int newHeight = int(ratio*img.rows);
59 |     int shiftX = int((1-ratio)*img.cols);
60 |     int shiftY = int((1-ratio)*img.rows);
61 |     cv::Mat imgCenter = img(cv::Rect(int(shiftX/2),int(shiftY/2),newWidth, newHeight)).clone() ;
62 |     augImgs.push_back(imgCenter);
63 |     cv::Mat imgCenterFlip;
64 |     cv::flip(imgCenter,imgCenterFlip,1);
65 |     augImgs.push_back(imgCenterFlip);
66 |     for(int i = 0; i < 2; i++){
67 |         for(int j = 0; j < 2; j++){
68 |             cv::Rect rt = cv::Rect(i*shiftX, j*shiftY, newWidth, newHeight);
69 |             cv::Mat imgCenter_;
70 |             imgCenter_ = img(rt).clone();
71 |             augImgs.push_back(imgCenter_);
72 |             cv::Mat imgCenterFlip_;
73 |             cv::flip(imgCenter_,imgCenterFlip_,1);
74 |             augImgs.push_back(imgCenterFlip_);
75 |         }
76 |     }
77 |     return 1;
78 | }
79 | 
80 | void makeDir(const std::string &path){
81 |     std::vector<std::string> pathSplit;
82 |     splitString(path, pathSplit, '/');
83 |     std::string filePath = "/";
84 |     for(int i = 0; i < pathSplit.size(); i++){
85 |         filePath += pathSplit[i];
86 |         if(!access(filePath.c_str(),0)){
87 |             filePath += "/";
88 |             continue;
89 |         }
90 |         std::string command = "mkdir " + filePath;
91 |         system(command.c_str());
92 |         filePath += "/";
93 |     }
94 | }


--------------------------------------------------------------------------------
/main.cpp:
--------------------------------------------------------------------------------
  1 | #include <iostream>
  2 | #include "classification.hpp"
  3 | #include "common.hpp"
  4 | int testNormal() {
  5 |     std::string modelPath = "./model/cls.pt";
  6 |     std::string sMean = "0.485,0.456,0.406";
  7 |     std::string sStd = "0.229,0.224,0.225";
  8 | 
  9 |     classification cls;
 10 |     cls.setMean(sMean);
 11 |     cls.setStd(sStd);
 12 |     cls.init(modelPath);
 13 | 
 14 |     std::string imgPath = "./data";
 15 |     std::vector<std::string> fileNames;
 16 |     getFileNameFromFolder(imgPath, fileNames);
 17 |     for(auto name : fileNames){
 18 |         cv::Mat img = cv::imread(name);
 19 |         classification::clsOutput out;
 20 |         cls.inference(img, out);
 21 |         softmax(out.beforeSoftmaxRes);
 22 |     }
 23 |     return 0;
 24 | }
 25 | 
 26 | enum eTestMode{
 27 |     IteratedData = 0, //迭代数据
 28 |     CleanData    = 1  //清洗数据
 29 | };
 30 | 
 31 | struct inputParam{
 32 |     eTestMode   mode             = CleanData;
 33 |     std::string modelProjectPath = "./model/cls1.pt";
 34 |     std::string modelProjectMean = "0.485,0.456,0.406";
 35 |     std::string modelProjectStd  = "0.229,0.224,0.225";
 36 |     cv::Size    modelProjectSize = cv::Size(112,48);
 37 |     std::string modelMerge1Path  = "./model/cls2.pt";
 38 |     std::string modelMerge1Mean  = "0.485,0.456,0.406";
 39 |     std::string modelMerge1Std   = "0.229,0.224,0.225";
 40 |     cv::Size    modelMerge1Size  = cv::Size(112,48);
 41 |     std::string modelMerge2Path  = "./model/cls3.pt";
 42 |     std::string modelMerge2Mean  = "0.485,0.456,0.406";
 43 |     std::string modelMerge2Std   = "0.229,0.224,0.225";
 44 |     cv::Size    modelMerge2Size  = cv::Size(112,48);
 45 |     std::string srcDataPath      = "./data/src";
 46 |     std::string dstDataPath      = "./data/save";
 47 |     int         classNum         = 2;
 48 |     float       thresh           = 0.8;
 49 | };
 50 | 
 51 | int inference(const inputParam& input){
 52 |     //模型初始化
 53 |     classification projectModel, merge1Model, merge2Model;
 54 |     if(IteratedData == input.mode){
 55 |         projectModel.init(input.modelProjectPath);
 56 |         projectModel.setMean(input.modelProjectMean);
 57 |         projectModel.setMean(input.modelProjectStd);
 58 |     }
 59 |     merge1Model.init(input.modelMerge1Path);
 60 |     merge1Model.setMean(input.modelMerge1Mean);
 61 |     merge1Model.setMean(input.modelMerge1Std);
 62 |     merge2Model.init(input.modelMerge2Path);
 63 |     merge2Model.setMean(input.modelMerge2Mean);
 64 |     merge2Model.setMean(input.modelMerge2Std);
 65 | 
 66 |     // 创建文件夹
 67 |     for(int i = 0; i < input.classNum; i++){
 68 |         makeDir(input.dstDataPath + "/" + std::to_string(i) + "N");
 69 |         for(int j = 0; j < input.classNum; j++){
 70 |             makeDir(input.dstDataPath + "/" + std::to_string(i) + "/" + std::to_string(j));
 71 |             makeDir(input.dstDataPath + "/" + std::to_string(i) + "/" + std::to_string(j) + "N");
 72 |         }
 73 |     }
 74 | 
 75 |     //数据处理
 76 |     std::vector<std::string> imgPaths;
 77 |     getFileNameFromFolder(input.srcDataPath, imgPaths);
 78 |     for(int i = 0; i < imgPaths.size(); i++){
 79 |         if(0 == i%100){
 80 |             std::cout << "processed " << i << " imgs..." << std::endl;
 81 |         }
 82 |         cv::Mat img = cv::imread(imgPaths[i], 1);
 83 |         if(img.empty()) continue;
 84 |         classification::clsOutput projectClsOutput;
 85 |         std::vector<classification::clsOutput> merge1Outputs, merge2Outputs;
 86 |         if(IteratedData == input.mode){
 87 |             cv::Mat imgProject;
 88 |             cv::resize(img, imgProject, input.modelProjectSize);
 89 |             projectModel.inference(img,projectClsOutput);
 90 |         }else{
 91 |             std::vector<std::string> pathSplit;
 92 |             splitString(imgPaths[i], pathSplit, '/');
 93 |             projectClsOutput.inferenceRes = std::atoi(pathSplit[pathSplit.size()-2].c_str());
 94 |             projectClsOutput.score = 1.0;
 95 |         }
 96 |         std::vector<cv::Mat> imgs;
 97 |         tenCrop(img, imgs, 0.9);
 98 | //        for(auto m : imgs){
 99 | //            cv::imshow("show", m);
100 | //            cv::waitKey();
101 | //        }
102 |         merge1Model.inferenceVector(imgs, merge1Outputs);
103 |         merge2Model.inferenceVector(imgs, merge2Outputs);
104 |         std::vector<float> mergeFeature(input.classNum);
105 |         for(int m = 0; m < input.classNum; m++){
106 |             for(int n = 0; n < merge1Outputs.size(); n++){
107 |                 mergeFeature[m] += merge1Outputs[n].beforeSoftmaxRes[m];
108 |                 mergeFeature[m] += merge2Outputs[n].beforeSoftmaxRes[m];
109 |             }
110 |             mergeFeature[m] /= 2*merge1Outputs.size();
111 |         }
112 |         softmax(mergeFeature);
113 |         int maxPosition = max_element(mergeFeature.begin(),mergeFeature.end()) - mergeFeature.begin();
114 |         if(projectClsOutput.score < input.thresh){
115 |             std::string dstPath = input.dstDataPath + "/" + std::to_string(projectClsOutput.inferenceRes) + "N";
116 |             std::string command = "cp " + imgPaths[i] + " " + dstPath;
117 |             system(command.c_str());
118 |         } else {
119 |             if(mergeFeature[maxPosition] < input.thresh){
120 |                 std::string dstPath = input.dstDataPath + "/" + std::to_string(projectClsOutput.inferenceRes)
121 |                                       + "/" + std::to_string(maxPosition) + "N";
122 |                 std::string command = "cp " + imgPaths[i] + " " + dstPath;
123 |                 system(command.c_str());
124 |             }else{
125 |                 std::string dstPath = input.dstDataPath + "/" + std::to_string(projectClsOutput.inferenceRes)
126 |                                       + "/" + std::to_string(maxPosition);
127 |                 std::string command = "cp " + imgPaths[i] + " " + dstPath;
128 |                 system(command.c_str());
129 |             }
130 |         }
131 |     }
132 |     return 1;
133 | }
134 | 
135 | int main(){
136 |     inputParam input;
137 |     inference(input);
138 | }
139 | 


--------------------------------------------------------------------------------