Access values in dict in tf.dataset.Dataset.map with tf.striing

Question

I am creating a tf.data.Dataset where I start with list_files to get all paths to my images. The annotations are stored on disc as json files. The structure of the json file is

{ 
   "img1.png": {
                data ...
               },
   "img2.png": ...
}

Hence the key-value is the image name.

I can easily extract the image names from the paths provided by list_files. However, that is tf.string, which cannot be used directly(?) to access the values in the annotation.

Is there an easy way to convert the tf.string to a python string so I can read the groundtruth data from the json file?

Alternatively convert the annotation to a proper tf type.

from typing import Mapping
from numpy import ndarray
import tensorflow as tf
import cv2 as cv
from pathlib import Path
from typing import Any, Mapping, NamedTuple
import json

class Point:
    x: float
    y: float

    def __init__(self, x: float, y: float):
        self.x = x
        self.y = y

class BoundingBox(NamedTuple):
    top: float
    left: float
    bottom: float
    right: float

class Annotation:
    image: tf.Tensor
    bounding_box: tf.Tensor
    is_visible: bool

    def __init__(self, image, bounding_box, is_visible):
        self.image = image
        self.bounding_box = bounding_box
        self.is_visible = is_visible

LABELS = {
    "NO_CLUB": 0,
    "CLUB": 1,
    "bbox": BoundingBox,
}


def is_in_split(image_path: tf.string, is_training: bool) -> bool:
    hash = tf.strings.to_hash_bucket_fast(image_path, 10)
    if is_training:
        return hash < 8
    else:
        return hash >= 8


def create_image_and_annotation(image_path: tf.string, annotation: Mapping[str, Any]):
    bits = tf.io.read_file(image_path)
    file_split = tf.strings.split(image_path, "/")
    image_name = file_split[-1]
    suffix = tf.strings.split(image_name, ".")[-1]

    jpeg = [
        tf.convert_to_tensor("jpg", dtype=tf.string),
        tf.convert_to_tensor("JPG", dtype=tf.string),
        tf.convert_to_tensor("jpeg", dtype=tf.string),
        tf.convert_to_tensor("JPEG", dtype=tf.string),
    ]
    is_jpeg = [tf.math.equal(suffix, s) for s in jpeg]
    png = [
        tf.convert_to_tensor("png", dtype=tf.string),
        tf.convert_to_tensor("PNG", dtype=tf.string),
    ]
    is_png = [tf.math.equal(suffix, s) for s in png]
    if tf.math.reduce_any(is_jpeg):
        image = tf.io.decode_jpeg(bits, channels=3)
    else:
        image = tf.io.decode_png(bits, channels=3)
    # Here I want to use image_name to access the annotation for the specific image! <---
    bounding_box = BoundingBox(0,0,10,10)
    return image, (bounding_box, True)


def createDataset(dir: Path, annotation: Mapping[str, Any], is_training: bool) -> tf.data.Dataset:
    image_path_png = str(dir / "images" / "*.png")
    image_path_PNG = str(dir / "images" / "*.PNG")
    image_path_jpg = str(dir / "images" / "*.jpg")
    image_path_JPG = str(dir / "images" / "*.JPG")
    image_path_jpeg = str(dir / "images" / "*.jpeg")
    image_path_JPEG = str(dir / "images" / "*.JPEG")
    image_dirs = [image_path_png, image_path_PNG, image_path_jpg, image_path_JPG, image_path_jpeg, image_path_JPEG]

    dataset = (tf.data.Dataset.list_files(image_dirs)
        .shuffle(1000)
        .map(lambda x: create_image_and_annotation(x, annotation))
        )

    for d in dataset:
        pass

    return dataset

def getDataset(data_root_path: Path, is_training: bool) -> tf.data.Dataset:
        dirs = [x for x in data_root_path.iterdir() if x.is_dir()]
        datasets = []
        for dir in dirs:
            json_path = dir / "annotations.json"
            with open(json_path) as json_file:
                annotation = json.load(json_file)
                createDataset(dir, annotation, is_training=is_training)



training_data = getDataset(Path("/home/erik/Datasets/ClubHeadDetection"), True)

Answer 1

It is possible when converting the feature name from string to int or you added a ragged tensor into it. You can map labels and data from the phase method. Train example

Sample: Same as mapping using dataset.map function but it specifically to convert to int or you can use a ragged tensor for < data, label >

import os
from os.path import exists

import tensorflow as tf
import tensorflow_io as tfio
from google.protobuf import json_format

import matplotlib.pyplot as plt

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
None
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
config = tf.config.experimental.set_memory_growth(physical_devices[0], True)
print(physical_devices)
print(config)

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Variables
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""           
image = tf.io.read_file( "F:\\Pictures\\Kids\\images.jpg" )
image = tf.io.decode_jpeg(image)                            
image = tf.cast( image, dtype=tf.int64 )
                    
example = tf.train.Example(
    features=tf.train.Features(
        feature={
            "1": tf.train.Feature(
            
                int64_list=tf.train.Int64List(
                    value=tf.constant( image, shape=( 183 * 275 * 3 ) ).numpy() )
                    
                    ),
            "2": tf.train.Feature(
            
                int64_list=tf.train.Int64List(
                    value=tf.constant( image, shape=( 183 * 275 * 3 ) ).numpy() )
                    
                    )       
            }))

data_string = json_format.MessageToJson(example)
example_binary = tf.io.decode_json_example(data_string)

example_phase = tf.io.parse_example(
    serialized=[example_binary.numpy()],
    features = {    "1": tf.io.FixedLenFeature(shape=[ 183 * 275 * 3 ], dtype=tf.int64),
                    "2": tf.io.FixedLenFeature(shape=[ 183 * 275 * 3 ], dtype=tf.int64)
                })

data = list(example_phase.items())
label = [ int(x[0]) for x in data ]

data = [ x[1] for x in data ]
data = [ x[0] for x in data ]

dataset = tf.data.Dataset.from_tensors(( data, label ))

# for d in dataset:
    # print( d )

plt.imshow( tf.constant( example_phase['1'], shape=( 183, 275, 3) ) )
plt.show()

Output:

(<tf.Tensor: shape=(2, 150975), dtype=int64, numpy=
array([[227, 162, 132, ..., 227, 169, 129],
       [227, 162, 132, ..., 227, 169, 129]], dtype=int64)>, <tf.Tensor: shape=(2,), dtype=int32, numpy=array([1, 2])>)

Input used

Answer 2

The easiest solution was to read the file using

annotation = tf.io.read_file(str(json_path))

Then the bounding box is created by calling

bbox = tf.py_function(
        create_bbox,
        inp=[annotation, image_name],
        Tout=[tf.float32, tf.float32, tf.float32, tf.float32, tf.float32],
    )

Inside create_bbox we can now call python functions. Such as

annotation_py = annotation.numpy()
annotation_json = json.loads(annotation_py)
key_py = image_name.numpy().decode("utf-8")