Caffe的图像转换工具convert_imageset源码分析

作者

        caffe提供了将图像转换为lmdb或者leveldb格式的工具，在tool文件夹下的convert_imageset中可以找到对应文件，这个工具在很多例子中都有用到，比如mnist和imagenet，由于caffe默认使用lmdb这种又快又小的格式，我们在处理大量图像时也会用到，而我们如果要根据自身需求进行修改，那就需要读源码了。

        这里忍不住扯一些别的，convert_imageset工具不支持浮点数和多标签，而博主的label都是浮点数和多标签，前两天准备全部使用hdf5来存储data和label，后来发现图片存储成hdf5是一件耗时费力占空间的事，于是考虑将标签存储为hdf5，而对UINT8类型的图像使用lmdb进行存储，要存成lmdb自然就会想到使用自带的工具，于是就来分析这个工具了。

        我们从convert_imageset第一行开始，细细的读：

#include <algorithm>
#include <fstream>  // NOLINT(readability/streams)
#include <string>
#include <utility>
#include <vector>

#include "boost/scoped_ptr.hpp"
#include "gflags/gflags.h"
#include "glog/logging.h"

#include "caffe/proto/caffe.pb.h"
#include "caffe/util/db.hpp"
#include "caffe/util/format.hpp"
#include "caffe/util/io.hpp"
#include "caffe/util/rng.hpp"

using namespace caffe;  // NOLINT(build/namespaces)
using std::pair;
using boost::scoped_ptr;

        上面是文件的声明部分，caffe的工具和caffe一样穿插着glog、gflags等等库，不得不承认google的库确实非常好用，在读caffe源码前请务必了解这几个库，顺带加上protobuf：

DEFINE_bool(gray, false,
    "When this option is on, treat images as grayscale ones");
DEFINE_bool(shuffle, false,
    "Randomly shuffle the order of images and their labels");
DEFINE_string(backend, "lmdb",
        "The backend {lmdb, leveldb} for storing the result");
DEFINE_int32(resize_width, 0, "Width images are resized to");
DEFINE_int32(resize_height, 0, "Height images are resized to");
DEFINE_bool(check_size, false,
    "When this option is on, check that all the datum have the same size");
DEFINE_bool(encoded, false,
    "When this option is on, the encoded image will be save in datum");
DEFINE_string(encode_type, "",
    "Optional: What type should we encode the image as ('png','jpg',...).");

