caffe神经网络框架的辅助工具（将图片转换为leveldb格式）

caffe中负责整个网络输入的datalayer是从leveldb里读取数据的，是一个google实现的很高效的kv数据库。

因此我们训练网络必须先把数据转成leveldb的格式。

这里我实现的是把一个目录的全部图片转成leveldb的格式。

工具使用命令格格式：convert_imagedata src_dir dst_dir attach_dir channel width height

例子：./convert_imagedata.bin /home/linger/imdata/collar_train/ /home/linger/linger/testfile/crop_train_db/ /home/linger/linger/testfile/crop_train_attachment/
3 50 50

源码：

#include <google/protobuf/text_format.h>

#include <glog/logging.h>

#include <leveldb/db.h>

#include <stdint.h>

#include <fstream>  // NOLINT(readability/streams)

#include <string>

#include <set>

#include <stdio.h>

#include <string.h>

#include <stdlib.h>

#include <dirent.h>

#include <sys/stat.h>

#include <unistd.h>

#include <sys/types.h>

#include "caffe/proto/caffe.pb.h"

#include <opencv2/highgui/highgui.hpp>

#include <opencv2/highgui/highgui_c.h>

#include <opencv2/imgproc/imgproc.hpp>

using std::string;

using namespace std;

set<string> all_class_name;

map<string,int> class2id;

/**

 * path:文件夹

 * files：用于保存文件名称的vector

 * r：是否须要遍历子文件夹

 * return:文件名称，不包括路径

 */

void list_dir(const char *path,vector<string>& files,bool r = false)

{

	DIR *pDir;

	struct dirent *ent;

	char childpath[512];

	pDir = opendir(path);

	memset(childpath, 0, sizeof(childpath));

	while ((ent = readdir(pDir)) != NULL)

	{

		if (ent->d_type & DT_DIR)

		{

			if (strcmp(ent->d_name, ".") == 0 || strcmp(ent->d_name, "..") == 0)

			{

				continue;

			}

			if(r) //假设须要遍历子文件夹

			{

				sprintf(childpath, "%s/%s", path, ent->d_name);

				list_dir(childpath,files);

			}

		}

		else

		{

			files.push_back(ent->d_name);

		}

	}

	sort(files.begin(),files.end());//排序

}

string get_classname(string path)

{

	int index = path.find_last_of('_');

	return path.substr(0, index);

}

int get_labelid(string fileName)

{

	string class_name_tmp = get_classname(fileName);

	all_class_name.insert(class_name_tmp);

	map<string,int>::iterator name_iter_tmp = class2id.find(class_name_tmp);

	if (name_iter_tmp == class2id.end())

	{

		int id = class2id.size();

		class2id.insert(name_iter_tmp, std::make_pair(class_name_tmp, id));

		return id;

	}

	else

	{

		return name_iter_tmp->second;

	}

}

void loadimg(string path,char* buffer)

{

	cv::Mat img = cv::imread(path, CV_LOAD_IMAGE_COLOR);

	string val;

	int rows = img.rows;

	int cols = img.cols;

	int pos=0;

	for (int c = 0; c < 3; c++)

	{

		for (int row = 0; row < rows; row++)

		{

			for (int col = 0; col < cols; col++)

			{

				buffer[pos++]=img.at<cv::Vec3b>(row,col)[c];

			}

		}

	}

}

void convert(string imgdir,string outputdb,string attachdir,int channel,int width,int height)

{

	leveldb::DB* db;

	leveldb::Options options;

	options.create_if_missing = true;

	options.error_if_exists = true;

	caffe::Datum datum;

	datum.set_channels(channel);

	datum.set_height(height);

	datum.set_width(width);

	int image_size = channel*width*height;

	char buffer[image_size];

	string value;

	CHECK(leveldb::DB::Open(options, outputdb, &db).ok());

	vector<string> filenames;

	list_dir(imgdir.c_str(),filenames);

	string img_log = attachdir+"image_filename";

	ofstream writefile(img_log.c_str());

	for(int i=0;i<filenames.size();i++)

	{

		string path= imgdir;

		path.append(filenames[i]);//算出绝对路径

		loadimg(path,buffer);

		int labelid = get_labelid(filenames[i]);

		datum.add_label(labelid);

		datum.set_data(buffer,image_size);

		datum.SerializeToString(&value);

		snprintf(buffer, image_size, "%05d", i);

		printf("\nclassid:%d classname:%s abspath:%s",labelid,get_classname(filenames[i]).c_str(),path.c_str());

		db->Put(leveldb::WriteOptions(),string(buffer),value);

		//printf("%d %s\n",i,fileNames[i].c_str());

		assert(writefile.is_open());

		writefile<<i<<" "<<filenames[i]<<"\n";

	}

	delete db;

	writefile.close();

	img_log = attachdir+"image_classname";

	writefile.open(img_log.c_str());

	set<string>::iterator iter = all_class_name.begin();

	while(iter != all_class_name.end())

	{

		assert(writefile.is_open());

		writefile<<(*iter)<<"\n";

		//printf("%s\n",(*iter).c_str());

		iter++;

	}

	writefile.close();

}

int main(int argc, char** argv)

{

	if (argc < 6)

	{

	    LOG(ERROR) << "convert_imagedata src_dir dst_dir attach_dir channel width height";

	    return 0;

	}

//./convert_imagedata.bin  /home/linger/imdata/collarTest/ /home/linger/linger/testfile/dbtest/  /home/linger/linger/testfile/test_attachment/ 3 250 250

	//   ./convert_imagedata.bin /home/linger/imdata/collar_train/ /home/linger/linger/testfile/crop_train_db/ /home/linger/linger/testfile/crop_train_attachment/ 3 50 50

	google::InitGoogleLogging(argv[0]);

	string src_dir = argv[1];

	string src_dst = argv[2];

	string attach_dir = argv[3];

	int channel = atoi(argv[4]);

	int width = atoi(argv[5]);

	int height = atoi(argv[6]);

	//for test

	/*

	src_dir = "/home/linger/imdata/collarTest/";

	src_dst = "/home/linger/linger/testfile/dbtest/";

	attach_dir = "/home/linger/linger/testfile/";

	channel = 3;

	width = 250;

	height = 250;

	 */

	convert(src_dir,src_dst,attach_dir,channel,width,height);

}

caffe神经网络框架的辅助工具（将图片转换为leveldb格式）的相关教程结束。