ganime/data/experimental.py

from abc import ABC, abstractclassmethod, abstractmethod
import glob
import math
import os
from typing import Dict
from typing_extensions import dataclass_transform

import numpy as np
import tensorflow as tf
from tqdm.auto import tqdm


def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):  # if value ist tensor
        value = value.numpy()  # get value of tensor
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _float_feature(value):
    """Returns a floast_list from a float / double."""
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))


def _int64_feature(value):
    """Returns an int64_list from a bool / enum / int / uint."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def serialize_array(array):
    array = tf.io.serialize_tensor(array)
    return array


class Dataset(ABC):
    def __init__(self, dataset_path: str):
        self.dataset_path = dataset_path

    @classmethod
    def _parse_single_element(cls, element) -> tf.train.Example:

        features = tf.train.Features(feature=cls._get_features(element))

        return tf.train.Example(features=features)

    @abstractclassmethod
    def _get_features(cls, element) -> Dict[str, tf.train.Feature]:
        pass

    @abstractclassmethod
    def _parse_tfr_element(cls, element):
        pass

    @classmethod
    def write_to_tfr(cls, data: np.ndarray, out_dir: str, filename: str):
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

        # Write all elements to a single tfrecord file
        single_file_name = cls.__write_to_single_tfr(data, out_dir, filename)

        # The optimal size for a single tfrecord file is around 100 MB. Get the number of files that need to be created
        number_splits = cls.__get_number_splits(single_file_name)

        if number_splits > 1:
            os.remove(single_file_name)
            cls.__write_to_multiple_tfr(data, out_dir, filename, number_splits)

    @classmethod
    def __write_to_multiple_tfr(
        cls, data: np.array, out_dir: str, filename: str, n_splits: int
    ):

        file_count = 0

        max_files = math.ceil(data.shape[0] / n_splits)

        print(f"Creating {n_splits} files with {max_files} elements each.")

        for i in tqdm(range(n_splits)):
            current_shard_name = os.path.join(
                out_dir,
                f"{filename}.tfrecords-{str(i).zfill(len(str(n_splits)))}-of-{n_splits}",
            )
            writer = tf.io.TFRecordWriter(current_shard_name)

            current_shard_count = 0
            while current_shard_count < max_files:  # as long as our shard is not full
                # get the index of the file that we want to parse now
                index = i * max_files + current_shard_count
                if index >= len(
                    data
                ):  # when we have consumed the whole data, preempt generation
                    break

                current_element = data[index]

                # create the required Example representation
                out = cls._parse_single_element(element=current_element)

                writer.write(out.SerializeToString())
                current_shard_count += 1
                file_count += 1

        writer.close()
        print(f"\nWrote {file_count} elements to TFRecord")
        return file_count

    @classmethod
    def __get_number_splits(cls, filename: str):
        target_size = 100 * 1024 * 1024  # 100mb

        single_file_size = os.path.getsize(filename)
        number_splits = math.ceil(single_file_size / target_size)
        return number_splits

    @classmethod
    def __write_to_single_tfr(cls, data: np.array, out_dir: str, filename: str):

        current_path_name = os.path.join(
            out_dir,
            f"{filename}.tfrecords-0-of-1",
        )

        writer = tf.io.TFRecordWriter(current_path_name)
        for element in tqdm(data):
            writer.write(cls._parse_single_element(element).SerializeToString())
        writer.close()

        return current_path_name

    def load(self) -> tf.data.TFRecordDataset:
        path = self.dataset_path
        dataset = None

        if os.path.isdir(path):
            dataset = self._load_folder(path)
        elif os.path.isfile(path):
            dataset = self._load_file(path)
        else:
            raise ValueError(f"Path {path} is not a valid file or folder.")

        dataset = dataset.map(self._parse_tfr_element)
        return dataset

    def _load_file(self, path) -> tf.data.TFRecordDataset:
        return tf.data.TFRecordDataset(path)

    def _load_folder(self, path) -> tf.data.TFRecordDataset:

        return tf.data.TFRecordDataset(
            glob.glob(os.path.join(path, "**/*.tfrecords*"), recursive=True)
        )


class VideoDataset(Dataset):
    @classmethod
    def _get_features(cls, element) -> Dict[str, tf.train.Feature]:
        return {
            "frames": _int64_feature(element.shape[0]),
            "height": _int64_feature(element.shape[1]),
            "width": _int64_feature(element.shape[2]),
            "depth": _int64_feature(element.shape[3]),
            "raw_video": _bytes_feature(serialize_array(element)),
        }

    @classmethod
    def _parse_tfr_element(cls, element):
        # use the same structure as above; it's kinda an outline of the structure we now want to create
        data = {
            "frames": tf.io.FixedLenFeature([], tf.int64),
            "height": tf.io.FixedLenFeature([], tf.int64),
            "width": tf.io.FixedLenFeature([], tf.int64),
            "raw_video": tf.io.FixedLenFeature([], tf.string),
            "depth": tf.io.FixedLenFeature([], tf.int64),
        }

        content = tf.io.parse_single_example(element, data)

        frames = content["frames"]
        height = content["height"]
        width = content["width"]
        depth = content["depth"]
        raw_video = content["raw_video"]

        # get our 'feature'-- our image -- and reshape it appropriately
        feature = tf.io.parse_tensor(raw_video, out_type=tf.uint8)
        feature = tf.reshape(feature, shape=[frames, height, width, depth])
        return feature


class ImageDataset(Dataset):
    @classmethod
    def _get_features(cls, element) -> Dict[str, tf.train.Feature]:
        return {
            "height": _int64_feature(element.shape[0]),
            "width": _int64_feature(element.shape[1]),
            "depth": _int64_feature(element.shape[2]),
            "raw_image": _bytes_feature(serialize_array(element)),
        }

    @classmethod
    def _parse_tfr_element(cls, element):
        # use the same structure as above; it's kinda an outline of the structure we now want to create
        data = {
            "height": tf.io.FixedLenFeature([], tf.int64),
            "width": tf.io.FixedLenFeature([], tf.int64),
            "raw_image": tf.io.FixedLenFeature([], tf.string),
            "depth": tf.io.FixedLenFeature([], tf.int64),
        }

        content = tf.io.parse_single_example(element, data)

        height = content["height"]
        width = content["width"]
        depth = content["depth"]
        raw_image = content["raw_image"]

        # get our 'feature'-- our image -- and reshape it appropriately
        feature = tf.io.parse_tensor(raw_image, out_type=tf.uint8)
        feature = tf.reshape(feature, shape=[height, width, depth])
        return feature