Add image sha256 to tfrecord (#1931)

main
zhiltsov-max 6 years ago committed by GitHub
parent f9fab2459d
commit a3448a20ad
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -5,6 +5,7 @@
import codecs import codecs
from collections import OrderedDict from collections import OrderedDict
import hashlib
import logging as log import logging as log
import os import os
import os.path as osp import os.path as osp
@ -180,15 +181,18 @@ class TfDetectionApiConverter(Converter):
features.update({ features.update({
'image/encoded': bytes_feature(b''), 'image/encoded': bytes_feature(b''),
'image/format': bytes_feature(b'') 'image/format': bytes_feature(b''),
'image/key/sha256': bytes_feature(b''),
}) })
if self._save_images: if self._save_images:
if item.has_image and item.image.has_data: if item.has_image and item.image.has_data:
buffer, fmt = self._save_image(item, filename) buffer, fmt = self._save_image(item, filename)
key = hashlib.sha256(buffer).hexdigest()
features.update({ features.update({
'image/encoded': bytes_feature(buffer), 'image/encoded': bytes_feature(buffer),
'image/format': bytes_feature(fmt.encode('utf-8')), 'image/format': bytes_feature(fmt.encode('utf-8')),
'image/key/sha256': bytes_feature(key.encode('utf8')),
}) })
else: else:
log.warning("Item '%s' has no image" % item.id) log.warning("Item '%s' has no image" % item.id)

@ -85,6 +85,10 @@ class TfDetectionApiExtractor(SourceExtractor):
'image/width': tf.io.FixedLenFeature([], tf.int64), 'image/width': tf.io.FixedLenFeature([], tf.int64),
'image/encoded': tf.io.FixedLenFeature([], tf.string), 'image/encoded': tf.io.FixedLenFeature([], tf.string),
'image/format': tf.io.FixedLenFeature([], tf.string), 'image/format': tf.io.FixedLenFeature([], tf.string),
# use varlen to avoid errors when this field is missing
'image/key/sha256': tf.io.VarLenFeature(tf.string),
# Object boxes and classes. # Object boxes and classes.
'image/object/bbox/xmin': tf.io.VarLenFeature(tf.float32), 'image/object/bbox/xmin': tf.io.VarLenFeature(tf.float32),
'image/object/bbox/xmax': tf.io.VarLenFeature(tf.float32), 'image/object/bbox/xmax': tf.io.VarLenFeature(tf.float32),

Loading…
Cancel
Save