|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
"""Tests for oid_tfrecord_creation.py.""" |
|
|
|
import pandas as pd |
|
import tensorflow as tf |
|
|
|
from object_detection.dataset_tools import oid_tfrecord_creation |
|
|
|
|
|
def create_test_data(): |
|
data = { |
|
'ImageID': ['i1', 'i1', 'i1', 'i1', 'i1', 'i2', 'i2'], |
|
'LabelName': ['a', 'a', 'b', 'b', 'c', 'b', 'c'], |
|
'YMin': [0.3, 0.6, 0.8, 0.1, None, 0.0, 0.0], |
|
'XMin': [0.1, 0.3, 0.7, 0.0, None, 0.1, 0.1], |
|
'XMax': [0.2, 0.3, 0.8, 0.5, None, 0.9, 0.9], |
|
'YMax': [0.3, 0.6, 1, 0.8, None, 0.8, 0.8], |
|
'IsOccluded': [0, 1, 1, 0, None, 0, 0], |
|
'IsTruncated': [0, 0, 0, 1, None, 0, 0], |
|
'IsGroupOf': [0, 0, 0, 0, None, 0, 1], |
|
'IsDepiction': [1, 0, 0, 0, None, 0, 0], |
|
'ConfidenceImageLabel': [None, None, None, None, 0, None, None], |
|
} |
|
df = pd.DataFrame(data=data) |
|
label_map = {'a': 0, 'b': 1, 'c': 2} |
|
return label_map, df |
|
|
|
|
|
class TfExampleFromAnnotationsDataFrameTests(tf.test.TestCase): |
|
|
|
def test_simple(self): |
|
label_map, df = create_test_data() |
|
|
|
tf_example = oid_tfrecord_creation.tf_example_from_annotations_data_frame( |
|
df[df.ImageID == 'i1'], label_map, 'encoded_image_test') |
|
self.assertProtoEquals( |
|
""" |
|
features { |
|
feature { |
|
key: "image/encoded" |
|
value { bytes_list { value: "encoded_image_test" } } } |
|
feature { |
|
key: "image/filename" |
|
value { bytes_list { value: "i1.jpg" } } } |
|
feature { |
|
key: "image/object/bbox/ymin" |
|
value { float_list { value: [0.3, 0.6, 0.8, 0.1] } } } |
|
feature { |
|
key: "image/object/bbox/xmin" |
|
value { float_list { value: [0.1, 0.3, 0.7, 0.0] } } } |
|
feature { |
|
key: "image/object/bbox/ymax" |
|
value { float_list { value: [0.3, 0.6, 1.0, 0.8] } } } |
|
feature { |
|
key: "image/object/bbox/xmax" |
|
value { float_list { value: [0.2, 0.3, 0.8, 0.5] } } } |
|
feature { |
|
key: "image/object/class/label" |
|
value { int64_list { value: [0, 0, 1, 1] } } } |
|
feature { |
|
key: "image/object/class/text" |
|
value { bytes_list { value: ["a", "a", "b", "b"] } } } |
|
feature { |
|
key: "image/source_id" |
|
value { bytes_list { value: "i1" } } } |
|
feature { |
|
key: "image/object/depiction" |
|
value { int64_list { value: [1, 0, 0, 0] } } } |
|
feature { |
|
key: "image/object/group_of" |
|
value { int64_list { value: [0, 0, 0, 0] } } } |
|
feature { |
|
key: "image/object/occluded" |
|
value { int64_list { value: [0, 1, 1, 0] } } } |
|
feature { |
|
key: "image/object/truncated" |
|
value { int64_list { value: [0, 0, 0, 1] } } } |
|
feature { |
|
key: "image/class/label" |
|
value { int64_list { value: [2] } } } |
|
feature { |
|
key: "image/class/text" |
|
value { bytes_list { value: ["c"] } } } } |
|
""", tf_example) |
|
|
|
def test_no_attributes(self): |
|
label_map, df = create_test_data() |
|
|
|
del df['IsDepiction'] |
|
del df['IsGroupOf'] |
|
del df['IsOccluded'] |
|
del df['IsTruncated'] |
|
del df['ConfidenceImageLabel'] |
|
|
|
tf_example = oid_tfrecord_creation.tf_example_from_annotations_data_frame( |
|
df[df.ImageID == 'i2'], label_map, 'encoded_image_test') |
|
self.assertProtoEquals(""" |
|
features { |
|
feature { |
|
key: "image/encoded" |
|
value { bytes_list { value: "encoded_image_test" } } } |
|
feature { |
|
key: "image/filename" |
|
value { bytes_list { value: "i2.jpg" } } } |
|
feature { |
|
key: "image/object/bbox/ymin" |
|
value { float_list { value: [0.0, 0.0] } } } |
|
feature { |
|
key: "image/object/bbox/xmin" |
|
value { float_list { value: [0.1, 0.1] } } } |
|
feature { |
|
key: "image/object/bbox/ymax" |
|
value { float_list { value: [0.8, 0.8] } } } |
|
feature { |
|
key: "image/object/bbox/xmax" |
|
value { float_list { value: [0.9, 0.9] } } } |
|
feature { |
|
key: "image/object/class/label" |
|
value { int64_list { value: [1, 2] } } } |
|
feature { |
|
key: "image/object/class/text" |
|
value { bytes_list { value: ["b", "c"] } } } |
|
feature { |
|
key: "image/source_id" |
|
value { bytes_list { value: "i2" } } } } |
|
""", tf_example) |
|
|
|
def test_label_filtering(self): |
|
label_map, df = create_test_data() |
|
|
|
label_map = {'a': 0} |
|
|
|
tf_example = oid_tfrecord_creation.tf_example_from_annotations_data_frame( |
|
df[df.ImageID == 'i1'], label_map, 'encoded_image_test') |
|
self.assertProtoEquals( |
|
""" |
|
features { |
|
feature { |
|
key: "image/encoded" |
|
value { bytes_list { value: "encoded_image_test" } } } |
|
feature { |
|
key: "image/filename" |
|
value { bytes_list { value: "i1.jpg" } } } |
|
feature { |
|
key: "image/object/bbox/ymin" |
|
value { float_list { value: [0.3, 0.6] } } } |
|
feature { |
|
key: "image/object/bbox/xmin" |
|
value { float_list { value: [0.1, 0.3] } } } |
|
feature { |
|
key: "image/object/bbox/ymax" |
|
value { float_list { value: [0.3, 0.6] } } } |
|
feature { |
|
key: "image/object/bbox/xmax" |
|
value { float_list { value: [0.2, 0.3] } } } |
|
feature { |
|
key: "image/object/class/label" |
|
value { int64_list { value: [0, 0] } } } |
|
feature { |
|
key: "image/object/class/text" |
|
value { bytes_list { value: ["a", "a"] } } } |
|
feature { |
|
key: "image/source_id" |
|
value { bytes_list { value: "i1" } } } |
|
feature { |
|
key: "image/object/depiction" |
|
value { int64_list { value: [1, 0] } } } |
|
feature { |
|
key: "image/object/group_of" |
|
value { int64_list { value: [0, 0] } } } |
|
feature { |
|
key: "image/object/occluded" |
|
value { int64_list { value: [0, 1] } } } |
|
feature { |
|
key: "image/object/truncated" |
|
value { int64_list { value: [0, 0] } } } |
|
feature { |
|
key: "image/class/label" |
|
value { int64_list { } } } |
|
feature { |
|
key: "image/class/text" |
|
value { bytes_list { } } } } |
|
""", tf_example) |
|
|
|
|
|
if __name__ == '__main__': |
|
tf.test.main() |
|
|