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); }