Upload 3 files

update CDM:
1. support chinese formula
2. update process speed
3. update match and check code
4. fix '\n' bug

latex2bbox_color.py

@@ -0,0 +1,215 @@
+import os
+import re
+import json
+import shutil
+import logging
+import subprocess
+import numpy as np
+
+from threading import Timer
+from PIL import Image, ImageDraw
+from modules.latex_processor import (
+    normalize_latex,
+    token_add_color_RGB,
+    clean_latex
+)
+from modules.tokenize_latex.tokenize_latex import tokenize_latex
+
+
+tabular_template = r"""
+\documentclass[12pt]{article}
+\usepackage[landscape]{geometry}
+\usepackage{geometry}
+\geometry{a<PaperSize>paper,scale=0.98}
+\pagestyle{empty}
+\usepackage{booktabs}
+\usepackage{multirow}
+\usepackage{amssymb}
+\usepackage{upgreek}
+\usepackage{amsmath}
+\usepackage{xcolor}
+\begin{document}
+\makeatletter
+\renewcommand*{\@textcolor}[3]{%%
+  \protect\leavevmode
+  \begingroup
+    \color#1{#2}#3%%
+  \endgroup
+}
+\makeatother
+\begin{displaymath}
+%s
+\end{displaymath}
+\end{document}
+"""
+
+formular_template = r"""
+\documentclass[12pt]{article}
+\usepackage[landscape]{geometry}
+\usepackage{geometry}
+\geometry{a<PaperSize>paper,scale=0.98}
+\pagestyle{empty}
+\usepackage{booktabs}
+\usepackage{amsmath}
+\usepackage{upgreek}
+\usepackage{amssymb}
+\usepackage{xcolor}
+\begin{document}
+\makeatletter
+\renewcommand*{\@textcolor}[3]{%%
+  \protect\leavevmode
+  \begingroup
+    \color#1{#2}#3%%
+  \endgroup
+}
+\makeatother
+\begin{displaymath}
+%s
+\end{displaymath}
+\end{document}
+"""
+
+
+def run_cmd(cmd, timeout_sec=30):
+    proc = subprocess.Popen(cmd, shell=True)
+    kill_proc = lambda p: p.kill()
+    timer = Timer(timeout_sec, kill_proc, [proc])
+    try:
+        timer.start()
+        stdout,stderr = proc.communicate()
+    finally:
+        timer.cancel()
+        
+def convert_pdf2img(pdf_filename, png_filename):
+    cmd = "magick -density 200 -quality 100 %s %s"%(pdf_filename, png_filename)
+    os.system(cmd)
+
+def crop_image(image_path, pad=8):
+    img = Image.open(image_path).convert("L")
+    img_data = np.asarray(img, dtype=np.uint8)
+    nnz_inds = np.where(img_data!=255)
+    if len(nnz_inds[0]) == 0:
+        y_min = 0
+        y_max = 10
+        x_min = 0
+        x_max = 10
+    else:
+        y_min = np.min(nnz_inds[0])
+        y_max = np.max(nnz_inds[0])
+        x_min = np.min(nnz_inds[1])
+        x_max = np.max(nnz_inds[1])
+        
+    img = Image.open(image_path).convert("RGB").crop((x_min-pad, y_min-pad, x_max+pad, y_max+pad))
+    img.save(image_path)
+    
+def extrac_bbox_from_color_image(image_path, color_list):
+    img = Image.open(image_path).convert("RGB")
+    W, H = img.size
+    pixels = list(img.getdata())
+    
+    bbox_list = []
+    for target_color in color_list:
+        target_pixels = [ i for i, pixel in enumerate(pixels)if pixel == target_color ]
+        x_list = []
+        y_list = []
+        for idx in target_pixels:
+            x_list.append(idx % W)
+            y_list.append(idx // W)
+        try:
+            y_min, y_max, x_min, x_max = min(y_list), max(y_list), min(x_list), max(x_list)
+            bbox_list.append([x_min-1, y_min-1, x_max+1, y_max+1])
+
+        except:
+            bbox_list.append([])
+            continue
+        
+    img = img.convert("L")
+    img_bw = img.point(lambda x: 255 if x == 255 else 0, '1')
+    img_bw.convert("RGB").save(image_path) 
+    return bbox_list
+
+
+def latex2bbox_color(input_arg):
+    latex, basename, output_path, temp_dir, total_color_list = input_arg
+    template = tabular_template if "tabular" in latex else formular_template
+    output_bbox_path = os.path.join(output_path, 'bbox', basename+'.jsonl')
+    output_vis_path = os.path.join(output_path, 'vis', basename+'.png')
+    output_base_path = os.path.join(output_path, 'vis', basename+'_base.png')
+    
+    if os.path.exists(output_bbox_path) and os.path.exists(output_vis_path) and os.path.exists(output_base_path):
+        return
+    
+    try:
+        ret, new_latex = tokenize_latex(latex, middle_file=os.path.join(temp_dir, basename+'.txt'))
+        if not(ret and new_latex):
+            log = f"ERROR, Tokenize latex failed: {basename}."
+            logging.info(log)
+            new_latex = latex
+        latex = normalize_latex(new_latex)
+        token_list = []
+        l_split = latex.strip().split(' ')
+        color_list = total_color_list[0:len(l_split)]
+        idx = 0
+        while idx < len(l_split):
+            l_split, idx, token_list = token_add_color_RGB(l_split, idx, token_list)
+
+        rgb_latex = " ".join(l_split)
+        for idx, color in enumerate(color_list):
+            R, G, B = color
+            rgb_latex = rgb_latex.replace(f"<color_{idx}>", f"{R},{G},{B}")
+
+        if len(token_list) > 1300:
+            paper_size = 3
+        elif len(token_list) > 600:
+            paper_size = 4
+        else:
+            paper_size = 5
+        final_latex = formular_template.replace("<PaperSize>", str(paper_size)) % rgb_latex
+        
+    except Exception as e:
+        log = f"ERROR, Preprocess latex failed: {basename}; {e}."
+        logging.info(log)
+        return
+    
+    pre_name = output_path.replace('/', '_').replace('.','_') + '_' + basename
+    tex_filename = os.path.join(temp_dir, pre_name+'.tex')
+    log_filename = os.path.join(temp_dir, pre_name+'.log')
+    aux_filename = os.path.join(temp_dir, pre_name+'.aux')
+    
+    with open(tex_filename, "w") as w: 
+        print(final_latex, file=w)
+    run_cmd(f"pdflatex -interaction=nonstopmode -output-directory={temp_dir} {tex_filename} >/dev/null")
+    try:
+        os.remove(tex_filename)
+        os.remove(log_filename)
+        os.remove(aux_filename)
+    except:
+        pass
+    pdf_filename = tex_filename[:-4]+'.pdf'
+    if not os.path.exists(pdf_filename):
+        log = f"ERROR, Compile pdf failed: {pdf_filename}"
+        logging.info(log)
+    else:
+        convert_pdf2img(pdf_filename, output_base_path)
+        os.remove(pdf_filename)
+        
+        crop_image(output_base_path)
+        bbox_list = extrac_bbox_from_color_image(output_base_path, color_list)
+        vis = Image.open(output_base_path)
+        draw = ImageDraw.Draw(vis)
+
+        with open(output_bbox_path, 'w') as f:
+            for token, box in zip(token_list, bbox_list):
+                item = {
+                    "bbox": box,
+                    "token": token
+                }
+                f.write(json.dumps(item)+'\n')
+
+                if not box:
+                    continue
+                x_min, y_min, x_max, y_max = box
+                draw.rectangle([x_min, y_min, x_max, y_max], fill=None, outline=(0,250,0), width=1)
+                draw.text((x_min, y_min), token, (250,0,0))
+            
+        vis.save(output_vis_path)

latex_processor.py

@@ -0,0 +1,536 @@
+import os
+import re
+import json
+import shutil
+import logging
+import numpy as np
+from PIL import Image
+
+
+SKIP_PATTERNS = [r'\{', r'\}', r'[\[\]]', r'\\begin\{.*?\}', r'\\end\{.*?\}', r'\^', r'\_', r'\\.*rule.*', r'\\.*line.*', r'\[[\-.0-9]+[epm][xtm]\]']
+SKIP_Tokens = ['\\', '\\\\', '\\index', '\\a', '&', '$', '\\multirow', '\\def', '\\raggedright', '\\url', '\\cr', '\\ensuremath', '\\left', '\\right', 
+               '\\mathchoice', '\\scriptstyle', '\\displaystyle', '\\qquad', '\\quad', '\\,', '\\!', '~', '\\boldmath']
+PHANTOM_Tokens = ['\\fontfamily', '\\vphantom', '\\phantom', '\\rowcolor', '\\ref']
+TWO_Tail_Tokens = ['\\frac', '\\binom']
+AB_Tail_Tokens = ['\\xrightarrow', '\\xleftarrow', '\\sqrt']        # special token \xxx [] {} 
+TWO_Tail_Invisb_Tokens = ['\\overset', '\\underset', '\\stackrel']
+ONE_Tail_Tokens = ['\\widetilde', '\\overline', '\\hat', '\\widehat', '\\tilde', '\\Tilde', '\\dot', '\\bar', '\\vec', '\\underline', '\\underbrace', '\\check',
+                   '\\breve', '\\Bar', '\\Vec', '\\mathring', '\\ddot']
+ONE_Tail_Invisb_Tokens = ['\\boldsymbol', '\\pmb', '\\textbf', '\\mathrm', '\\mathbf', '\\mathbb', '\\mathcal', '\\textmd', '\\texttt', '\\textnormal', 
+                          '\\text', '\\textit', '\\textup', '\\mathop', '\\mathbin', '\\smash', '\\operatorname', '\\textrm', '\\mathfrak', '\\emph',
+                          '\\textsf', '\\textsc']
+
+
+def flatten_multiline(latex):
+    brace_map = {
+        "\\left(": "\\right)",
+        "\\left[": "\\right]",
+        "\\left{": "\\right}",
+    }
+    l_split = latex.split(' ')
+    if l_split[0] == "\\begin{array}":
+        if l_split[-1] == "\\end{array}":
+            l_split = l_split[2:-1]
+        else:
+            l_split = l_split[2:]
+    
+    idx = 0
+    while idx < len(l_split):
+        token = l_split[idx]
+        if token.startswith("\\left") and token in brace_map.keys():
+            end_idx = find_matching_brace(l_split, idx, brace=[token, brace_map[token]])
+            if end_idx != -1:
+                idx = end_idx
+        elif token in ["\\\\", "~", "\\qquad"]:
+            l_split = l_split[0:idx] + l_split[idx+1:]
+            idx -= 1
+        idx += 1
+    latex = ' '.join(l_split)
+    return "$ "+latex+" $"
+    
+    
+def clean_latex(text):
+    # TODO 让GPT写的去空格函数, 初步测了是没问题的, 不确定是否完全没有bug
+    cleaned_text = re.sub(r'(?<=[^\\])\s+(?=[^\\])', '', text)
+    # TODO 有一些不能去掉的空格给补充回来
+    for item in ["\\hline", "\\midrule", "\\times", "\\bf", "\\footnotesize", "\\cr", '\\log']:
+        cleaned_text = cleaned_text.replace(item, item+" ")
+    cleaned_text = cleaned_text.replace(" \\mathcolor{black}", "\\mathcolor{black}")
+    return cleaned_text
+
+def remove_trailing_latex(formula):
+    pattern = r'(\\(hspace\*?\{[^{}]*?\}|vspace\*?\{[^{}]*?\}|smallskip|medskip|quad|qquad|bigskip|[;,])|\~|\.)*$'
+    # Replace the matched pattern with an empty string
+    cleaned_formula = re.sub(pattern, '', formula, count=1)
+    return cleaned_formula
+
+def find_matching_brace(sequence, start_index, brace=['{', '}']):
+    # Finds the index of the matching brace for the one at start_index
+    left_brace, right_brace = brace
+    depth = 0
+    for i, char in enumerate(sequence[start_index:], start=start_index):
+        if char == left_brace:
+            depth += 1
+        elif char == right_brace:
+            depth -= 1
+            if depth == 0:
+                return i
+    if depth > 0:
+        error_info = "Warning! found no matching brace in sequence !"
+        raise ValueError(error_info)
+    return -1
+
+def normalize_latex(l, rm_trail=False):
+    if "tabular" in l:
+        latex_type = "tabular"
+    else:
+        latex_type = "formula"
+        
+    if rm_trail:
+        l = remove_trailing_latex(l)
+    l = l.strip().replace(r'\pmatrix', r'\mypmatrix').replace(r'\matrix', r'\mymatrix')
+    
+    # TODO \raggedright \arraybackslash, these align method, difficult to handle, remove it.
+    for item in ['\\raggedright', '\\arraybackslash']:
+        l = l.replace(item, "")
+    
+    for item in ['\\lowercase', '\\uppercase']:
+        l = l.replace(item, "")
+        
+    # TODO \hspace {1 . 5 cm}, for formula, change to \hspace{1.5cm}, for table, remove it.
+    pattern = r'\\[hv]space { [.0-9a-z ]+ }'
+    old_token = re.findall(pattern, l, re.DOTALL)
+    if latex_type == "tabular":
+        new_token = ["" for item in old_token]
+    else:
+        new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft)
+        
+    # TODO take \begin {tabular} {} as one token
+    # TODO there are \begin{array} in table too，so the process should run in both formula and table.
+    if latex_type == "tabular":
+        l = l.replace("\\begin {tabular}", "\\begin{tabular}")
+        l = l.replace("\\end {tabular}", "\\end{tabular}")
+        l = l.replace("\\begin {array}", "\\begin{array}")
+        l = l.replace("\\end {array}", "\\end{array}")
+        l_split = l.split(' ')
+        idx = 0
+        while idx < len(l_split):
+            token = l_split[idx]
+            if token == "\\begin{tabular}":
+                sub_idx = idx + 1
+                end_idx = find_matching_brace(l_split, sub_idx)
+                new_token = "".join(l_split[idx: end_idx+1])
+                l_split = l_split[0:idx] + [new_token] + l_split[end_idx+1:]
+                break
+            idx += 1
+        l = ' '.join(l_split)
+        
+        # TODO some complex format, hart to deal with re.match, so using brace match, such as：\cmidrule ( l { 3 p t } r { 3 p t } ) { 1 - 1 }
+        l_split = l.split(' ')
+        idx = 0
+        while idx < len(l_split):
+            token = l_split[idx]
+            if token in ["\\cmidrule", "\\cline"]:
+                sub_idx = idx + 1
+                if l_split[sub_idx] == "(":
+                    mid_end = find_matching_brace(l_split, sub_idx, brace=['(', ')'])
+                    end_idx = find_matching_brace(l_split, mid_end+1)
+                else:
+                    end_idx = find_matching_brace(l_split, sub_idx)
+                new_token = "".join(l_split[idx: end_idx+1])
+                l_split = l_split[0:idx] + [new_token] + l_split[end_idx+1:]
+            idx += 1
+        l = ' '.join(l_split)
+    
+    pattern = r'\\begin{array} { [lrc ]+ }'
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace("\\begin{array} ", "<s>").replace(" ", "").replace("<s>", "\\begin{array} ") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft)
+    
+    # TODO token such \not= should be one token
+    pattern = r'\\not [<>+=\-]'
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft)
+        
+    # TODO tokens such as \dots \exp \sinh, split them to parts, so the bbox match will be easier.
+    
+    l = " "+l+" "
+    l = l.replace(" \\ldots ", " . . . ")
+    l = l.replace(" \\cdots ", " . . . ")
+    l = l.replace(" \\dots ", " . . . ")
+    l = l.replace(" \\dotsb ", " . . . ")
+    l = l.replace(" \\log ", " \\mathrm { l o g } ")
+    l = l.replace(" \\exp ", " \\mathrm { e x p } ")
+    l = l.replace(" \\sin ", " \\mathrm { s i n } ")
+    l = l.replace(" \\cos ", " \\mathrm { c o s } ")
+    l = l.replace(" \\tan ", " \\mathrm { t a n } ")
+    l = l.replace(" \\tanh ", " \\mathrm { t a n h } ")
+    l = l.replace(" \\cosh ", " \\mathrm { c o s h } ")
+    l = l.replace(" \\sinh ", " \\mathrm { s i n h } ")
+        
+    # ** token such as \big( should be one token
+    pattern = r'\\[Bb]ig[g]?[glrm]? [(){}|\[\]] '
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft+" ")
+        
+    pattern = r'\\[Bb]ig[g]?[glrm]? \\.*? '
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft+" ")
+        
+    # TODO when \operatorname * meets mathcolor it comes error, yet the * is useless, so we simply remove it bynow.
+    pattern = r'\\operatorname \*'
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = ["\\operatorname" for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft)
+    
+    # TODO \lefteqn will lead to letter overlap, it's harmfull for render, so simply remove it.
+    l = l.replace("\\lefteqn", "")
+    
+    # TODO \footnote can not seem as ONE_Tail_Invisb_Tokens(usually this type token add color by \mathrm {\color(x)}, yet \footnode should be \color{\footnote{x}}), so we simple change it to "^".
+    l = l.replace("\\footnote ", "^ ")
+    
+    # TODO \' can not be rendered separately(cause to different visulize performence), so we take these tokens as one token such as \' e -> \'e, on the other hand, if { after \' then render them separately.
+    pattern = r'\\\' [^{] '
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft+" ")
+    
+    # TODO [ -1.5ex ] [ 1.5pt ] [ 3 mm ] some layout adjustment, no need to render. combine them as one token.
+    if latex_type == "tabular":
+        pattern = r'\[ [\-.0-9 ]+[exptcm ]+ \]'
+        old_token = re.findall(pattern, l, re.DOTALL)
+        new_token = [item.replace(" ", "") for item in old_token]
+        for bef, aft in zip(old_token, new_token):
+            l = l.replace(bef, aft)
+    
+    # ** \parbox { 3cm } {} shoudle be combined as one token
+    pattern = r'\\parbox {[^{]+}'
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft)
+    
+    # ** \raisebox{<lift>}[<height>][<depth>] {} shoudle be combined as one token, \raisebox{-1.5ex}[0pt]
+    pattern = r'\\raisebox {[^{]+} [\[\]0-9 exptcm]+{'
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft[0:-1]+" {")
+        
+    # ** \char shoudle be combined as one token
+    pattern = r'{ \\char[0-9\' ]+}'
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, "{ "+aft[1:-1]+" }")
+        
+    # ** \not xx shoudle be combined as one token
+    pattern = r'\\not [\\=\<\>][^ ]+ '
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft+" ")
+        
+    # ** \specialrule{1pt}{2pt}{2pt}, special lines, shoudle be combined as one token
+    pattern = r'\\specialrule {[ .0-9a-z]+} {[ .0-9a-z]+} {[ .0-9a-z]+}'
+    old_token = re.findall(pattern, l, re.DOTALL)
+    new_token = [item.replace(" ", "") for item in old_token]
+    for bef, aft in zip(old_token, new_token):
+        l = l.replace(bef, aft)
+        
+    # ** for easier add color, the original color should be removed, there are two type of color for now: \color[rgb]{0, 1, 0} and \color{red}
+    pattern = r'\\colorbox[ \[\]RGBrgb]+{ [A-Za-z 0-9,!]+ } |\\color[ \[\]RGBrgb]+{ [A-Za-z 0-9,!]+ } |\\textcolor[ \[\]RGBrgb]+{ [A-Za-z 0-9,!]+ } |\\cellcolor[ \[\]RGBrgb]+{ [A-Za-z 0-9,!]+ } '
+    old_token = re.findall(pattern, l, re.DOTALL)
+    for bef in old_token:
+        l = l.replace(bef, "")
+    
+    # ** filling the missing brace [] and {} according to token.
+    l_split = l.split(' ')
+    idx = 0
+    while idx < len(l_split):
+        token = l_split[idx]
+        if token in ONE_Tail_Tokens + ONE_Tail_Invisb_Tokens:
+        # ** normalize tokens such as \hat, fill missing the {}, such as \hat \lambda -> \hat {\lambda}
+            sub_idx = idx + 1
+            while sub_idx < len(l_split) and l_split[sub_idx] in ONE_Tail_Tokens+ONE_Tail_Invisb_Tokens:
+                sub_idx += 1
+            new_split = l_split[0:idx]
+            for ii in range(idx, sub_idx):
+                new_split = new_split + [l_split[ii], "{"]
+            if l_split[sub_idx] != "{":
+                new_split = new_split + [l_split[sub_idx]] + ["}"]*(sub_idx-idx)
+                l_split = new_split + l_split[sub_idx+1:]
+            else:
+                end_idx = find_matching_brace(l_split, sub_idx)
+                new_split = new_split + l_split[sub_idx+1:end_idx] + ["}"]*(sub_idx-idx)
+                l_split = new_split + l_split[end_idx+1:]
+        elif token in AB_Tail_Tokens:
+        # ** normalize special tokens such as \sqrt, fill the missing [] {} in \sqrt [] {}, yet the [] is optional, for example: \sqrt A B -> \sqrt {A} B and \sqrt [A] B -> \sqrt [A] {B}
+            if l_split[idx + 1] != "[" and l_split[idx + 1] != "{":
+                l_split = l_split[0:idx+1] + ["{"] + [l_split[idx+1]] + ["}"] + l_split[idx+2:]
+            else:
+                if l_split[idx + 1] == "[":
+                    end1 = find_matching_brace(l_split, idx+1, brace=['[', ']'])
+                else:
+                    end1 = idx
+                if l_split[end1 + 1] != "{":
+                    l_split = l_split[0:end1+1] + ["{"] + [l_split[end1+1]] + ["}"] + l_split[end1+2:]
+        elif token in TWO_Tail_Tokens + TWO_Tail_Invisb_Tokens:
+        # ** normalize special tokens such as \frac, add missing brace in \frac {A} {B} for example: \frac {\lambda} 2 -> \frac {\lambda} {2}
+            if l_split[idx + 1] != "{":
+                l_split = l_split[0:idx+1] + ["{"] + [l_split[idx+1]] + ["}"] + l_split[idx+2:]
+            end1 = find_matching_brace(l_split, idx+1)
+            if l_split[end1 + 1] != "{":
+                l_split = l_split[0:end1+1] + ["{"] + [l_split[end1+1]] + ["}"] + l_split[end1+2:]
+            
+        idx += 1
+    l = ' '.join(l_split)
+    
+    return l
+
+def token_add_color(l_split, idx, render_dict):
+    token = l_split[idx]
+    if token in PHANTOM_Tokens:
+        # ** special tokens that do not need render, skip it 
+        if l_split[idx + 1] == '{':
+            brace_end = find_matching_brace(l_split, idx + 1)
+        else:
+            brace_end = idx + 1
+        next_idx = brace_end + 1
+    elif token in TWO_Tail_Tokens:
+        # ** tokens such as \frac A B, and the token needs render too.
+        num_start = idx + 1
+        num_end = find_matching_brace(l_split, num_start)
+        den_start = num_end + 1
+        den_end = find_matching_brace(l_split, den_start)
+        l_split_copy = l_split[:idx] + [r'\mathcolor{black}{'+token+'{'] + \
+                        [r'\mathcolor{gray}{'] + l_split[num_start + 1:num_end] + \
+                        ['}'] + [r'}{'] + [r'\mathcolor{gray}{'] + l_split[den_start + 1:den_end] + \
+                        ['}'] + ['}'] + ['}'] + l_split[den_end + 1:]
+                        
+        l_new = ' '.join(l_split_copy)
+        l_new = r'\mathcolor{gray}{ ' + l_new + ' }'
+        render_dict[str(idx)] = l_new, token
+        next_idx = idx + 1
+    elif token in ONE_Tail_Tokens:
+        # ** tokens such as \hat A, and the token needs render too.
+        num_start = idx + 1
+        num_end = find_matching_brace(l_split, num_start)
+        l_split_copy = l_split[:idx] + [r'\mathcolor{black}{'] + l_split[idx: num_start+1] + \
+                        [r'\mathcolor{gray}{'] + l_split[num_start+1: num_end] + \
+                        ['}'] + l_split[num_end: num_end+1] + ['}'] + l_split[num_end+1:]
+        l_new = ' '.join(l_split_copy)
+        l_new = r'\mathcolor{gray}{ ' + l_new + ' }'
+        render_dict[str(idx)] = l_new, token
+        next_idx = idx + 1
+    elif token in ONE_Tail_Invisb_Tokens:
+        # ** tokens such as \text A B, and the token does not need render.
+        num_start = idx + 1
+        num_end = find_matching_brace(l_split, num_start)
+        sub_idx = num_start+1
+        if num_end-num_start == 2:
+            l_split_copy = l_split.copy()
+            l_split_copy[sub_idx] = r'{\mathcolor{black}{' + l_split_copy[sub_idx] + '}}'
+            l_new = ' '.join(l_split_copy)
+            l_new = r'\mathcolor{gray}{ ' + l_new + ' }'
+            render_dict[str(idx)] = l_new, l_split[sub_idx]
+            next_idx = num_end
+        else:
+            while sub_idx < num_end:
+                l_split, sub_idx, render_dict = token_add_color(l_split, sub_idx, render_dict)
+        next_idx = num_end + 1
+    elif token in AB_Tail_Tokens:
+        # ** special token \xrightarrow, could be \xrightarrow [] {} or \xrightarrow {}, process method are different.
+        if l_split[idx+1] == '{':
+            num_start = idx + 1
+            num_end = find_matching_brace(l_split, num_start)
+            l_split_copy = l_split[:idx] + [r'\mathcolor{black}{'] + l_split[idx: idx+2] \
+                        + [r'\mathcolor{gray}{'] + l_split[num_start+1: num_end] + ['}}'] + l_split[num_end:]
+            l_new = ' '.join(l_split_copy)
+            l_new = r'\mathcolor{gray}{ ' + l_new + ' }'
+            render_dict[str(idx)] = l_new, token
+            sub_idx = num_start+1
+            while sub_idx < num_end:
+                l_split, sub_idx, render_dict = token_add_color(l_split, sub_idx, render_dict)
+            next_idx = num_end + 1
+        elif l_split[idx+1] == '[':
+            num_start = idx + 1
+            num_end = find_matching_brace(l_split, num_start, brace=['[', ']'])
+            den_start = num_end + 1
+            den_end = find_matching_brace(l_split, den_start)
+            l_split_copy = l_split[:idx] + [r'{\mathcolor{black}{'] + l_split[idx: idx+2] \
+                        + [r'\mathcolor{gray}{'] + l_split[idx+2: num_end] + ['}'] + l_split[num_end:den_start+1] \
+                        + [r'\mathcolor{gray}{'] + l_split[den_start+1: den_end] + ['}'] + l_split[den_end: den_end+1] \
+                        + ['}}'] + l_split[den_end+1:]
+            l_new = ' '.join(l_split_copy)
+            l_new = r'\mathcolor{gray}{ ' + l_new + ' }'
+            render_dict[str(idx)] = l_new, token
+            sub_idx = num_start + 1
+            while sub_idx < num_end:
+                l_split, sub_idx, render_dict = token_add_color(l_split, sub_idx, render_dict)
+            sub_idx = den_start + 1
+            while sub_idx < den_end:
+                l_split, sub_idx, render_dict = token_add_color(l_split, sub_idx, render_dict)
+            next_idx = den_end + 1
+    elif token in ["\\multicolumn", "\\multirow"]:
+        # ** tokens with three {}, such as \multicolumn {} {} {}, the text in third {} need be rendered.
+        first_start = idx + 1
+        first_end = find_matching_brace(l_split, first_start)
+        second_start = first_end + 1
+        second_end = find_matching_brace(l_split, second_start)
+        third_start = second_end + 1
+        third_end = find_matching_brace(l_split, third_start)
+        
+        sub_idx = third_start+1
+        while sub_idx < third_end:
+            l_split, sub_idx, render_dict = token_add_color(l_split, sub_idx, render_dict)
+        next_idx = third_end + 1
+    elif token in SKIP_Tokens+TWO_Tail_Invisb_Tokens or any(re.match(pattern, token) for pattern in SKIP_PATTERNS):
+        # ** tokens no need render, just skip
+        # print('skip', idx, token)
+        # TODO special case :[], could be single, or in \sqrt[]{}.
+        if (token == "[" and l_split[idx-1]!="\\sqrt") or (token == "]" and idx>=3 and l_split[idx-3]!="\\sqrt"):
+            l_split_copy = l_split.copy()
+            l_split_copy[idx] = r'\mathcolor{black}{ ' + l_split_copy[idx] + ' }'
+            l_new = ' '.join(l_split_copy)
+            l_new = r'\mathcolor{gray}{ ' + l_new + ' }'
+            render_dict[str(idx)] = l_new, token
+            next_idx = idx + 1
+        else:
+            next_idx = idx + 1
+    else:
+        # ** nomal token
+        l_split_copy = l_split.copy()
+        # TODO sometimes there is translation after add color, the exp prove that \mathcolor{black}{ A } is better than \mathcolor{black}{A}
+        l_split_copy[idx] = r'\mathcolor{black}{ ' + l_split_copy[idx] + ' }'
+
+        l_new = ' '.join(l_split_copy)
+        l_new = r'\mathcolor{gray}{ ' + l_new + ' }'
+        render_dict[str(idx)] = l_new, token
+        next_idx = idx + 1
+        
+    return l_split, next_idx, render_dict
+
+
+def token_add_color_RGB(l_split, idx, token_list, brace_color=False):
+    """using \mathcolor[RGB]{r,g,b} to render latex. 
+    """
+    token = l_split[idx]
+    if not token:
+        next_idx = idx + 1
+    elif token in PHANTOM_Tokens:
+        # ** special tokens that do not need render, skip it 
+        if l_split[idx + 1] == '{':
+            brace_end = find_matching_brace(l_split, idx + 1)
+        else:
+            brace_end = idx + 1
+        next_idx = brace_end + 1
+    elif token in TWO_Tail_Tokens:
+        # ** tokens such as \frac A B, and the token needs render too.
+        num_start = idx + 1
+        num_end = find_matching_brace(l_split, num_start)
+        den_start = num_end + 1
+        den_end = find_matching_brace(l_split, den_start)
+        color_token = "\\mathcolor[RGB]{<color_<idx>>}{".replace("<idx>", str(len(token_list)))
+        l_split = l_split[:idx] + [color_token+token] + l_split[idx+1: den_end+1] + ["}"] + l_split[den_end+1:]
+        token_list.append(token)
+        next_idx = idx + 1
+    elif token in ONE_Tail_Tokens:
+        # ** tokens such as \hat A, and the token needs render too.
+        num_start = idx + 1
+        num_end = find_matching_brace(l_split, num_start)
+        color_token = "\\mathcolor[RGB]{<color_<idx>>}{".replace("<idx>", str(len(token_list)))
+        if token != "\\underbrace" and num_end+1 < len(l_split) and l_split[num_end+1] == "_":
+            l_split = l_split[:idx] + ["{"+color_token+token] + l_split[idx+1: num_end+1] + ["}}"] + l_split[num_end+1:]
+        else:
+            l_split = l_split[:idx] + [color_token+token] + l_split[idx+1: num_end+1] + ["}"] + l_split[num_end+1:]
+        token_list.append(token)
+        next_idx = idx + 1
+    elif token in ONE_Tail_Invisb_Tokens:
+        # ** tokens such as \text A B, and the token does not need render.
+        num_start = idx + 1
+        num_end = find_matching_brace(l_split, num_start)
+        sub_idx = num_start+1
+        if num_end-num_start == 2:
+            color_token = "\\mathcolor[RGB]{<color_<idx>>}{".replace("<idx>", str(len(token_list)))
+            token_list.append(l_split[num_start+1])
+            l_split = l_split[:num_start+1] + [color_token+l_split[num_start+1]+"}"] + l_split[num_end:]
+        else:
+            while sub_idx < num_end:
+                l_split, sub_idx, token_list = token_add_color_RGB(l_split, sub_idx, token_list)
+        next_idx = num_end + 1
+    elif token in AB_Tail_Tokens:
+        # ** special token \xrightarrow, could be \xrightarrow [] {} or \xrightarrow {}, process method are different.
+        if l_split[idx+1] == '{':
+            num_start = idx + 1
+            num_end = find_matching_brace(l_split, num_start)
+            color_token = "\\mathcolor[RGB]{<color_<idx>>}{".replace("<idx>", str(len(token_list)))
+            l_split = l_split[:idx] + [color_token+token] + l_split[idx+1: num_end+1] + ["}"] + l_split[num_end+1:]
+            token_list.append(token)
+            sub_idx = num_start+1
+            while sub_idx < num_end:
+                l_split, sub_idx, token_list = token_add_color_RGB(l_split, sub_idx, token_list)
+            next_idx = num_end + 1
+        elif l_split[idx+1] == '[':
+            num_start = idx + 1
+            num_end = find_matching_brace(l_split, num_start, brace=['[', ']'])
+            den_start = num_end + 1
+            den_end = find_matching_brace(l_split, den_start)
+            color_token = "\\mathcolor[RGB]{<color_<idx>>}{".replace("<idx>", str(len(token_list)))
+            l_split = l_split[:idx] + [color_token+token] + l_split[idx+1: den_end+1] + ["}"] + l_split[den_end+1:]
+            token_list.append(token)
+            sub_idx = num_start + 1
+            while sub_idx < num_end:
+                l_split, sub_idx, token_list = token_add_color_RGB(l_split, sub_idx, token_list, brace_color=True)
+            sub_idx = den_start + 1
+            while sub_idx < den_end:
+                l_split, sub_idx, token_list = token_add_color_RGB(l_split, sub_idx, token_list)
+            next_idx = den_end + 1
+    elif token in ["\\multicolumn", "\\multirow"]:
+        # ** tokens with three {}, such as \multicolumn {} {} {}, the text in third {} need be rendered.
+        first_start = idx + 1
+        first_end = find_matching_brace(l_split, first_start)
+        second_start = first_end + 1
+        second_end = find_matching_brace(l_split, second_start)
+        third_start = second_end + 1
+        third_end = find_matching_brace(l_split, third_start)
+        
+        sub_idx = third_start+1
+        while sub_idx < third_end:
+            l_split, sub_idx, token_list = token_add_color_RGB(l_split, sub_idx, token_list)
+        next_idx = third_end + 1
+    elif token in SKIP_Tokens+TWO_Tail_Invisb_Tokens or any(re.match(pattern, token) for pattern in SKIP_PATTERNS):
+        # ** tokens no need render, just skip
+        # print('skip', idx, token)
+        # TODO special case :[], could be single, or in \sqrt[]{}.
+        if (token == "[" and l_split[idx-1]!="\\sqrt") or (token == "]" and idx>=3 and l_split[idx-3]!="\\sqrt"):
+            color_token = "\\mathcolor[RGB]{<color_<idx>>}{".replace("<idx>", str(len(token_list)))
+            l_split = l_split[:idx] + [color_token + l_split[idx] + "}"] + l_split[idx+1:]
+            token_list.append(token)
+            next_idx = idx + 1
+        else:
+            next_idx = idx + 1
+    else:
+        # ** nomal token
+        if brace_color or (idx > 1 and l_split[idx-1] == "_"):
+            color_token = "\\mathcolor[RGB]{<color_<idx>>}{".replace("<idx>", str(len(token_list)))
+            l_split = l_split[:idx] + ["{" + color_token + l_split[idx] + "}}"] + l_split[idx+1:]
+            token_list.append(token)
+            next_idx = idx + 1
+        else:
+            color_token = "\\mathcolor[RGB]{<color_<idx>>}{".replace("<idx>", str(len(token_list)))
+            l_split = l_split[:idx] + [color_token + l_split[idx] + "}"] + l_split[idx+1:]
+            token_list.append(token)
+            next_idx = idx + 1
+    return l_split, next_idx, token_list

