如何可视化TFRecord文件？附实现代码与校验要点

我之前在别的论坛被问到过这个问题，觉得很有必要分享给所有在TFRecords上踩坑的朋友。

常见问题1：标签不匹配 #

如果TFRecord文件里的标签和labels.pbtxt文件中的标签对不上，TensorFlow Object Detection API可能会出现奇怪的表现——程序能正常运行，损失值也会下降，但就是输出不了有效的检测结果。这种隐性问题特别容易被忽略，一定要仔细核对两者的标签ID和对应名称。

常见问题2：坐标系混淆 #

我自己老是搞混X-Y坐标系和行列空间，所以每次都会反复检查标注是否对应图像的正确区域。

验证标注的实用代码 #

我发现最靠谱的方式是用TF工具解码TFRecord并可视化标注，这样能直观确认标注是否正确。下面是我常用的代码：

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from object_detection.utils import visualization_utils as vu
from object_detection.protos import string_int_label_map_pb2 as pb
from object_detection.data_decoders.tf_example_decoder import TfExampleDecoder as TfDecoder
from google.protobuf import text_format

def main(tfrecords_filename, label_map=None):
    if label_map is not None:
        label_map_proto = pb.StringIntLabelMap()
        with tf.gfile.GFile(label_map,'r') as f:
            text_format.Merge(f.read(), label_map_proto)
        class_dict = {}
        for entry in label_map_proto.item:
            class_dict[entry.id] = {'name':entry.display_name}
    
    sess = tf.Session()
    decoder = TfDecoder(label_map_proto_file=label_map, use_display_name=False)
    sess.run(tf.tables_initializer())
    
    for record in tf.python_io.tf_record_iterator(tfrecords_filename):
        example = decoder.decode(record)
        host_example = sess.run(example)
        scores = np.ones(host_example['groundtruth_boxes'].shape[0])
        
        vu.visualize_boxes_and_labels_on_image_array(
            host_example['image'],
            host_example['groundtruth_boxes'],
            host_example['groundtruth_classes'],
            scores,
            class_dict,
            max_boxes_to_draw=None,
            use_normalized_coordinates=True
        )
        
        plt.imshow(host_example['image'])
        plt.show()

内容的提问来源于stack exchange，提问作者Steve Goley