177 lines
7.1 KiB
C++
177 lines
7.1 KiB
C++
#include <sstream>
|
|
#include <iostream>
|
|
|
|
#include "opencv2/quality.hpp"
|
|
#include "opencv2/quality/quality_utils.hpp"
|
|
#include "opencv2/imgcodecs.hpp"
|
|
#include "opencv2/ml.hpp"
|
|
|
|
/*
|
|
BRISQUE Trainer using LIVE DB R2
|
|
http://live.ece.utexas.edu/research/Quality/subjective.htm
|
|
H.R. Sheikh, Z.Wang, L. Cormack and A.C. Bovik, "LIVE Image Quality Assessment Database Release 2", http://live.ece.utexas.edu/research/quality .
|
|
H.R. Sheikh, M.F. Sabir and A.C. Bovik, "A statistical evaluation of recent full reference image quality assessment algorithms", IEEE Transactions on Image Processing, vol. 15, no. 11, pp. 3440-3451, Nov. 2006.
|
|
Z. Wang, A.C. Bovik, H.R. Sheikh and E.P. Simoncelli, "Image quality assessment: from error visibility to structural similarity," IEEE Transactions on Image Processing , vol.13, no.4, pp. 600- 612, April 2004.
|
|
*/
|
|
|
|
/*
|
|
Copyright (c) 2011 The University of Texas at Austin
|
|
All rights reserved.
|
|
|
|
Permission is hereby granted, without written agreement and without license or royalty fees, to use, copy,
|
|
modify, and distribute this code (the source files) and its documentation for
|
|
any purpose, provided that the copyright notice in its entirety appear in all copies of this code, and the
|
|
original source of this code, Laboratory for Image and Video Engineering (LIVE, http://live.ece.utexas.edu)
|
|
and Center for Perceptual Systems (CPS, http://www.cps.utexas.edu) at the University of Texas at Austin (UT Austin,
|
|
http://www.utexas.edu), is acknowledged in any publication that reports research using this code. The research
|
|
is to be cited in the bibliography as:
|
|
|
|
1) A. Mittal, A. K. Moorthy and A. C. Bovik, "BRISQUE Software Release",
|
|
URL: http://live.ece.utexas.edu/research/quality/BRISQUE_release.zip, 2011
|
|
|
|
2) A. Mittal, A. K. Moorthy and A. C. Bovik, "No Reference Image Quality Assessment in the Spatial Domain"
|
|
submitted
|
|
|
|
IN NO EVENT SHALL THE UNIVERSITY OF TEXAS AT AUSTIN BE LIABLE TO ANY PARTY FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL,
|
|
OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OF THIS DATABASE AND ITS DOCUMENTATION, EVEN IF THE UNIVERSITY OF TEXAS
|
|
AT AUSTIN HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
|
|
THE UNIVERSITY OF TEXAS AT AUSTIN SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
|
|
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE DATABASE PROVIDED HEREUNDER IS ON AN "AS IS" BASIS,
|
|
AND THE UNIVERSITY OF TEXAS AT AUSTIN HAS NO OBLIGATION TO PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
|
|
*/
|
|
|
|
/* Original Paper: @cite Mittal2 and Original Implementation: @cite Mittal2_software */
|
|
|
|
namespace {
|
|
|
|
#define CATEGORIES 5
|
|
#define IMAGENUM 982
|
|
#define JP2KNUM 227
|
|
#define JPEGNUM 233
|
|
#define WNNUM 174
|
|
#define GBLURNUM 174
|
|
#define FFNUM 174
|
|
|
|
// collects training data from LIVE R2 database
|
|
// returns {features, responses}, 1 row per image
|
|
std::pair<cv::Mat, cv::Mat> collect_data_live_r2(const std::string& foldername)
|
|
{
|
|
FILE* fid = nullptr;
|
|
|
|
//----------------------------------------------------
|
|
// class is the distortion category, there are 982 images in LIVE database
|
|
std::vector<std::string> distortionlabels;
|
|
distortionlabels.push_back("jp2k");
|
|
distortionlabels.push_back("jpeg");
|
|
distortionlabels.push_back("wn");
|
|
distortionlabels.push_back("gblur");
|
|
distortionlabels.push_back("fastfading");
|
|
|
|
int imnumber[5] = { 0,227,460,634,808 };
|
|
|
|
std::vector<int>categorylabels;
|
|
categorylabels.insert(categorylabels.end(), JP2KNUM, 0);
|
|
categorylabels.insert(categorylabels.end(), JPEGNUM, 1);
|
|
categorylabels.insert(categorylabels.end(), WNNUM, 2);
|
|
categorylabels.insert(categorylabels.end(), GBLURNUM, 3);
|
|
categorylabels.insert(categorylabels.end(), FFNUM, 4);
|
|
|
|
int iforg[IMAGENUM];
|
|
fid = fopen((foldername + "orgs.txt").c_str(), "r");
|
|
for (int itr = 0; itr < IMAGENUM; itr++)
|
|
CV_Assert( fscanf(fid, "%d", iforg + itr) > 0);
|
|
fclose(fid);
|
|
|
|
float dmosscores[IMAGENUM];
|
|
fid = fopen((foldername + "dmos.txt").c_str(), "r");
|
|
for (int itr = 0; itr < IMAGENUM; itr++)
|
|
CV_Assert( fscanf(fid, "%f", dmosscores + itr) > 0 );
|
|
fclose(fid);
|
|
|
|
// features vector, 1 row per image
|
|
cv::Mat features(0, 0, CV_32FC1);
|
|
|
|
// response vector, 1 row per image
|
|
cv::Mat responses(0, 1, CV_32FC1);
|
|
|
|
for (int itr = 0; itr < IMAGENUM; itr++)
|
|
{
|
|
//Dont compute features for original images
|
|
if (iforg[itr])
|
|
continue;
|
|
|
|
// append dmos score
|
|
float score = dmosscores[itr];
|
|
responses.push_back(cv::Mat(1, 1, CV_32FC1, (void*)&score));
|
|
|
|
// load image, calc features
|
|
std::string imname = "";
|
|
imname.append(foldername);
|
|
imname.append("/");
|
|
imname.append(distortionlabels[categorylabels[itr]].c_str());
|
|
imname.append("/img");
|
|
imname += std::to_string((itr - imnumber[categorylabels[itr]] + 1));
|
|
imname.append(".bmp");
|
|
|
|
cv::Mat im_features;
|
|
cv::quality::QualityBRISQUE::computeFeatures(cv::imread(imname), im_features); // outputs a row vector
|
|
|
|
features.push_back(im_features.row(0)); // append row vector
|
|
}
|
|
|
|
return std::make_pair(std::move(features), std::move(responses));
|
|
} // collect_data_live_r2
|
|
}
|
|
|
|
inline void printHelp()
|
|
{
|
|
using namespace std;
|
|
cout << " Demo of training BRISQUE quality assessment model using LIVE R2 database." << endl;
|
|
cout << " A. Mittal, A. K. Moorthy and A. C. Bovik, 'No Reference Image Quality Assessment in the Spatial Domain'" << std::endl << std::endl;
|
|
|
|
cout << " Usage: program <live_r2_db_path> <output_model_path> <output_range_path>" << endl << endl;
|
|
}
|
|
|
|
int main(int argc, const char * argv[])
|
|
{
|
|
using namespace cv::ml;
|
|
|
|
if (argc != 4)
|
|
{
|
|
printHelp();
|
|
exit(1);
|
|
}
|
|
|
|
std::cout << "Training BRISQUE on database at " << argv[1] << "..." << std::endl;
|
|
|
|
// collect data from the data set
|
|
auto data = collect_data_live_r2( std::string( argv[1] ) + "/" );
|
|
|
|
// extract column ranges for features
|
|
const auto range = cv::quality::quality_utils::get_column_range(data.first);
|
|
|
|
// scale all features from -1 to 1
|
|
cv::quality::quality_utils::scale<float>(data.first, range, -1.f, 1.f);
|
|
|
|
// do training, output train file
|
|
// libsvm call from original BRISQUE impl: svm-train -s 3 -g 0.05 -c 1024 -b 1 -q train_scale allmodel
|
|
auto svm = SVM::create();
|
|
svm->setType(SVM::Types::EPS_SVR);
|
|
svm->setKernel(SVM::KernelTypes::RBF);
|
|
svm->setGamma(0.05);
|
|
svm->setC(1024.);
|
|
svm->setTermCriteria(cv::TermCriteria(cv::TermCriteria::Type::EPS, 1000, 0.001));
|
|
svm->setP(.1);// default p (epsilon) from libsvm
|
|
|
|
svm->train(data.first, cv::ml::ROW_SAMPLE, data.second);
|
|
svm->save( argv[2] ); // save to location specified in argv[2]
|
|
|
|
// output scale file to argv[3]
|
|
cv::Mat range_mat(range);
|
|
cv::FileStorage fs(argv[3], cv::FileStorage::WRITE );
|
|
fs << "range" << range_mat;
|
|
|
|
return 0;
|
|
}
|