visual_matcher.py

@@ -0,0 +1,191 @@
+import time
+import numpy as np
+from PIL import Image
+from scipy.spatial.distance import cdist
+from scipy.optimize import linear_sum_assignment
+
+
+class SimpleAffineTransform:
+    """
+    simple affine transform, only translation and scale.
+    """
+    def __init__(self, translation=(0, 0), scale=1.0):
+        self.translation = np.array(translation)
+        self.scale = scale
+
+    def estimate(self, src, dst):
+        src_center = np.mean(src, axis=0)
+        dst_center = np.mean(dst, axis=0)
+        self.translation = dst_center - src_center
+
+        src_dists = np.linalg.norm(src - src_center, axis=1)
+        dst_dists = np.linalg.norm(dst - dst_center, axis=1)
+        self.scale = np.mean(dst_dists) / (np.mean(src_dists) + 1e-10)
+
+    def inverse(self):
+        inverse_transform = AffineTransform(-self.translation, 1.0/self.scale)
+        return inverse_transform
+
+    def __call__(self, coords):
+        return self.scale * (coords - np.mean(coords, axis=0)) + np.mean(coords, axis=0) + self.translation
+
+    def residuals(self, src, dst):
+        return np.sqrt(np.sum((self(src) - dst) ** 2, axis=1))
+
+
+def norm_coords(x, left, right):
+    if x < left:
+        return left
+    if x > right:
+        return right
+    return x
+
+def norm_same_token(token):
+    special_map = {
+        "\\cdot": ".",
+        "\\mid": "|",
+        "\\to": "\\rightarrow",
+        "\\top": "T",
+        "\\Tilde": "\\tilde",
+        "\\cdots": "\\dots",
+        "\\prime": "'",
+        "\\ast": "*",
+        "\\left<": "\\langle",
+        "\\right>": "\\rangle"
+    }
+    if token in special_map.keys():
+        token = special_map[token]
+    if token.startswith('\\left') or token.startswith('\\right'):
+        token = token.replace("\\left", "").replace("\\right", "")
+    if token.startswith('\\big') or token.startswith('\\Big'):
+        if "\\" in token[4:]:
+            token = "\\"+token[4:].split("\\")[-1]
+        else:
+            token = token[-1]
+    
+    if token in ['\\leq', '\\geq']:
+        return token[0:-1]
+    if token in ['\\lVert', '\\rVert', '\\Vert']:
+        return '\\|'
+    if token in ['\\lvert', '\\rvert', '\\vert']:
+        return '|'
+    if token.endswith("rightarrow"):
+        return "\\rightarrow"
+    if token.endswith("leftarrow"):
+        return "\\leftarrow"
+    if token.startswith('\\wide'):
+        return token.replace("wide", "")
+    if token.startswith('\\var'):
+        return token.replace("\\var", "")
+    return token
+
+
+class HungarianMatcher:
+    def __init__(
+        self, 
+        cost_token: float = 1,
+        cost_position: float = 0.05,
+        cost_order: float = 0.15,
+    ):
+        self.cost_token = cost_token
+        self.cost_position = cost_position
+        self.cost_order = cost_order
+        self.cost = {}
+    
+    def calculate_token_cost_old(self, box_gt, box_pred):
+        token_cost = np.ones((len(box_gt), len(box_pred)))
+        for i in range(token_cost.shape[0]):
+            box1 = box_gt[i]
+            for j in range(token_cost.shape[1]):
+                box2 = box_pred[j]
+                if box1['token'] == box2['token']:
+                    token_cost[i, j] = 0
+                elif norm_same_token(box1['token']) == norm_same_token(box2['token']):
+                    token_cost[i, j] = 0.05
+        return np.array(token_cost)
+        
+    def calculate_token_cost(self, box_gt, box_pred):
+        token2id = {}
+        for data in box_gt+box_pred:
+            if data['token'] not in token2id:
+                token2id[data['token']] = len(token2id)
+        num_classes = len(token2id)
+        
+        token2id_norm = {}
+        for data in box_gt+box_pred:
+            if norm_same_token(data['token']) not in token2id_norm:
+                token2id_norm[norm_same_token(data['token'])] = len(token2id_norm)
+        num_classes_norm = len(token2id_norm)
+        
+        gt_token_array = []
+        norm_gt_token_array = []    
+        for data in box_gt:
+            gt_token_array.append(token2id[data['token']])
+            norm_gt_token_array.append(token2id_norm[norm_same_token(data['token'])])
+            
+        pred_token_logits = []
+        norm_pred_token_logits = []
+        for data in box_pred:
+            logits = [0] * num_classes
+            logits[token2id[data['token']]] = 1
+            pred_token_logits.append(logits)
+            
+            logits_norm = [0] * num_classes_norm
+            logits_norm[token2id_norm[norm_same_token(data['token'])]] = 1
+            norm_pred_token_logits.append(logits_norm)
+            
+        gt_token_array = np.array(gt_token_array)
+        pred_token_logits = np.array(pred_token_logits)
+        
+        norm_gt_token_array = np.array(norm_gt_token_array)
+        norm_pred_token_logits = np.array(norm_pred_token_logits)
+        
+        token_cost = 1.0 - pred_token_logits[:, gt_token_array]
+        norm_token_cost = 1.0 - norm_pred_token_logits[:, norm_gt_token_array]
+
+        token_cost[np.logical_and(token_cost==1, norm_token_cost==0)] = 0.05
+        return token_cost.T
+        
+        
+    def box2array(self, box_list, size):
+        W, H = size
+        box_array = []
+        for box in box_list:
+            x_min, y_min, x_max, y_max = box['bbox']
+            box_array.append([x_min/W, y_min/H, x_max/W, y_max/H])
+        return np.array(box_array)
+        
+    def order2array(self, box_list):
+        order_array = []
+        for idx, box in enumerate(box_list):
+            order_array.append([idx / len(box_list)])
+        return np.array(order_array)
+    
+    def calculate_l1_cost(self, gt_array, pred_array):
+        scale = gt_array.shape[-1]
+        l1_cost = cdist(gt_array, pred_array, 'minkowski', p=1)
+        return l1_cost / scale
+        
+    def __call__(self, box_gt, box_pred, gt_size, pred_size):
+        aa = time.time()
+        gt_box_array = self.box2array(box_gt, gt_size)
+        pred_box_array = self.box2array(box_pred, pred_size)
+        gt_order_array = self.order2array(box_gt)
+        pred_order_array = self.order2array(box_pred)
+
+        token_cost = self.calculate_token_cost(box_gt, box_pred)
+        position_cost = self.calculate_l1_cost(gt_box_array, pred_box_array)
+        order_cost = self.calculate_l1_cost(gt_order_array, pred_order_array)
+
+        self.cost["token"] = token_cost
+        self.cost["position"] = position_cost
+        self.cost["order"] = order_cost
+        
+        cost = self.cost_token * token_cost + self.cost_position * position_cost + self.cost_order * order_cost
+        cost[np.isnan(cost) | np.isinf(cost)] = 100
+        indexes = linear_sum_assignment(cost)
+        matched_idxes = []
+        for a, b in zip(*indexes):
+            matched_idxes.append((a, b))
+        
+        return matched_idxes