【嚴肅臉】使用caffe實現色情圖片的識別

ChenJoya發表於2016-10-04

前言

前幾天看到了雅虎開源了一個色情圖片的識別模型新聞,上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;
		}
	}
}

效果:

使用了某預告片來做顯示:




相關文章