Protobuf Python 示例

接触了Protobuf之后就迷上它了。虽然一开始觉得有点麻烦，但是上手之后就觉得比XML好用太多了。以至于现在不管啥都想用Protobuf。之前都是用的Protobuf C++，现在挑个简单的任务在Python上练练手吧~

背景非常简单，我在做身份证号码的OCR，现在经过切分和人工标记，已经有了这样的一个训练集：

分类文件夹

现在，我希望把这个训练集中的内容写成prototxt格式。proto的定义如下：

// [START declaration]
syntax = "proto3";
package ocr;
// [END declaration]

// [START messages]
message DataPair {
  string data_path = 1;
  int32 data_label = 2;
}

message DataPairList {
  repeated DataPair data_pair = 1;
}
// [END messages]

其实非常简单，就是一个数据路径加一个数据标签。那我们现在就来看看如何来写一个python程序，将文件夹中的数据读出，写成prototxt格式吧：

Step 1: 生成Datapair_pb2.py

生成的方式是利用protobuf给的生成工具，protoc，格式为：

protoc -I=$SRC_DIR --python_out=$DST_DIR $SRC_DIR/$PROTO_NAME

下面是我的例子：

protoc -I=E:\Pictures\ImageDataBase\IDCard\train\id-num\for-train --python_out=D:\PycharmProjects\FileDirOp E:\Pictures\ImageDataBase\IDCard\train\id-num\for-train\Datapair.proto

这样，就在D:\PycharmProjects\FileDirOp生成了一个Datapair_pb2.py文件。

Step 2: 安装Python的Protobuf包

如果没有安装包，会发现Datapair_pb2.py文件中

from google.protobuf import descriptor as _descriptor
from google.protobuf import message as _message
from google.protobuf import reflection as _reflection
from google.protobuf import symbol_database as _symbol_database

是无法使用的。如此，则需要安装protobuf。安装方法如下

> cd $PROTOBUF_DIR$
> cd python
> python setup.py build
> python setup.py test
> python setup.py install

Step3: 撰写Python代码

代码如下：

#!/usr/bin/python
#  -*- coding:utf-8 -*-

import os
import Datapair_pb2
from os.path import join, splitext, abspath
import google.protobuf


def read_data(dir_name):
    pairs_list = []
    cwd = os.getcwd()
    os.chdir(dir_name)
    subdir_list = [d for d in os.listdir('.') if os.path.isdir(d)]
    print 'classes:', subdir_list
    label = 0
    for d in subdir_list:
        data_paths = [abspath(join(d, f)) for f in os.listdir(d) if splitext(f)[-1] == '.bmp']
        pairs = [(path, label) for path in data_paths]
        for p in pairs:
            pairs_list.append(p)
        label += 1
    os.chdir(cwd)
    return pairs_list


def write_pairs_to_prototxt(pairs, filename):
    data_pair_list = Datapair_pb2.DataPairList()
    for p in pairs:
        data_pair = data_pair_list.data_pair.add()
        data_pair.data_path = p[0]
        data_pair.data_label = p[1]
    # print data_pair_list
    with open(filename, 'w') as f:
        f.write(str(data_pair_list))


if __name__ == "__main__":
    print 'The root directory of train set is?'
    root_dir_name = raw_input('(dir name): ')
    while not os.path.isdir(root_dir_name):
        print 'The directory you entered is not a directory, please check and retype'
        root_dir_name = raw_input('(dir name): ')
    data_pairs = read_data(root_dir_name)
    print 'Total number of data pairs',len(data_pairs)

    print 'The name of .prototxt is :'
    prototxt_name = raw_input('(default dir: ' + root_dir_name+ ') > ')
    if len(prototxt_name) < 10 or prototxt_name[-9:] != ".prototxt":
        prototxt_name += '.prototxt'
    if not os.path.isabs(prototxt_name):
        prototxt_name = join(root_dir_name, prototxt_name)
    write_pairs_to_prototxt(data_pairs, prototxt_name)

运行：

The root directory of train set is?
(dir name): E:\Pictures\ImageDataBase\IDCard\train\id-num\for-train
classes: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'X']
Total number of data pairs 220
The name of .prototxt is :
(default dir: E:\Pictures\ImageDataBase\IDCard\train\id-num\for-train) > id_num_train_data
 
Process finished with exit code 0

生成文件没有问题。

虽然不确定

with open(filename, 'w') as f:
    f.write(str(data_pair_list))

用这样一句话来生成.prototxt文件是否正确，但至少，生成的文件是可用的。与C++相比，不禁惊叹Python的简洁性！

参考： https://developers.google.com/protocol-buffers/docs/pythontutorial