TensorFlow 简单 pb (pbtxt) 文件读写

        在使用 TensorFlow 的时候,很多地方都会遇到 protobuf (.proto) 文件,比如配置 TensorFlow Detection API 的过程中需要执行如下语句:

$ protoc object_detection/protos/*.proto --python_out=.

表示把文件夹 protos 下的所有 .proto 文件转化对应的 .py 文件(配置过程见 TensorFlow 训练自己的目标检测器)。 之后,再借助这些转化来的 .py 文件就可以读取特定格式的 .pbtxt, .config 等文件了。比如,可以使用 string_int_label_map.proto 转化来的 string_int_label_map_pb2.py 文件来读取目标检测与实例分割 的类名与类标号配置文件(这种文件以 .pbtxt 作为后缀名,假设要检测 person 和 car 等类目标):

item {
    id: 1
    name: 'person'
}

item {
    id: 2
    name: 'car'
}

...

        Protocol Buffer(protobuf) 是谷歌开源的一种数据存储语言,每一个文件以 .proto 为后缀,它不依赖特定的语言与平台,扩展性极强。TensorFlow 内部的数据存储(比如 .pb.ckpt 等)基本都使用 protobuf 格式。

        下面以一个简单的例子来说明 protobuf 文件的读取。

1.protobuf 数据结构定义

        假设我们的目的是读取如下文件(命名为 students.pbtxt,使用文本编辑器编辑):

student_info {
    name: 'Zhang San';
    age: 20;
    sex: 0;
}

student_info {
    name: 'Li Si';
    age: 25;
    sex: 0;
}

student_info {
    name: 'Wang Wu';
    age: 18;
    sex: 1;
}

显然,这是一份简单的结构化数据,但若使用传统的数据读取方式且要快速方便的解析出其中的学生信息,却不容易。此时,如果使用 protobuf 则相当便捷。

        首先,用文本编辑器编辑一个 .proto 文件(命名为:student_info.proto):

syntax = "proto3";

package proto_test;

message Student {
    string name = 1;
    int32 age = 2;
    int32 sex = 3;        
}

message StudentInfo {
    repeated Student student_info = 1;        
}

其中的 syntax 指定使用的 protobuf 版本,可以填写 proto2 (protobuf 2)和 proto3 (protobuf 3),这里使用的是后者。下面的 package 指定 student_info.proto 文件所在的文件夹名字。接下来,以关键字 message 开头定义了一个简单的数据结构 Student,里面包括三个可选字段 name,age 和 sex,后面的数字 1,2,3 指定这三个字段在编码序列化后的二进制数据中的顺序,因此在同一个 message 内部,它们是不允许重复的。如果是 protobuf 2 的版本,需要在可选字段前面加上 optional 关键字(关键字包括 optionalrequiredrepeated)。定义好 Student 结构之后, students.pbtxt 文件的内容基本是重复这个结构,因此还需要定义一个新的 message StudentInfo,它包含一个可重复的 (repeated) 字段 student_info。至此,数据解析格式定义完了,接下来要将它转化为 Python 格式语言,执行:

$ protoc proto_test/*.proto --python_out=.

会自动生成一个 student_info_pb2.py 文件,前几行如下:

# Generated by the protocol buffer compiler.  DO NOT EDIT!
# source: proto_test/student_info.proto

import sys
_b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1'))
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
# @@protoc_insertion_point(imports)

_sym_db = _symbol_database.Default()




DESCRIPTOR = _descriptor.FileDescriptor(
  name='proto_test/student_info.proto',
  package='proto_test',
  syntax='proto3',
...

基于这个 student_info_pb2.py 就可以方便的解析 students.pbtxt 中的内容了。

2.读取 .pbtxt 文件

        有了转化来的 student_info_pb2.py 文件,解析 students.pbtxt 就轻而易举了,代码如下(命名为:read_student_info.py):

# -*- coding: utf-8 -*-
"""
Created on Sat Sep  8 19:18:37 2018

@author: shirhe-lyh
"""

import tensorflow as tf

from google.protobuf import text_format

import student_info_pb2


def load_pbtxt_file(path):
    """Read .pbtxt file.
    
    Args: 
        path: Path to StringIntLabelMap proto text file (.pbtxt file).
        
    Returns:
        A StringIntLabelMapProto.
        
    Raises:
        ValueError: If path is not exist.
    """
    if not tf.gfile.Exists(path):
        raise ValueError('`path` is not exist.')
        
    with tf.gfile.GFile(path, 'r') as fid:
        pbtxt_string = fid.read()
        pbtxt = student_info_pb2.StudentInfo()
        try:
            text_format.Merge(pbtxt_string, pbtxt)
        except text_format.ParseError:
            pbtxt.ParseFromString(pbtxt_string)
    return pbtxt


def get_student_info_dict(path):
    """Reads a .pbtxt file and returns a dictionary.
    
    Args:
        path: Path to StringIntLabelMap proto text file.
        
    Returns:
        A dictionary mapping class names to indices.
    """
    pbtxt = load_pbtxt_file(path)
    
    result_dict = {}
    for student_info in pbtxt.student_info:
        result_dict[student_info.name] = [student_info.age, student_info.sex]
    return result_dict

首先, 使用 tf.gfile.GFile 将指定文件读入成字符串,然后定义一个 student_info_pb2.StudentInfo() 结构,这样使用 google.protobuftext_format.Merge 直接将 student_info 一个一个解析成 Student 结构,此时读取 name、age 和 sex 字段只需要通过 . 属性即可。

        如开头所言,我们来读取 students.pbtxt 文件:

student_info {
    name: 'Zhang San';
    age: 20;
    sex: 0;
}

student_info {
    name: 'Li Si';
    age: 25;
    sex: 0;
}

student_info {
    name: 'Wang Wu';
    age: 18;
    sex: 1;
}

读取代码如下:

# -*- coding: utf-8 -*-
"""
Created on Sat Sep  8 19:26:54 2018

@author: shirhe-lyh
"""

import read_student_info


if __name__ == '__main__':
    student_info_path = './students.pbtxt'
    students_dict = read_student_info.get_student_info_dict(student_info_path)
    print(students_dict)

执行后输出:

{'Li Si': [25, 0], 'Wang Wu': [18, 1], 'Zhang San': [20, 0]}
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 216,240评论 6 498
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 92,328评论 3 392
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 162,182评论 0 353
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 58,121评论 1 292
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 67,135评论 6 388
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 51,093评论 1 295
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 40,013评论 3 417
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,854评论 0 273
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 45,295评论 1 310
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,513评论 2 332
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,678评论 1 348
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 35,398评论 5 343
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,989评论 3 325
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,636评论 0 22
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,801评论 1 268
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,657评论 2 368
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,558评论 2 352

推荐阅读更多精彩内容