quantizer.py

import torch
import numpy as np
import copy
from collections import OrderedDict
import json
from datasets import ClassLabel
import random
import math
from functools import lru_cache
from matplotlib import font_manager
from colorama import Fore, Style, init 


class BaseQuantizer:
    @property
    def ignore_tokens(self):
        if self.num_mask_tokens > 0:
            if self.mask_type == 'cm3':
                return [self.predict_start_token] + self.mask_tokens
            elif self.mask_type == 'mask_aug':
                return [self.mask_aug_token]
            else:
                raise ValueError(f'Invalid mask type {self.mask_type}')
        else:
            return []

    def __init__(self, simplify_json=False, mask_all=False,
                    num_mask_tokens=0, mask_type='cm3', **kwargs):
        self.simplify_json=simplify_json
        self.io_ignore_replace_tokens = ['<split-text>']
        self.mask_all = mask_all
        self.num_mask_tokens = num_mask_tokens
        self.mask_type = mask_type
        if self.mask_type == 'mask_aug':
            self.mask_aug_token = '<mask-aug>'
        elif self.mask_type == 'cm3':
            self.predict_start_token = '<pred-start>'
        else:
            raise ValueError(f'Invalid mask type {self.mask_type}')

    def get_additional_mask_tokens(self):
        if self.mask_type == 'cm3': # 两种配置：1. ['<pred-start>'] + '<mask-%d>'，数量和self.num_mask_tokens相关 2. ['<mask-aug>']
            self.mask_tokens = ['<mask-%d>' % i for i in range(self.num_mask_tokens)]
            return [self.predict_start_token] + self.mask_tokens
        elif self.mask_type == 'mask_aug':
            return [self.mask_aug_token]
        else:
            raise ValueError(f'Invalid mask type {self.mask_type}')

    def dump2json(self, json_example):
        if self.simplify_json: # 将 dict 转化为 str， 如果simplify_json is True，那么缩减空格和换行，删除token的双引号
            content = json.dumps(json_example, separators=(',',':'))
            for token in self.additional_special_tokens:
                content = content.replace(f'"{token}"', token)
        else:
            content = json.dumps(json_example)
        return content

    def load_json(self, content): # 将str转化为json
        replace_tokens = set(self.additional_special_tokens) - set(self.io_ignore_replace_tokens) # sirui change 
        if self.simplify_json:
            for token in replace_tokens: # 如果simplify_json is True，那么为 token 添加双引号
                content = content.replace(token, f'"{token}"')
        return json.loads(content)

    def apply_masking(self, 
                  json_example, 
                  mask_all=None, 
                  return_meta=False, 
                  target_keys=['width', 'height', 'left', 'top'],
                  target_element_types=None
                  ):
        if mask_all is None:
            mask_all = self.mask_all
        json_example = copy.deepcopy(json_example)
        target_keys = set(target_keys)
        target_tokens = []
        for shape_i, shape in enumerate(json_example['layers']['textlayer']):
            # element_type = self.general_dequantize(shape['type'],'type',to_float=False)
            # if target_element_types is not None:
                # if element_type not in target_element_types:
                    # continue
            for key_i, key in enumerate(shape.keys()):
                if key in target_keys:
                    target_tokens.append((shape_i, key_i, key, shape[key]))
        if not mask_all:
            target_num_mask_tokens = random.randint(1, self.num_mask_tokens)
            if len(target_tokens) > target_num_mask_tokens:
                random.shuffle(target_tokens)
                target_tokens = target_tokens[:target_num_mask_tokens]
                # sort by shape_i and key_i
                target_tokens = sorted(target_tokens, key=lambda x: x[0]*100+x[1])
        else:
            if len(target_tokens) > self.num_mask_tokens:
                # 取最后面几个
                target_tokens = target_tokens[-self.num_mask_tokens:]

        tuples = []
        meta_infos = []
        for mask_i, (shape_i, key_i, key, value) in enumerate(target_tokens):
            if self.mask_type == 'cm3':
                mask_token = self.mask_tokens[mask_i]
            elif self.mask_type == 'mask_aug':
                mask_token = self.mask_aug_token
            else:
                raise ValueError(f'Invalid mask type {self.mask_type}')
            # <one-1><decimal0-1><decimal1-2>
            if '<' in value:
                num_token = value.count('<')
            else:
                num_token = value.count(' ')
            json_example['layers']['textlayer'][shape_i][key] = mask_token
            tuples.append((mask_token, value, num_token))
            meta_infos.append((shape_i,key))
        if return_meta:
            return json_example, tuples, meta_infos
        else:
            return json_example, tuples

    def make_prediction_postfix(self, tuples):
        postfix = self.predict_start_token
        for mask_token, value, num_token in tuples:
            postfix = postfix+ f'{mask_token}{value}'
        return postfix

