SmilingWolf/danbooru-v4-tfrecord

NUM_CLASSES = 12822


# Function to convert value to bytes_list
def _bytes_feature(value):
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy()
    elif isinstance(value, str):
        value = value.encode()
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


# Function to convert bool/enum/int/uint to int64_list
def _int64_feature(value):
    int64_list = tf.train.Int64List(value=tf.reshape(value, (-1,)))
    return tf.train.Feature(int64_list=int64_list)


# Function to create a tf.train.Example message
def serialize_example(image_id, image_bytes, label_indexes, tag_string):
    feature = {
        "image_id": _int64_feature(image_id),
        "image_bytes": _bytes_feature(image_bytes),
        "label_indexes": _int64_feature(label_indexes),
        "tag_string": _bytes_feature(tag_string),
    }
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()


# Function to deserialize a single tf.train.Example message
def deserialize_example(example_proto):
    feature_description = {
        "image_id": tf.io.FixedLenFeature([], tf.int64),
        "image_bytes": tf.io.FixedLenFeature([], tf.string),
        "label_indexes": tf.io.VarLenFeature(tf.int64),
        "tag_string": tf.io.FixedLenFeature([], tf.string),
    }

    # Parse the input 'tf.train.Example' proto using the dictionary above.
    parsed_example = tf.io.parse_single_example(example_proto, feature_description)
    image_tensor = tf.io.decode_jpeg(parsed_example["image_bytes"], channels=3)

    # We only stored label indexes in the TFRecords to save space
    # Emulate MultiLabelBinarizer to get a tensor of 0s and 1s
    label_indexes = tf.sparse.to_dense(
        parsed_example["label_indexes"],
        default_value=0,
    )
    one_hots = tf.one_hot(label_indexes, NUM_CLASSES)
    labels = tf.reduce_max(one_hots, axis=0)
    labels = tf.cast(labels, tf.float32)

    sample = {
        "image_ids": parsed_example["image_id"],
        "images": image_tensor,
        "labels": labels,
        "tags": parsed_example["tag_string"],
    }
    return sample