接觸了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