# specs={
#     "width":"size",
#     "height":"size",
#     "left":"pos",
#     "top":"pos",
#     "x":"pos", # center x
#     "y":"pos", # center y
#     "opacity":"opacity",
#     "color":"color",
#     "angle":"angle",
#     "font_size":"font_size",
#     'ratio':'ratio',
#     'letter_spacing': 'spacing',
#     'textlen': 'textlen'
# }

specs={
    "width":"size",
    "height":"size",
    "x":"pos", # center x
    "y":"pos", # center y
    "color":"color",
    "font":"font"
}

# TODO change min_max_bins
# min_max_bins = {
#     'size':(0,2,256),
#     'pos':(-1,1,256),
#     # 'opacity':(0,1,8),
#     'opacity':(0,255,8),
#     'color':(0,255,32),
#     'angle':(0,2*np.pi,64),
#     'font_size':(2,200,100),
#     'spacing': (0,1,40),
#     'textlen': (1,20,20)
# }
min_max_bins = {
    'size': (0,1,256),
    'pos': (0,1,256),
    'color': (0,137,138),
    'font': (0,511,512)
}

import numpy as np

# pre 和 post 分别代表 10 的幂，分别对应大数和小数部分，参数代表位数
def get_keys_and_multipliers(pre_decimal=3, post_decimal=2):
    pre_keys = ['one', 'ten', 'hundred', 'thousand']
    pre_multiplers = [1, 10, 100, 1000]
    assert pre_decimal <= len(pre_keys)
    pre_keys = pre_keys[:pre_decimal][::-1]
    pre_multiplers = pre_multiplers[:pre_decimal][::-1]

    post_keys = [f'decimal{x}' for x in range(post_decimal)]
    post_multiplers = [10 ** -(x+1) for x in range(post_decimal)]

    keys =  pre_keys + post_keys
    multiplers = pre_multiplers + post_multiplers
    return keys, multiplers

class DecimalQuantizer:
    def __init__(self, max_pre_decimal=3, max_post_decimal=2):
        self.max_pre_decimal = max_pre_decimal
        self.max_post_decimal = max_post_decimal
        self.keys, self.multiplers = get_keys_and_multipliers(max_pre_decimal, max_post_decimal)
        self.symbols = {
            -1: '<symbol-1>',
            1: '<symbol-0>',
        }

    def get_vocab(self):
        special_tokens = [*self.symbols.values()]   # ['<symbol-1>', '<symbol-0>']
        for key in self.keys:  # ['one', 'ten', 'hundred', 'thousand'] + ['decimal0', 'decimal1]
            special_tokens.extend([f'<{key}-{i}>' for i in range(10)])
        return special_tokens

    def check_valid(self, token):
        prefix = token.lstrip('<').split('-')[0] # '<symbol-1>' -> 'symbol-1>' -> ['symbol', '1>']
        if prefix =='symbol' or prefix in self.keys:
            return True
        else:
            return False

    # 小数点后保留两位
    def __call__(self, val, pre_decimal=None, post_decimal=None, need_symbol=False): # 100.00
        if pre_decimal is None:
            pre_decimal = self.max_pre_decimal
        if post_decimal is None:
            post_decimal = self.max_post_decimal
        
        assert pre_decimal <= self.max_pre_decimal
        assert post_decimal <= self.max_post_decimal

        keys, multiplers = get_keys_and_multipliers(pre_decimal, post_decimal)
        
        symbol = int(np.sign(val)) # 返回一个浮点数（1.0, -1.0 或 0.0），代表正负和0
        if symbol == 0: # 两类：>= 0  &  < 0
            symbol = 1
        val = round(abs(val), post_decimal) # 将 val 的绝对值四舍五入到 post_decimal 位小数
        
        tokens = []
        if need_symbol: # self.symbols = {-1: '<symbol-1>', 1: '<symbol-0>',}
            symbol_type = self.symbols[symbol]
            tokens.append(symbol_type)
        else:
            assert symbol >= 0

        for key, multipler in zip(keys, multiplers):
            # 用于获取对于给定数值 val，每一位的数字，并且生成为'<one-7>'这样的token
            v = math.floor(val / multipler)
            if v > 9:
                raise ValueError(f'Invalid value {val} for {pre_decimal} pre_decimal and {post_decimal} post_decimal')
            val = val - v * multipler
            tokens.append(f'<{key}-{v}>')
            
        # 对于val，生成每一位数字对应的token，如果need_symbol = True，还会在前面加上 标识 >= 0 和 < 0 的 symbol-1 和 symbol-0
        return ''.join(tokens)

    def parse_token(self, token):
        # <hundred-1> -> hundred, 1
        key, val = token[1:-1].split('-')
        return key, int(val)
    
    def decode(self, tokens_str): # 将token_str用 > 先拆开，再添上 > ，然后转化为 list
        tokens = tokens_str.split('>')
        tokens = [x+'>' for x in tokens if x != '']
        if tokens[0].startswith('<symbol'):
            symbol_type = tokens[0]
            tokens = tokens[1:]
            inv_map = {v: k for k, v in self.symbols.items()} # 和 原字典 键、值 对调
            symbol = inv_map[symbol_type]
        else:
            symbol = 1

        accumulater = 0
        for token in tokens:
            key, val = self.parse_token(token)
            multipler_index = self.keys.index(key)
            multipler = self.multiplers[multipler_index]
            actual_val = val * multipler
            # print(key, val, multipler, actual_val)
            accumulater += actual_val
        accumulater = accumulater * symbol
        
        # 还原出原来的整数，带有符号，并且精度 由 pre/post_decimal位数控制
        return accumulater

