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

 

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

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

 

#!/usr/bin/env sh
python ./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 Predict(const cv::Mat& img);
	void WrapInputLayer(std::vector* input_channels);
	void Preprocess(const cv::Mat& img,	std::vector* input_channels);

private:
	shared_ptr > 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(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* 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 output = Predict(img);
	return output[1];
}


std::vector Classifier::Predict(const cv::Mat& img) {
	Blob* input_layer = net_->input_blobs()[0];
	input_layer->Reshape(1, num_channels_,
		input_geometry_.height, input_geometry_.width);
	net_->Reshape();

	std::vector input_channels;
	WrapInputLayer(&input_channels);

	Preprocess(img, &input_channels);

	net_->Forward();

	Blob* output_layer = net_->output_blobs()[0];
	const float* begin = output_layer->cpu_data();
	const float* end = begin + output_layer->channels();
	return std::vector(begin, end);
}

void Classifier::WrapInputLayer(std::vector* input_channels) {
	Blob* 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* 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);
	
	vectorchannels_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(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