object_detection_dataset.py 文件源码

def read_record(self, record):
        """Parse record TFRecord into a set a set of values, names and types
        that can be queued and then read.

        Returns:
            - queue_values: Dict with tensor values.
            - queue_names: Names for each tensor.
            - queue_types: Types for each tensor.
        """
        # We parse variable length features (bboxes in a image) as sequence
        # features
        context_example, sequence_example = tf.parse_single_sequence_example(
            record,
            context_features=self.CONTEXT_FEATURES,
            sequence_features=self.SEQUENCE_FEATURES
        )

        # Decode image
        image_raw = tf.image.decode_image(
            context_example['image_raw'], channels=3
        )

        image = tf.cast(image_raw, tf.float32)

        height = tf.cast(context_example['height'], tf.int32)
        width = tf.cast(context_example['width'], tf.int32)
        image_shape = tf.stack([height, width, 3])
        image = tf.reshape(image, image_shape)

        label = self._sparse_to_tensor(sequence_example['label'])
        xmin = self._sparse_to_tensor(sequence_example['xmin'])
        xmax = self._sparse_to_tensor(sequence_example['xmax'])
        ymin = self._sparse_to_tensor(sequence_example['ymin'])
        ymax = self._sparse_to_tensor(sequence_example['ymax'])

        # Stack parsed tensors to define bounding boxes of shape (num_boxes, 5)
        bboxes = tf.stack([xmin, ymin, xmax, ymax, label], axis=1)

        image, bboxes, preprocessing_details = self.preprocess(image, bboxes)

        filename = tf.cast(context_example['filename'], tf.string)

        # TODO: Send additional metadata through the queue (scale_factor,
        # applied_augmentations)

        queue_dtypes = [tf.float32, tf.int32, tf.string]
        queue_names = ['image', 'bboxes', 'filename']
        queue_values = {
            'image': image,
            'bboxes': bboxes,
            'filename': filename,
        }

        return queue_values, queue_dtypes, queue_names