Resnet改进的色情图片检测-nsfw

来源:互联网 发布:java调用kettle 编辑:程序博客网 时间:2024/05/21 14:02


nsfw(Not suitable for work classifier)是雅虎开源的进行色情检测识别的一个网络,该网络基于Resnet50修改,取得了非常不错的色情检测效果。

在安装caffe之后就可以运行上面的网络,运行shell脚本如下。

#!/usr/bin/env shpython ./classify_nsfw.py \ --model_defnsfw_model/deploy.prototxt \ --pretrained_modelnsfw_model/resnet_50_1by2_nsfw.caffemodel \ test.jpg

根据作者提供的python程序实现的基于c++的图形化显示,程序如下,

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_ONLYCaffe::set_mode(Caffe::CPU);#elseCaffe::set_mode(Caffe::GPU);#endifnet_.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);elsesample = img;cv::Mat sample_resized;if (sample.size() != input_geometry_)cv::resize(sample, sample_resized, input_geometry_);elsesample_resized = sample;cv::Mat sample_float;if (num_channels_ == 3)sample_resized.convertTo(sample_float, CV_32FC3);elsesample_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);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;}

实验结果如下,

程序链接:http://download.csdn.net/detail/qq_14845119/9751834

reference

https://github.com/yahoo/open_nsfw


1 0
原创粉丝点击