【严肃脸】使用caffe实现色情图片的识别

前言

前几天看到了雅虎开源了一个色情图片的识别模型新闻，上Github一看，是基于caffe的。试了试，模型效果很赞。Github地址：https://github.com/yahoo/open_nsfw

至于测试的数据集，就自行找图吧（逃

关于在程序中使用caffe可以戳我的这一篇博客：http://blog.csdn.net/mr_curry/article/details/52443126
（如何在程序中像使用OpenCV一样使用caffe）

准备

下载好模型和配置文件，观察网络结构。

输入的图片需为彩色图片，尺寸为224*224（Vgg的网络也是224*224）.根据最后一层，Softmax将会输出一个概率（图片有多色？）

再打开压缩包中的.py文件，我们可以观察到图像的均值：

我们对网络结构做如下修改，使用MemoryData层：

代码：

为了表示的更为明显，可以用line函数进行画线：波动越大，表示当前图片越...

caffe_predefine.h:

#include "caffe/layers/input_layer.hpp"
#include "caffe/layers/inner_product_layer.hpp"
#include "caffe/layers/dropout_layer.hpp"
#include "caffe/layers/conv_layer.hpp"
#include "caffe/layers/relu_layer.hpp"
#include <iostream>
#include "caffe/caffe.hpp"
#include <opencv.hpp>
#include <caffe/layers/memory_data_layer.hpp>
#include "caffe/layers/pooling_layer.hpp"
#include "caffe/layers/lrn_layer.hpp"
#include "caffe/layers/softmax_layer.hpp"
#include <caffe/layers/batch_norm_layer.hpp>
#include <caffe/layers/scale_layer.hpp>
#include <caffe/layers/eltwise_layer.hpp>
#include <caffe/layers/bias_layer.hpp>

caffe::MemoryDataLayer<float> *memory_layer;
caffe::Net<float>* net;

DrawLine.h:

#include <opencv.hpp>

using namespace cv;
using namespace std;

void DrawLine(Mat T,vector<Point> point_array);

load_model.h:

#include <opencv.hpp>
using namespace cv;
using namespace std;

void Caffe_Predefine();
float getProb(Mat source);

load_model.cpp:

#include <caffe_predefine.h>
#include <load_model.h>
namespace caffe
{
	extern INSTANTIATE_CLASS(InputLayer);
	extern INSTANTIATE_CLASS(InnerProductLayer);
	extern INSTANTIATE_CLASS(DropoutLayer);
	extern INSTANTIATE_CLASS(ConvolutionLayer);
	REGISTER_LAYER_CLASS(Convolution);
	extern INSTANTIATE_CLASS(ReLULayer);
	REGISTER_LAYER_CLASS(ReLU);
	extern INSTANTIATE_CLASS(PoolingLayer);
	REGISTER_LAYER_CLASS(Pooling);
	extern INSTANTIATE_CLASS(LRNLayer);
	REGISTER_LAYER_CLASS(LRN);
	extern INSTANTIATE_CLASS(SoftmaxLayer);
	REGISTER_LAYER_CLASS(Softmax);
	extern INSTANTIATE_CLASS(MemoryDataLayer);
	extern INSTANTIATE_CLASS(BatchNormLayer);
	extern INSTANTIATE_CLASS(ScaleLayer);
	extern INSTANTIATE_CLASS(EltwiseLayer);
	extern INSTANTIATE_CLASS(BiasLayer);
}

template <typename Dtype>
caffe::Net<Dtype>* Net_Init_Load(std::string param_file, std::string pretrained_param_file, caffe::Phase phase)
{
	caffe::Net<Dtype>* net(new caffe::Net<Dtype>(param_file, caffe::TEST));
	net->CopyTrainedLayersFrom(pretrained_param_file);
	return net;
}

void Caffe_Predefine()//when our code begining run must add it
{
	caffe::Caffe::set_mode(caffe::Caffe::CPU);
	net = Net_Init_Load<float>("open_nsfw_memorydata.prototxt", "resnet_50_1by2_nsfw.caffemodel", caffe::TEST);
	memory_layer = (caffe::MemoryDataLayer<float> *)net->layers()[0].get();
}

float getProb(Mat source)
{
	vector<Mat> test;
	vector<int> label;
	test.push_back(source);
	label.push_back(0);
	memory_layer->AddMatVector(test, label);// memory_layer and net , must be define be a global variable.
	std::vector<caffe::Blob<float>*> input_vec;
	net->Forward(input_vec);
	boost::shared_ptr<caffe::Blob<float>  > prob = net->blob_by_name("prob");
	return prob->data_at(0, 0, 1, 0);
}

DrawLine.cpp:

#include <DrawLine.h>

void DrawLine(Mat T, vector<Point> point_array)
{
	for (int i = 1; i < point_array.size();i++)
		line(T, point_array[i-1], point_array[i], Scalar(0, 0, 255), 3);
}

Main.cpp:

#include <load_model.h>

#include <DrawLine.h>

#define X 0
#define Y 200
int main()
{
	Caffe_Predefine();
	VideoCapture cap("test.mp4");
	Mat frame;
	float x = 1, y;

	vector<Point> point_array;

	Point T_s(X, Y);
	point_array.push_back(T_s);
	while (true)
	{
		cap >> frame;
		if (!frame.empty())
		{
			y = getProb(frame);
			cout <<"当前概率为"<< y << endl;
			Point T_l(X+x++, (Y-100*y));//*100为了更为明显显示
			point_array.push_back(T_l);
			DrawLine(frame, point_array);
			imshow("NSFW", frame);
			waitKey(1);
		}
		else
		{
			break;
		}
	}
}

效果：

使用了某预告片来做显示：