# min_max_bins = {
#     'size': (0,1,256),
#     'pos': (0,1,256),
#     'color': (0,137,138),
#     'font': (0,511,512)
# }
pre_post_decimals={
    'size': {
        'pre_decimal': 1,
        'post_decimal': 2,
        'need_symbol': False
    },
    'pos': {
        'pre_decimal': 1,
        'post_decimal': 2,
        'need_symbol': True
    },
    'opacity': {
        'pre_decimal': 1,
        'post_decimal': 1,
        'need_symbol': False
    },
    'color':{
        'pre_decimal': 3,
        'post_decimal': 0,
        'need_symbol': False
    },
    'angle':{
        'pre_decimal': 1,
        'post_decimal': 2,
        'need_symbol': False
    },
    'font_size':{
        'pre_decimal': 3,
        'post_decimal': 0,
        'need_symbol': False
    },
}

class QuantizerV4(BaseQuantizer):
    def __init__(self, quant=True,
                 decimal_quantize_types = [],
                 decimal_quantize_kwargs = {'max_pre_decimal':3, 'max_post_decimal':2},
                 mask_values=False,
                 **kwargs):
        super().__init__(**kwargs)
        self.min = min
        self.max = max
        self.quant = quant
        self.mask_values = mask_values
        self.text_split_token = '<split-text>'
        self.decimal_quantize_types = decimal_quantize_types
        self.decimal_quantize = len(decimal_quantize_types) > 0
        if len(decimal_quantize_types) > 0:
            print('decimal quantize types', decimal_quantize_types)
            self.decimal_quantizer = DecimalQuantizer(**decimal_quantize_kwargs)
        else:
            self.decimal_quantizer = None
            
        self.set_min_max_bins(min_max_bins) 
        # min_max_bins = {
        #     'size': (0,1,256),
        #     'pos': (0,1,256),
        #     'color': (0,137,138),
        #     'font': (0,511,512)
        # }
        self.width = kwargs.get('width', 1456)
        self.height = kwargs.get('height', 1457)
        self.width = int(self.width)
        self.height = int(self.height)

    def set_min_max_bins(self, min_max_bins): # 检查 n_bins是否是偶数，然后将其 +1
        min_max_bins = copy.deepcopy(min_max_bins)
        # adjust the bins to plus one
        for type_name, (min_val, max_val, n_bins) in min_max_bins.items():
            assert n_bins % 2 == 0 # must be even
            min_max_bins[type_name] = (min_val, max_val, n_bins+1)
        self.min_max_bins = min_max_bins

    def setup_tokenizer(self, tokenizer):
        # 整个函数生成additional_special_tokens：1. '<split-text>' 2.<one-1> <symbol-1> : decimal quantizer 3. <size-255> quantizerV4 4.self.get_additional_mask_tokens()
        # 然后tokenizer.add_special_tokens({'additional_special_tokens': additional_special_tokens})
        additional_special_tokens = [self.text_split_token]  # self.text_split_token = '<split-text>'
        if self.decimal_quantize:
            special_tokens = self.decimal_quantizer.get_vocab() # <one-1> <symbol-1> 
            self.io_ignore_replace_tokens += special_tokens # self.io_ignore_replace_tokens = ['<split-text>'] 在BaseQuantizer中声明
            additional_special_tokens += special_tokens
        # the order must be preserved, other wise the tokenizer will be wrong
        rest_types = [key for key in self.min_max_bins.keys() if key not in self.decimal_quantize_types]
        for type_name in rest_types:
            min_val, max_val, n_bins = self.min_max_bins[type_name]
            additional_special_tokens += [f'<{type_name}-{i}>' for i in range(n_bins)] # <size-256>

        if self.num_mask_tokens > 0:
            additional_special_tokens.extend(self.get_additional_mask_tokens())
        
        print('additional_special_tokens', additional_special_tokens)
        
        tokenizer.add_special_tokens({'additional_special_tokens': additional_special_tokens})
        self.additional_special_tokens = set(additional_special_tokens)
        return tokenizer

    @lru_cache(maxsize=128) # 缓存函数的返回值，以提高性能。maxsize=128 表示缓存最多存储 128 个不同的输入结果
    def get_bins(self, real_type): # real_type: size, pos, font, color
        # 返回 最小值，最大值，等距数组
        min_val, max_val, n_bins = self.min_max_bins[real_type]
        return min_val, max_val, np.linspace(min_val, max_val, n_bins)

    def quantize(self, x, type): # (0.25, 'y') -> (<size-50>)
        if not self.quant:
            return x
        """Quantize a float array x into n_bins discrete values."""
        real_type = specs[type] # x, y, width, height, color, font -> size, pos, font, color
        min_val, max_val, bins = self.get_bins(real_type)
        x = np.clip(float(x), min_val, max_val) # 确保 x 的值在 [min_val, max_val] 范围内，否则截断
        if self.decimal_quantize and real_type in self.decimal_quantize_types:
            return self.decimal_quantizer(x, **pre_post_decimals[real_type])
        val = np.digitize(x, bins) - 1  # val是一个整数，取值范围在[0, len(bins)]，换句话说就是bins数组的索引
        n_bins = len(bins)
        assert val >= 0 and val < n_bins
        return f'<{real_type}-{val}>' # <size-255>
    
    def dequantize(self, x): # （<size-255> -> 0.99?)
        # <pos-1>->1
        val = x.split('-')[1].strip('>')
        # <pos-1>->pos
        real_type = x.split('-')[0][1:]
        if self.decimal_quantize and self.decimal_quantizer.check_valid(x):
            return self.decimal_quantizer.decode(x)
        min_val, max_val, bins = self.get_bins(real_type)
        return bins[int(val)]

    def construct_map_dict(self):
        map_dict = {}
        for i in range(self.min_max_bins['size'][2]): # 'size': (0, 1, 256),
            name = "<size-%d>" % i
            value = self.dequantize(name)
            map_dict[name] = str(value) # 255 -> 0.99?
        for i in range(self.min_max_bins['pos'][2]):
            name = "<pos-%d>" % i
            value = self.dequantize(name)
            map_dict[name] = str(value)
        return map_dict
    
    def postprocess_colorandfont(self, json_example):
        # 将其中的 正则 匹配部分 用双引号包裹
        import re
        json_example = re.sub(r'(<font-\d+>)', r'"\1"', json_example)
        json_example = re.sub(r'(<color-\d+>)', r'"\1"', json_example)
        return json_example
   
    def to_str(self, x, type):
        feature = self.get_feature(type)
        return feature.int2str(x)

    def convert2layout(self, example): # 将原始的数据转化为 <size-255> 的 token形式
        new_example = OrderedDict()
        new_example['wholecaption'] = example['wholecaption']
        new_layout = []
        for meta_layer in example['layout']:
            new_layout.append({
                "layer": meta_layer["layer"],
                "x": self.quantize(meta_layer["x"]/self.width, 'x'),
                "y": self.quantize(meta_layer["y"]/self.height, 'y'),
                "width": self.quantize(meta_layer["width"]/self.width, 'width'),
                "height": self.quantize(meta_layer["height"]/self.height, 'height')
            })
        new_example['layout'] = new_layout
        return new_example
    
    def apply_masking(self, 
                  json_example, 
                  mask_all=None, 
                  return_meta=False, 
                  # target_keys=['width', 'height', 'left', 'top'], # useless
                  # target_element_types=None, # useless
                  mask_values = True
                  ):
        if mask_all is None:
            mask_all = self.mask_all
        
        json_example = copy.deepcopy(json_example)
        
        # 这段内容对json中的一些 value 替换为 <mask-i>，并用self.num_mask_tokens限制mask的数量，根据参数还可能进行随机mask
        # 并记录 <mask-i> & value & num_token = value.count('<') 的 三元tuple 
        target_tokens = []
        if self.mask_values and mask_values:
            target_tokens.append((-1,-1,'globalcaption', json_example['globalcaption']))
            target_tokens.append((-1,-1,'canvas_width', json_example['canvas_width']))
            target_tokens.append((-1,-1,'canvas_height', json_example['canvas_height']))
            target_tokens.append((-1,-1,'category', json_example['category']))
            target_tokens.append((-1,-1,'keywords', json_example['keywords']))
            target_tokens.append((-1,-1,'bgcaption', json_example['layers']['bglayer']['bgcaption']))
            target_tokens.append((-1,-1,'flag', json_example['layers']['objlayer']['flag']))
            target_tokens.append((-1,-1,'objcaption', json_example['layers']['objlayer']['objcaption']))
        for layer_i, textlayer in enumerate(json_example['layers']['textlayer']):
            target_tokens.append((layer_i, -1, 'text', json_example['layers']['textlayer'][textlayer]))
        if not mask_all: # 随机取值 target_num_mask_tokens, 上界是self.num_mask_tokens
            target_num_mask_tokens = random.randint(1, self.num_mask_tokens)
            if len(target_tokens) > target_num_mask_tokens:
                random.shuffle(target_tokens)
                target_tokens = target_tokens[:target_num_mask_tokens]
                # sort by shape_i and key_i
                target_tokens = sorted(target_tokens, key=lambda x: x[0]*100+x[1])
        else: # 取定值 num_mask_tokens
            if len(target_tokens) > self.num_mask_tokens:
                # 取最后面几个
                target_tokens = target_tokens[-self.num_mask_tokens:]

        tuples = []
        meta_infos = []
        layer_list = ['heading', 'subheading', 'body']
        for mask_i, (shape_i, key_i, key, value) in enumerate(target_tokens):
            if self.mask_type == 'cm3':
                mask_token = self.mask_tokens[mask_i]
            elif self.mask_type == 'mask_aug':
                mask_token = self.mask_aug_token
            else:
                raise ValueError(f'Invalid mask type {self.mask_type}')
            # <one-1><decimal0-1><decimal1-2>
            if '<' in value:
                num_token = value.count('<')
            else:
                num_token = value.count(' ') + 1
            if shape_i == -1:
                if key in ['bgcaption']:
                    json_example['layers']['bglayer']['bgcaption'] = mask_token
                elif key in ['objcaption']:
                    json_example['layers']['objlayer']['objcaption'] = mask_token
                elif key in ['flag']:
                    json_example['layers']['objlayer']['flag'] = mask_token
                else:
                    json_example[key] = mask_token
            else:
                curlayer = layer_list[shape_i]
                json_example['layers']['textlayer'][curlayer] = mask_token
            tuples.append((mask_token, value, num_token)) 
            meta_infos.append((shape_i,key))
        if return_meta:
            return json_example, tuples, meta_infos
        else:
            return json_example, tuples


# useless orginally used for render
def is_font_exists(font_name):  
    font_list = font_manager.findSystemFonts() 
    # print("\nfont_list: ",font_list)
    for font in font_list:  
        if font_name.lower() in font.lower():  
            return True  
    return False 

def print_info(msg):  
    print(Fore.GREEN + "[INFO] " + msg)  
  
def print_warning(msg):  
    print(Fore.YELLOW + "[WARNING] " + msg)  
  
def print_error(msg):  
    print(Fore.RED + "[ERROR] " + msg)  

def load_feature(path):
    with open(path) as f:
        content = f.read()
    content = json.loads(content)
    names = [content[str(i)] for i in range(len(content))]
    return ClassLabel(num_classes= len(names), names=names)

def get_quantizer(version='v1', update_vocab=False, **kwargs):
    """ if kwargs.pop('separate_alpha', False):  # useless
        kwargs['n_visual_tokens'] *= 2 """
    if version == 'v4':
        quantizer = QuantizerV4(**kwargs)
    else:
        raise NotImplementedError

    return quantizer