SVGBuilder/data_processing.py at main · amcghm/SVGBuilder · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
"""
Data Processing and Tokenization for SVG Generation

This module handles the transformation of raw dataset examples into tokenized
tensors ready for model training and evaluation. It includes building the
vocabulary, generating incremental image features, and tokenizing text/SVG
sequences.
"""
import torch
from utils.common import svg_to_img, png_to_jpg_white, scaled_paths_to_incremental_svgs

def build_vocab(train_ds):
    """Builds the vocabulary from the training dataset."""
    path_set = set(path['d'] for example in train_ds.dataset for path in example['svg'])

    # Special tokens, path, offset_x, offset_y, scale, r, g, b
    vocab_list = ['[PAD]', '[UNK]', '[BOS]', '[EOS]'] + \
        list(path_set) + \
        list(train_ds.offset_x_set) + list(train_ds.offset_y_set) + list(train_ds.scale_set) + \
        list(train_ds.r_set) + list(train_ds.g_set) +list(train_ds.b_set)

    vocab = {word: index for index, word in enumerate(vocab_list)}
    id2token = {v: k for k, v in vocab.items()}

    return vocab, id2token

class DataProcessor:
    def __init__(self, vocab, text_tokenizer, clip_model, clip_processor, config):
        self.vocab = vocab
        self.text_tokenizer = text_tokenizer
        self.clip_model = clip_model
        self.clip_processor = clip_processor
        self.config = config
        self.device = clip_model.device

    def _pad_and_truncate_list(self, ls, max_length, pad_value):
        if len(ls) >= max_length:
            return ls[:max_length]

        padding_length = max_length - len(ls)
        # For features
        if isinstance(pad_value, list):
            return ls + [pad_value] * padding_length
        # For token IDs
        else:
            return ls + [pad_value] * padding_length

    @torch.no_grad()
    def add_incremental_svg_imgs(self, example):
        """Generates and adds incremental image features for an example."""
        svgs = scaled_paths_to_incremental_svgs(example['svg'], add_blank_first=True)
        imgs = list(map(svg_to_img, svgs))
        imgs = list(map(png_to_jpg_white, imgs))
        inputs = self.clip_processor(images=imgs, return_tensors="pt").to(self.device)
        features = self.clip_model.get_image_features(**inputs)

        if hasattr(features, 'pooler_output'):
            features = features.pooler_output
        elif hasattr(features, 'image_embeds'):
            features = features.image_embeds
        elif isinstance(features, tuple):
            features = features[0]

        example['imgs_features'] = features.cpu().tolist()
        return example

    def tokenize_train(self, example):
        """Tokenizes a single training example."""
        encoding = {}
        # Text
        text_encoding = self.text_tokenizer(
            example['text'], padding='max_length', truncation=True, max_length=self.config['text_max_length']
        )
        encoding['text_input_ids'] = text_encoding['input_ids']
        encoding['text_attention_mask'] = text_encoding['attention_mask']

        # SVG label
        svg_sequence = []
        for path in example['svg']:
            svg_sequence.extend([
                self.vocab[path['d']],
                self.vocab[f"x{path['offset_x']}"], self.vocab[f"y{path['offset_y']}"],
                self.vocab[f"s{path['scale']}"],
                self.vocab[f"r{path['rgb_fill'][0]}"], self.vocab[f"g{path['rgb_fill'][1]}"], self.vocab[f"b{path['rgb_fill'][2]}"]
            ])

        # label needs a [EOS] token at end
        encoding['label'] = self._pad_and_truncate_list(
            svg_sequence + [self.vocab['[EOS]']], self.config['svg_max_length'], self.vocab['[PAD]']
        )
        # svg input ids need a [BOS] token at begin
        encoding['svg_input_ids'] = self._pad_and_truncate_list(
            [self.vocab['[BOS]']] + svg_sequence, self.config['svg_max_length'], self.vocab['[PAD]']
        )
        # attention mask
        encoding['svg_attention_mask'] = [1 if input_id != self.vocab['[PAD]'] else 0 for input_id in encoding['svg_input_ids']]

        # incremental svg imgs
        encoding['imgs_features'] = self._pad_and_truncate_list(
            example['imgs_features'], self.config['component_max_length'], [0] * self.config['d_model']
        )
        encoding['svg_str'] = example['svg_str']

        return encoding

    def tokenize_test_and_eval(self, examples):
        """Tokenizes a single evaluation/test example."""
        encoding = {}
        text_encoding = self.text_tokenizer(
            examples['text'], padding='max_length', truncation=True, max_length=self.config['text_max_length']
        )
        encoding['text_input_ids'] = text_encoding['input_ids']
        encoding['text_attention_mask'] = text_encoding['attention_mask']
        encoding['text'] = examples['text']
        encoding['svg_str'] = examples['svg_str']
        return encoding