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