add dataset helper

dataset_helper.py

@@ -0,0 +1,326 @@
+import random
+from typing import Tuple, Any
+
+import numpy as np
+import pandas as pd
+import torch
+
+# SEQUENCE GENERATION
+PADDING_VALUE = float('-100')
+
+# ANIMATION_PARAMETER_INDICES = {
+#     0: [],  # EOS
+#     1: [10, 11, 12, 13],  # translate: begin, dur, x, y
+#     2: [10, 11, 14, 15],  # curve: begin, dur, via_x, via_y
+#     3: [10, 11, 16], # scale: begin, dur, from_factor
+#     4: [10, 11, 17], # rotate: begin, dur, from_degree
+#     5: [10, 11, 18], # skewX: begin, dur, from_x
+#     6: [10, 11, 19], # skewY: begin, dur, from_y
+#     7: [10, 11, 20, 21, 22], # fill: begin, dur, from_r, from_g, from_b
+#     8: [10, 11, 23], # opcaity: begin, dur, from_f
+#     9: [10, 11, 24], # blur: begin, dur, from_f
+# }
+
+ANIMATION_PARAMETER_INDICES = {
+    0: [],  # EOS
+    1: [0, 1, 2, 3],  # translate: begin, dur, x, y
+    2: [0, 1, 4, 5],  # curve: begin, dur, via_x, via_y
+    3: [0, 1, 6], # scale: begin, dur, from_factor
+    4: [0, 1, 7], # rotate: begin, dur, from_degree
+    5: [0, 1, 8], # skewX: begin, dur, from_x
+    6: [0, 1, 9], # skewY: begin, dur, from_y
+    7: [0, 1, 10, 11, 12], # fill: begin, dur, from_r, from_g, from_b
+    8: [0, 1, 13], # opcaity: begin, dur, from_f
+    9: [0, 1, 14], # blur: begin, dur, from_f
+}
+
+
+def unpack_embedding(embedding: torch.Tensor, dim=0, device="cpu") -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Args:
+        device: cpu / gpu
+        dim: dimension where the embedding is positioned
+        embedding: embedding of dimension 270
+
+    Returns: tuple of tensors: deep-svg embedding, type of prediction, animation parameters
+
+    """
+    if embedding.shape[dim] != 282:
+        print(embedding.shape)
+        raise ValueError('Dimension of 270 required.')
+
+    if dim == 0:
+        deep_svg = embedding[: -26].to(device)
+        types = embedding[-26: -15].to(device)
+        parameters = embedding[-15:].to(device)
+
+    elif dim == 1:
+        deep_svg = embedding[:, : -26].to(device)
+        types = embedding[:, -26: -15].to(device)
+        parameters = embedding[:, -15:].to(device)
+
+    elif dim == 2:
+        deep_svg = embedding[:, :, : -26].to(device)
+        types = embedding[:, :, -26: -15].to(device)
+        parameters = embedding[:, :, -15:].to(device)
+
+    else:
+        raise ValueError('Dimension > 2 not possible.')
+    return deep_svg, types, parameters
+
+
+def generate_dataset(dataframe_index: pd.DataFrame,
+                     input_sequences_dict_used: dict,
+                     input_sequences_dict_unused: dict,
+                     output_sequences: pd.DataFrame,
+                     logos_list: dict,
+                     sequence_length_input: int,
+                     sequence_length_output: int,
+                     ) -> dict:
+    """
+    Builds the dataset and returns it
+
+    Args:
+        input_sequences_dict_used: dictionary containing input sequences per logo
+        input_sequences_dict_unused: dictionary containing all unused paths
+        dataframe_index: dataframe containing the relevant indexes for the dataframes
+        output_sequences: dataframe containing animations
+        logos_list: dictionary in train/test split containing list for logo ids
+        sequence_length_input: length of input sequence for padding
+        sequence_length_output: length of output sequence for padding
+
+    Returns: dictionary containing the dataset for training/testing
+
+    """
+    dataset = {
+        "is_bucketing": False,
+        "train": {
+            "input": [],
+            "output": []
+        },
+        "test": {
+            "input": [],
+            "output": []
+        }
+    }
+    for i, logo_info in dataframe_index.iterrows():
+        logo = logo_info['filename']  # e.g. logo_1
+        file = logo_info['file']  # e.g. logo_1_animation_2
+        oversample = logo_info['repeat']
+        print(f"Processing {logo} with {file}: ")
+
+        if input_sequences_dict_used.keys().__contains__(logo) and input_sequences_dict_unused.keys().__contains__(logo):
+            for j in range(oversample):
+                input_tensor = _generate_input_sequence(
+                    input_sequences_dict_used[logo].copy(),
+                    input_sequences_dict_unused[logo].copy(),
+                    #pd.DataFrame(),
+                    null_features=26,  # TODO depends on architecture later
+                    sequence_length=sequence_length_input,
+                    # is_randomized=True, always now
+                    is_padding=True
+                )
+
+                output_tensor = _generate_output_sequence(
+                    output_sequences[(output_sequences['filename'] == logo) & (output_sequences['file'] == file)].copy(),
+                    sequence_length=sequence_length_output,
+                    is_randomized=False,
+                    is_padding=True
+                )
+                # append to lists
+                if logo in logos_list["train"]:
+                    random_index = random.randint(0, len(dataset["train"]["input"]))
+                    dataset["train"]["input"].insert(random_index, input_tensor)
+                    dataset["train"]["output"].insert(random_index, output_tensor)
+
+                elif logo in logos_list["test"]:
+                    dataset["test"]["input"].append(input_tensor)
+                    dataset["test"]["output"].append(output_tensor)
+                    break  # no oversampling in testing
+
+                else:
+                    print(f"Some problem with {logo}. Neither in train or test set list.")
+                    break
+
+    dataset["train"]["input"] = torch.stack(dataset["train"]["input"])
+    dataset["train"]["output"] = torch.stack(dataset["train"]["output"])
+    dataset["test"]["input"] = torch.stack(dataset["test"]["input"])
+    dataset["test"]["output"] = torch.stack(dataset["test"]["output"])
+
+    return dataset
+
+
+def _generate_input_sequence(logo_embeddings_used: pd.DataFrame,
+                             logo_embeddings_unused: pd.DataFrame,
+                             null_features: int,
+                             sequence_length: int,
+                             is_padding: bool) -> torch.Tensor:
+    """
+    Build a torch tensor for the transformer input sequences.
+    Includes
+    - Ensuring all used embeddings are included
+    - Filling the remainder with unused embeddings up to sequence length
+    - Generation of padding
+
+    Args:
+        logo_embeddings (pd.DataFrame): DataFrame containing logo embeddings.
+        null_features (int): Number of null features to add to each embedding.
+        sequence_length (int): Target length for padding sequences.
+        is_padding: if true, function adds padding
+
+    Returns:
+        torch.Tensor: Tensor representing the input sequences.
+    """
+    logo_embeddings_used.drop(columns=['filename', 'animation_id'], inplace=True)
+    logo_embeddings_unused.drop(columns=['filename', 'animation_id'], inplace=True)
+
+    # Combine used and unused. Fill used with random unused samples
+    logo_embeddings = logo_embeddings_unused
+    remaining_slots = sequence_length - len(logo_embeddings)
+    if remaining_slots > 0:
+        sample_size = min(len(logo_embeddings_unused), remaining_slots)
+        additional_embeddings = logo_embeddings_unused.sample(n=sample_size, replace=False)
+        logo_embeddings = pd.concat([logo_embeddings, additional_embeddings], ignore_index=True)
+        logo_embeddings.reset_index()
+
+    # Randomization
+    logo_embeddings = logo_embeddings.sample(frac=1).reset_index(drop=True)
+
+    # Null Features
+    if null_features > 0:
+        logo_embeddings = pd.concat([logo_embeddings,
+                                     pd.DataFrame(0,
+                                                  index=logo_embeddings.index,
+                                                  columns=range(logo_embeddings.shape[1],
+                                                                logo_embeddings.shape[1] + null_features))],
+                                    axis=1,
+                                    ignore_index=True)
+
+    if is_padding:
+        logo_embeddings = _add_padding(logo_embeddings, sequence_length)
+
+    return torch.tensor(logo_embeddings.values)
+
+
+def _generate_output_sequence(animation: pd.DataFrame,
+                              sequence_length: int,
+                              is_randomized: bool,
+                              is_padding: bool) -> torch.Tensor:
+    """
+    Build a torch tensor for the transformer output sequences.
+    Includes
+    - Randomization (later, when same start time)
+    - Generation of padding
+    - Add EOS Token
+
+    Args:
+        animation (pd.DataFrame): DataFrame containing logo embeddings.
+        sequence_length (int): Target length for padding sequences.
+        is_randomized: shuffle order of paths, applies when same start time
+        is_padding: if true, function adds padding
+
+    Returns:
+        torch.Tensor: Tensor representing the input sequences.
+    """
+    if is_randomized:
+        animation = animation.sample(frac=1).reset_index(drop=True)
+        print("Note: Randomization not implemented yet")
+
+    animation.sort_values(by=['a10'], inplace=True)  # again ordered by time start.
+    animation.drop(columns=['file', 'filename', "Unnamed: 0",	"id"], inplace=True)
+
+    # Append the EOS row to the DataFrame
+    sos_eos_row = {col: 0 for col in animation.columns}
+    sos_eos_row["a0"] = 1
+    sos_eos_row = pd.DataFrame([sos_eos_row])
+    animation = pd.concat([sos_eos_row, animation, sos_eos_row],
+                          ignore_index=True)
+
+    # Padding Generation: Add padding rows or cut off excess rows
+    if is_padding:
+        animation = _add_padding(animation, sequence_length)
+
+    return torch.Tensor(animation.values)
+
+
+def _add_padding(dataframe: pd.DataFrame, sequence_length: int) -> pd.DataFrame:
+    """
+    Add padding to a dataframe
+
+    Args:
+        dataframe: dataframe to add padding to
+        sequence_length: length of final sequences
+
+    Returns:
+
+    """
+    if len(dataframe) < sequence_length:
+        padding_rows = pd.DataFrame([[PADDING_VALUE] * len(dataframe.columns)] * (sequence_length - len(dataframe)),
+                                    columns=dataframe.columns)
+        dataframe = pd.concat([dataframe, padding_rows], ignore_index=True)
+    elif len(dataframe) > sequence_length:
+        # Cut off excess rows
+        dataframe = dataframe.iloc[:sequence_length]
+
+    return dataframe
+
+
+# BUCKETING
+def generate_buckets_2D(dataset, column1, column2, quantiles1, quantiles2, print_histogram=True):
+    """
+
+    Args:
+        dataset: dataset to generate buckets for
+        column1: first column name
+        column2: second column name
+        quantiles1: initial quantiles for column1
+        quantiles2: initial quantiles for column2
+        print_histogram: if true, a histogram of the 2D buckets is printed
+
+    Returns: dictionary object with bucket edges
+
+    """
+    x_edges = dataset[column1].quantile(quantiles1)
+    y_edges = dataset[column2].quantile(quantiles2)
+
+    x_edges = np.array(x_edges)
+    y_edges = np.unique(y_edges)
+
+    if print_histogram:
+        hist, x_edges, y_edges = np.histogram2d(dataset[column1],
+                                                dataset[column2],
+                                                bins=[x_edges, y_edges])
+        print(hist)
+
+    return {
+        "input_edges": list(x_edges),
+        "output_edges": list(y_edges)
+    }
+
+
+def get_bucket(input_length, output_length, buckets):
+    bucket_name = ""
+
+    for i, input_edge in enumerate(buckets["input_edges"]):
+        # print(f"{i}: {input_length} < {input_edge}")
+        if input_length > input_edge:
+            continue
+
+        bucket_name = bucket_name + str(int(i))  # chr(ord('A')+i)
+        break
+
+    bucket_name = bucket_name + "-"
+
+    for i, output_edge in enumerate(buckets["output_edges"]):
+        if output_length > output_edge:
+            continue
+
+        bucket_name = bucket_name + str(int(i))
+        break
+
+    return bucket_name
+
+
+def warn_if_contains_NaN(dataset: torch.Tensor):
+    if torch.isnan(dataset).any():
+        print("There are NaN values in the dataset")