Wangyida
9 years ago
parent
9fdc78a302
commit
391596a1b2
10 changed files with 564 additions and 149 deletions
Binary file not shown.
Binary file not shown.
@ -0,0 +1,202 @@ |
||||
#include <stdio.h> // for snprintf |
||||
#include <string> |
||||
#include <vector> |
||||
|
||||
#include "boost/algorithm/string.hpp" |
||||
#include "google/protobuf/text_format.h" |
||||
#define CPU_ONLY |
||||
#include "caffe/blob.hpp" |
||||
#include "caffe/common.hpp" |
||||
#include "caffe/net.hpp" |
||||
#include "caffe/proto/caffe.pb.h" |
||||
#include "caffe/util/db.hpp" |
||||
#include "caffe/util/io.hpp" |
||||
#include "caffe/vision_layers.hpp" |
||||
|
||||
using caffe::Blob; |
||||
using caffe::Caffe; |
||||
using caffe::Datum; |
||||
using caffe::Net; |
||||
using boost::shared_ptr; |
||||
using std::string; |
||||
namespace db = caffe::db; |
||||
|
||||
template<typename Dtype> |
||||
int feature_extraction_pipeline(int argc, char** argv); |
||||
int main(int argc, char** argv) { |
||||
return feature_extraction_pipeline<float>(argc, argv); |
||||
// return feature_extraction_pipeline<double>(argc, argv);
|
||||
} |
||||
|
||||
template<typename Dtype> |
||||
int feature_extraction_pipeline(int argc, char** argv) { |
||||
::google::InitGoogleLogging(argv[0]); |
||||
const int num_required_args = 7; |
||||
/*if (argc < num_required_args) {
|
||||
LOG(ERROR)<< |
||||
"This program takes in a trained network and an input data layer, and then" |
||||
" extract features of the input data produced by the net.\n" |
||||
"Usage: extract_features pretrained_net_param" |
||||
" feature_extraction_proto_file extract_feature_blob_name1[,name2,...]" |
||||
" save_feature_dataset_name1[,name2,...] num_mini_batches db_type" |
||||
" [CPU/GPU] [DEVICE_ID=0]\n" |
||||
"Note: you can extract multiple features in one pass by specifying" |
||||
" multiple feature blob names and dataset names seperated by ','." |
||||
" The names cannot contain white space characters and the number of blobs" |
||||
" and datasets must be equal."; |
||||
return 1; |
||||
}*/ |
||||
int arg_pos = num_required_args; |
||||
|
||||
arg_pos = num_required_args; |
||||
//if (argc > arg_pos && strcmp(argv[arg_pos], "GPU") == 0) {
|
||||
if (argc > arg_pos && strcmp("CPU", "GPU") == 0) { |
||||
LOG(ERROR)<< "Using GPU"; |
||||
int device_id = 0; |
||||
if (argc > arg_pos + 1) { |
||||
device_id = atoi(argv[arg_pos + 1]); |
||||
CHECK_GE(device_id, 0); |
||||
} |
||||
LOG(ERROR) << "Using Device_id=" << device_id; |
||||
Caffe::SetDevice(device_id); |
||||
Caffe::set_mode(Caffe::GPU); |
||||
} else { |
||||
LOG(ERROR) << "Using CPU"; |
||||
Caffe::set_mode(Caffe::CPU); |
||||
} |
||||
|
||||
arg_pos = 0; // the name of the executable
|
||||
//std::string pretrained_binary_proto(argv[++arg_pos]);
|
||||
std::string pretrained_binary_proto("/home/wangyida/Desktop/caffe/examples/triplet/3d_triplet_iter_200.caffemodel"); |
||||
|
||||
// Expected prototxt contains at least one data layer such as
|
||||
// the layer data_layer_name and one feature blob such as the
|
||||
// fc7 top blob to extract features.
|
||||
/*
|
||||
layers { |
||||
name: "data_layer_name" |
||||
type: DATA |
||||
data_param { |
||||
source: "/path/to/your/images/to/extract/feature/images_leveldb" |
||||
mean_file: "/path/to/your/image_mean.binaryproto" |
||||
batch_size: 128 |
||||
crop_size: 227 |
||||
mirror: false |
||||
} |
||||
top: "data_blob_name" |
||||
top: "label_blob_name" |
||||
} |
||||
layers { |
||||
name: "drop7" |
||||
type: DROPOUT |
||||
dropout_param { |
||||
dropout_ratio: 0.5 |
||||
} |
||||
bottom: "fc7" |
||||
top: "fc7" |
||||
} |
||||
*/ |
||||
//std::string feature_extraction_proto(argv[++arg_pos]);
|
||||
std::string feature_extraction_proto("/home/wangyida/Desktop/caffe/examples/triplet/3d_triplet.prototxt"); |
||||
boost::shared_ptr<Net<Dtype> > feature_extraction_net( |
||||
new Net<Dtype>(feature_extraction_proto, caffe::TEST)); |
||||
feature_extraction_net->CopyTrainedLayersFrom(pretrained_binary_proto); |
||||
|
||||
//std::string extract_feature_blob_names(argv[++arg_pos]);
|
||||
std::string extract_feature_blob_names("feat"); |
||||
std::vector<std::string> blob_names; |
||||
boost::split(blob_names, extract_feature_blob_names, boost::is_any_of(",")); |
||||
|
||||
//std::string save_feature_dataset_names(argv[++arg_pos]);
|
||||
std::string save_feature_dataset_names("/home/wangyida/Desktop/caffe/examples/triplet/feature_extracted"); |
||||
std::vector<std::string> dataset_names; |
||||
boost::split(dataset_names, save_feature_dataset_names, |
||||
boost::is_any_of(",")); |
||||
CHECK_EQ(blob_names.size(), dataset_names.size()) << |
||||
" the number of blob names and dataset names must be equal"; |
||||
size_t num_features = blob_names.size(); |
||||
|
||||
for (size_t i = 0; i < num_features; i++) { |
||||
CHECK(feature_extraction_net->has_blob(blob_names[i])) |
||||
<< "Unknown feature blob name " << blob_names[i] |
||||
<< " in the network " << feature_extraction_proto; |
||||
} |
||||
|
||||
//int num_mini_batches = atoi(argv[++arg_pos]);
|
||||
int num_mini_batches = atoi("6"); |
||||
|
||||
/*std::vector<shared_ptr<db::DB> > feature_dbs;
|
||||
std::vector<shared_ptr<db::Transaction> > txns; |
||||
const char* db_type = argv[++arg_pos]; |
||||
for (size_t i = 0; i < num_features; ++i) { |
||||
LOG(INFO)<< "Opening dataset " << dataset_names[i]; |
||||
shared_ptr<db::DB> db(db::GetDB(db_type)); |
||||
db->Open(dataset_names.at(i), db::NEW); |
||||
feature_dbs.push_back(db); |
||||
shared_ptr<db::Transaction> txn(db->NewTransaction()); |
||||
txns.push_back(txn); |
||||
}*/ |
||||
std::vector<FILE*> files; |
||||
for (size_t i = 0; i < num_features; ++i) |
||||
{ |
||||
LOG(INFO) << "Opening file " << dataset_names[i]; |
||||
FILE * temp = NULL; |
||||
temp = fopen(dataset_names[i].c_str(), "wb"); |
||||
files.push_back(temp); |
||||
} |
||||
|
||||
|
||||
LOG(ERROR)<< "Extacting Features"; |
||||
|
||||
//bool header_flag = true;
|
||||
|
||||
Datum datum; |
||||
//const int kMaxKeyStrLength = 100;
|
||||
//char key_str[kMaxKeyStrLength];
|
||||
std::vector<Blob<float>*> input_vec; |
||||
std::vector<int> image_indices(num_features, 0); |
||||
for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index) { |
||||
feature_extraction_net->Forward(input_vec); |
||||
for (int i = 0; i < num_features; ++i) { |
||||
const boost::shared_ptr<Blob<Dtype> > feature_blob = feature_extraction_net |
||||
->blob_by_name(blob_names[i]); |
||||
int batch_size = feature_blob->num(); |
||||
int dim_features = feature_blob->count() / batch_size; |
||||
if (batch_index == 0) |
||||
{ |
||||
int fea_num = batch_size*num_mini_batches; |
||||
fwrite(&dim_features, sizeof(int), 1, files[i]); |
||||
fwrite(&fea_num, sizeof(int), 1, files[i]); |
||||
//bool header_flag = false;
|
||||
} |
||||
const Dtype* feature_blob_data; |
||||
for (int n = 0; n < batch_size; ++n) { |
||||
|
||||
feature_blob_data = feature_blob->cpu_data() + |
||||
feature_blob->offset(n); |
||||
fwrite(feature_blob_data, sizeof(Dtype), dim_features, files[i]); |
||||
|
||||
++image_indices[i]; |
||||
if (image_indices[i] % 1000 == 0) { |
||||
LOG(ERROR)<< "Extracted features of " << image_indices[i] << |
||||
" query images for feature blob " << blob_names[i]; |
||||
} |
||||
} // for (int n = 0; n < batch_size; ++n)
|
||||
} // for (int i = 0; i < num_features; ++i)
|
||||
} // for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index)
|
||||
// write the last batch
|
||||
for (int i = 0; i < num_features; ++i) { |
||||
/* if (image_indices[i] % 1000 != 0) {
|
||||
txns.at(i)->Commit(); |
||||
} |
||||
LOG(ERROR)<< "Extracted features of " << image_indices[i] << |
||||
" query images for feature blob " << blob_names[i]; |
||||
feature_dbs.at(i)->Close();*/ |
||||
fclose(files[i]); |
||||
} |
||||
|
||||
LOG(ERROR)<< "Successfully extracted the features!"; |
||||
return 0; |
||||
} |
||||
|
||||
|
||||
@ -0,0 +1,191 @@ |
||||
#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; |
||||
|
||||
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.set_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); |
||||
|
||||
|
||||
|
||||
} |
||||
|
||||
Loading…
Reference in new issue