int main(int argc, char** argv) {
#ifdef USE_OPENCV
  ::google::InitGoogleLogging(argv[0]);
  // Print output to stderr (while still logging)
  FLAGS_alsologtostderr = 1;

#ifndef GFLAGS_GFLAGS_H_
  namespace gflags = google;
#endif

  gflags::SetUsageMessage("Convert a set of images to the leveldb/lmdb\n"
        "format used as input for Caffe.\n"
        "Usage:\n"
        "    convert_imageset [FLAGS] ROOTFOLDER/ LISTFILE DB_NAME\n"
        "The ImageNet dataset for the training demo is at\n"
        "    http://www.image-net.org/download-images\n");
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  if (argc < 4) {
    gflags::ShowUsageWithFlagsRestrict(argv[0], "tools/convert_imageset");
    return 1;
  }

        上面这部分源码就是在解析参数，主要是在使用gflags进行操作，在最开始通过宏的方式定义了几个全局参数，这些参数都会在命令行或者脚本传入参数的时候被赋值，这些变量都会被加上前缀FLAGS_，下面是一段工具使用时的命令参数，来感受一下：

GLOG_logtostderr=1 $TOOLS/convert_imageset \
    --resize_height=$RESIZE_HEIGHT \
    --resize_width=$RESIZE_WIDTH \
    --shuffle \
    $TRAIN_DATA_ROOT \
    $DATA/train.txt \
    $EXAMPLE/ilsvrc12_train_lmdb

        gflags就是这样通过脚本或者命令行给相应变量赋值的工具，进入主函数之后的代码也是在做解析命令。我们继续往下读：

  const bool is_color = !FLAGS_gray;
  const bool check_size = FLAGS_check_size;
  const bool encoded = FLAGS_encoded;
  const string encode_type = FLAGS_encode_type;

  std::ifstream infile(argv[2]);
  std::vector<std::pair<std::string, int> > lines;
  std::string filename;
  int label;
  while (infile >> filename >> label) {
    lines.push_back(std::make_pair(filename, label));
  }

        好了，这里我们就来好好分析几个参数：

        gray              用来标识输入的是否是灰度图像

        check_size  用来标识是否检查数据

        encoded       用来标识编码，而encoder_type则用来标识编码的类型，好像默认是jpg，如果你的图像是其它格式，你需要对这个参数进行修改。

        shuffele         用来标识是否被洗牌（序列随机化）

        backend       用来标识数据类型，默认为lmdb

        resize参数    用来自动将图像进行缩放，缩放至同样大小，因为caffe接收的图像一定是同样大小的（如果你没有修改caffe源码的话）

        在上面那段代码最后的while循环，就是在取出图像的名字和标签，这也是这个工具不支持多标签的原因，你可以从这里开始修改源码，让这个工具来支持多标签（你还需要修改caffe的源码，这部分内容请点击菜单栏的caffe笔记，我把资料整理在那里）

  if (FLAGS_shuffle) {
    // randomly shuffle data
    LOG(INFO) << "Shuffling data";
    shuffle(lines.begin(), lines.end());
  }
  LOG(INFO) << "A total of " << lines.size() << " images.";

  if (encode_type.size() && !encoded)
    LOG(INFO) << "encode_type specified, assuming encoded=true.";

  int resize_height = std::max<int>(0, FLAGS_resize_height);
  int resize_width = std::max<int>(0, FLAGS_resize_width);

        上面这段源码也就比较好懂了，根据输入参数做相应的洗牌操作等。

  // Create new DB
  scoped_ptr<db::DB> db(db::GetDB(FLAGS_backend));
  db->Open(argv[3], db::NEW);
  scoped_ptr<db::Transaction> txn(db->NewTransaction());

  // Storing to db
  std::string root_folder(argv[1]);
  Datum datum;
  int count = 0;
  int data_size = 0;
  bool data_size_initialized = false;

  for (int line_id = 0; line_id < lines.size(); ++line_id) {
    bool status;
    std::string enc = encode_type;
    if (encoded && !enc.size()) {
      // Guess the encoding type from the file name
      string fn = lines[line_id].first;
      size_t p = fn.rfind('.');
      if ( p == fn.npos )
        LOG(WARNING) << "Failed to guess the encoding of '" << fn << "'";
      enc = fn.substr(p);
      std::transform(enc.begin(), enc.end(), enc.begin(), ::tolower);
    }
    status = ReadImageToDatum(root_folder + lines[line_id].first,
        lines[line_id].second, resize_height, resize_width, is_color,
        enc, &datum);
    if (status == false) continue;
    if (check_size) {
      if (!data_size_initialized) {
        data_size = datum.channels() * datum.height() * datum.width();
        data_size_initialized = true;
      } else {
        const std::string& data = datum.data();
        CHECK_EQ(data.size(), data_size) << "Incorrect data field size "
            << data.size();
      }
    }
    // sequential
    string key_str = caffe::format_int(line_id, 8) + "_" + lines[line_id].first;

    // Put in db
    string out;
    CHECK(datum.SerializeToString(&out));
    txn->Put(key_str, out);

    if (++count % 1000 == 0) {
      // Commit db
      txn->Commit();
      txn.reset(db->NewTransaction());
      LOG(INFO) << "Processed " << count << " files.";
    }
  }

        根据要存储的类型创建相应的数据库（lmdb或者leveldb），这里的核心就是这个最外面的大for循环了，这里根据文件名读取图像并缩放到同样大小，根据需要检查数据，然后将其写入数据库，每1000副图像一个Transaction，而最后一批次可能不到1000，在for循环外有：

// write the last batch
  if (count % 1000 != 0) {
    txn->Commit();
    LOG(INFO) << "Processed " << count << " files.";
  }

        这样我们就明白这个工具的工作流程了。

        See You Next Chapter!

分享到微信

打开微信“扫一扫”，打开网页后点击屏幕右上角分享按钮
分享到微博

《Caffe的图像转换工具convert_imageset源码分析》有2条评论

发表评论