open_nsfw.zip

#define USE_OPENCV 1 #define CPU_ONLY 1 #include <caffe/caffe.hpp> #ifdef USE_OPENCV #include <opencv2/core/core.hpp> #include <opencv2/highgui/highgui.hpp> #include <opencv2/imgproc/imgproc.hpp> #endif // USE_OPENCV #include <algorithm> #include <iosfwd> #include <memory> #include <string> #include <utility> #include <vector> #include "head.h" using namespace caffe; using std::string; class Classifier { public: Classifier(const string& model_file,const string& trained_file); float Classify(const cv::Mat& img); private: std::vector<float> Predict(const cv::Mat& img); void WrapInputLayer(std::vector<cv::Mat>* input_channels); void Preprocess(const cv::Mat& img, std::vector<cv::Mat>* input_channels); private: shared_ptr<Net<float> > net_; cv::Size input_geometry_; int num_channels_; cv::Mat mean_; }; Classifier::Classifier(const string& model_file,const string& trained_file) { #ifdef CPU_ONLY Caffe::set_mode(Caffe::CPU); #else Caffe::set_mode(Caffe::GPU); #endif net_.reset(new Net<float>(model_file, TEST)); net_->CopyTrainedLayersFrom(trained_file); CHECK_EQ(net_->num_inputs(), 1) << "Network should have exactly one input."; CHECK_EQ(net_->num_outputs(), 1) << "Network should have exactly one output."; Blob<float>* input_layer = net_->input_blobs()[0]; num_channels_ = input_layer->channels(); CHECK(num_channels_ == 3 || num_channels_ == 1) << "Input layer should have 1 or 3 channels."; input_geometry_ = cv::Size(input_layer->width(), input_layer->height()); } float Classifier::Classify(const cv::Mat& img) { std::vector<float> output = Predict(img); return output[1]; } std::vector<float> Classifier::Predict(const cv::Mat& img) { Blob<float>* input_layer = net_->input_blobs()[0]; input_layer->Reshape(1, num_channels_, input_geometry_.height, input_geometry_.width); net_->Reshape(); std::vector<cv::Mat> input_channels; WrapInputLayer(&input_channels); Preprocess(img, &input_channels); net_->Forward(); Blob<float>* output_layer = net_->output_blobs()[0]; const float* begin = output_layer->cpu_data(); const float* end = begin + output_layer->channels(); return std::vector<float>(begin, end); } void Classifier::WrapInputLayer(std::vector<cv::Mat>* input_channels) { Blob<float>* input_layer = net_->input_blobs()[0]; int width = input_layer->width(); int height = input_layer->height(); float* input_data = input_layer->mutable_cpu_data(); for (int i = 0; i < input_layer->channels(); ++i) { cv::Mat channel(height, width, CV_32FC1, input_data); input_channels->push_back(channel); input_data += width * height; } } void Classifier::Preprocess(const cv::Mat& img, std::vector<cv::Mat>* input_channels) { cv::Mat sample; if (img.channels() == 3 && num_channels_ == 1) cv::cvtColor(img, sample, cv::COLOR_BGR2GRAY); else if (img.channels() == 4 && num_channels_ == 1) cv::cvtColor(img, sample, cv::COLOR_BGRA2GRAY); else if (img.channels() == 4 && num_channels_ == 3) cv::cvtColor(img, sample, cv::COLOR_BGRA2BGR); else if (img.channels() == 1 && num_channels_ == 3) cv::cvtColor(img, sample, cv::COLOR_GRAY2BGR); else sample = img; cv::Mat sample_resized; if (sample.size() != input_geometry_) cv::resize(sample, sample_resized, input_geometry_); else sample_resized = sample; cv::Mat sample_float; if (num_channels_ == 3) sample_resized.convertTo(sample_float, CV_32FC3); else sample_resized.convertTo(sample_float, CV_32FC1); vector<cv::Mat>channels_mean(3); channels_mean[0] = cv::Mat::ones(sample_float.rows, sample_float.cols, CV_32FC1) * 104; channels_mean[1] = cv::Mat::ones(sample_float.rows, sample_float.cols, CV_32FC1) * 117; channels_mean[2] = cv::Mat::ones(sample_float.rows, sample_float.cols, CV_32FC1) * 123; cv::merge(channels_mean, mean_); cv::Mat sample_normalized; cv::subtract(sample_float, mean_, sample_normalized); cv::split(sample_normalized, *input_channels); CHECK(reinterpret_cast<float*>(input_channels->at(0).data) == net_->input_blobs()[0]->cpu_data()) << "Input channels are not wrapping the input layer of the network."; } int main(int argc, char** argv) { ::google::InitGoogleLogging(argv[0]); string model_file = "deploy.prototxt"; string trained_file = "resnet_50_1by2_nsfw.caffemodel"; Classifier classifier(model_file, trained_file); string file = "1.jpg"; std::cout << "---------- Prediction for "<< file << " ----------" << std::endl; cv::Mat img = cv::imread(file, -1); CHECK(!img.empty()) << "Unable to decode image " << file; float prediction = classifier.Classify(img); char str_head[50] = "NSFW Score:"; char str_pre[10]; sprintf_s(str_pre, "%.4f", prediction); std::cout << "Scores < 0.2----->safe" << std::endl; std::cout << "Scores > 0.8----->NSFW" << std::endl; std::cout << "binned for different NSFW levels" << std::endl; std::cout << "Score:" << prediction << std::endl; cv::putText(img, std::strcat(str_head,str_pre), cv::Point(10, img.rows - 20), 3, 1, cv::Scalar(0, 0, 255)); cv::imshow("result", img); cv::waitKey(); return 0; }