Upload 3 files

modules/latex2bbox_color.py

@@ -1,5 +1,6 @@
 import os
 import re
+import cv2
 import json
 import shutil
 import logging
@@ -69,6 +70,33 @@ formular_template = r"""
 \end{document}
 """
 
+formular_template_zh = 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{CJK}
+\usepackage{amssymb}
+\usepackage{xcolor}
+\begin{document}
+\makeatletter
+\renewcommand*{\@textcolor}[3]{%%
+  \protect\leavevmode
+  \begingroup
+    \color#1{#2}#3%%
+  \endgroup
+}
+\makeatother
+\begin{CJK}{UTF8}{gkai}
+%s
+\end{CJK}
+\end{document}
+"""
+
 
 def run_cmd(cmd, timeout_sec=30):
     proc = subprocess.Popen(cmd, shell=True)
@@ -101,37 +129,41 @@ def crop_image(image_path, pad=8):
         
     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())
-    
+    img = cv2.imread(image_path)
     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:
+        r, g, b = target_color
+        target_rgb = np.array([b, g, r], dtype=np.uint8)
+        mask = np.all(img == target_rgb, axis=2)
+        coords = np.argwhere(mask)
+        if coords.size > 0:
+            x_min, y_min = coords[:, 1].min(), coords[:, 0].min()
+            x_max, y_max = coords[:, 1].max(), coords[:, 0].max()
+            bbox_list.append([int(x_min-1), int(y_min-1), int(x_max+1), int(y_max+1)])
+        else:
             bbox_list.append([])
-            continue
-        
-    img = img.convert("L")
+
+    img = Image.open(image_path).convert("RGB").convert("L")
     img_bw = img.point(lambda x: 255 if x == 255 else 0, '1')
-    img_bw.convert("RGB").save(image_path) 
+    img_bw.convert("RGB").save(image_path)
     return bbox_list
 
+def contains_chinese(text):
+    # 匹配中文字符的正则表达式范围
+    return re.search(r'[\u4e00-\u9fff]', text) is not None
 
 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
+    if "tabular" in latex:
+        template = tabular_template
+    else:
+        if contains_chinese(latex):
+            template = formular_template_zh
+            latex = latex.replace("，", ", ").replace("：", ": ").replace("；", "; ")
+        else:
+            template = 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')
@@ -140,6 +172,7 @@ def latex2bbox_color(input_arg):
         return
     
     try:
+        latex = latex.replace("\n", " ")
         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}."
@@ -164,7 +197,7 @@ def latex2bbox_color(input_arg):
             paper_size = 4
         else:
             paper_size = 5
-        final_latex = formular_template.replace("<PaperSize>", str(paper_size)) % rgb_latex
+        final_latex = template.replace("<PaperSize>", str(paper_size)) % rgb_latex
         
     except Exception as e:
         log = f"ERROR, Preprocess latex failed: {basename}; {e}."
@@ -198,18 +231,21 @@ def latex2bbox_color(input_arg):
         vis = Image.open(output_base_path)
         draw = ImageDraw.Draw(vis)
 
-        with open(output_bbox_path, 'w') as f:
+        with open(output_bbox_path, 'w', encoding='utf-8') as f:
             for token, box in zip(token_list, bbox_list):
                 item = {
                     "bbox": box,
                     "token": token
                 }
-                f.write(json.dumps(item)+'\n')
+                f.write(json.dumps(item, ensure_ascii=False)+'\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))
+                try:
+                    draw.text((x_min, y_min), token, (250,0,0))
+                except:
+                    pass
             
         vis.save(output_vis_path)

modules/latex_processor.py

@@ -8,14 +8,14 @@ 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']
+SKIP_Tokens = ['\\', '\\\\', '\\index', '\\a', '&', '$', '\\multirow', '\\def', '\\edef', '\\raggedright', '\\url', '\\cr', '\\ensuremath', '\\left', '\\right', 
+               '\\mathchoice', '\\scriptstyle', '\\displaystyle', '\\qquad', '\\quad', '\\,', '\\!', '~', '\\boldmath', '\\gdef', '\\today', '\\the']
+PHANTOM_Tokens = ['\\fontfamily', '\\vphantom', '\\phantom', '\\rowcolor', '\\ref', '\\thesubequation', '\\global', '\\theboldgroup']
 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']
+                   '\\breve', '\\Bar', '\\Vec', '\\mathring', '\\ddot', '\\Ddot', '\\dddot', '\\ddddot']
 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']
@@ -150,29 +150,74 @@ def normalize_latex(l, rm_trail=False):
     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 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 \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+" ")
         
     # 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 } ")
-        
+    l = re.sub(r'(?<=\s)--(?=\s)', r'- -', l)
+    l = re.sub(r'(?<=\s)---(?=\s)', r'- - -', l)
+    l = re.sub(r'(?<=\s)…(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\ldots(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\hdots(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\cdots(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\dddot(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\dots(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\dotsc(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\dotsi(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\dotsm(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\dotso(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\dotsb(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\mathellipsis(?=\s)', r'. . .', l)
+    l = re.sub(r'(?<=\s)\\ex(?=\s)', r'\\mathrm { e x }', l)
+    l = re.sub(r'(?<=\s)\\ln(?=\s)', r'\\mathrm { l n }', l)
+    l = re.sub(r'(?<=\s)\\lg(?=\s)', r'\\mathrm { l g }', l)
+    l = re.sub(r'(?<=\s)\\cot(?=\s)', r'\\mathrm { c o t }', l)
+    l = re.sub(r'(?<=\s)\\mod(?=\s)', r'\\mathrm { m o d }', l)
+    l = re.sub(r'(?<=\s)\\bmod(?=\s)', r'\\mathrm { m o d }', l)
+    l = re.sub(r'(?<=\s)\\pmod(?=\s)', r'\\mathrm { m o d }', l)  # \pmod 其实和mod不一样，但是不太好处理，暂时替换为\mod
+    l = re.sub(r'(?<=\s)\\min(?=\s)', r'\\mathrm { m i n }', l) 
+    l = re.sub(r'(?<=\s)\\max(?=\s)', r'\\mathrm { m a x }', l) 
+    l = re.sub(r'(?<=\s)\\ker(?=\s)', r'\\mathrm { k e r }', l) 
+    l = re.sub(r'(?<=\s)\\hom(?=\s)', r'\\mathrm { h o m }', l)
+    l = re.sub(r'(?<=\s)\\sec(?=\s)', r'\\mathrm { s e c }', l)
+    l = re.sub(r'(?<=\s)\\scs(?=\s)', r'\\mathrm { s c s }', l)
+    l = re.sub(r'(?<=\s)\\csc(?=\s)', r'\\mathrm { c s c }', l)
+    l = re.sub(r'(?<=\s)\\deg(?=\s)', r'\\mathrm { d e g }', l)
+    l = re.sub(r'(?<=\s)\\arg(?=\s)', r'\\mathrm { a r g }', l)
+    l = re.sub(r'(?<=\s)\\log(?=\s)', r'\\mathrm { l o g }', l)
+    l = re.sub(r'(?<=\s)\\dim(?=\s)', r'\\mathrm { d i m }', l)
+    l = re.sub(r'(?<=\s)\\exp(?=\s)', r'\\mathrm { e x p }', l)
+    l = re.sub(r'(?<=\s)\\sin(?=\s)', r'\\mathrm { s i n }', l)
+    l = re.sub(r'(?<=\s)\\cos(?=\s)', r'\\mathrm { c o s }', l)
+    l = re.sub(r'(?<=\s)\\tan(?=\s)', r'\\mathrm { t a n }', l)
+    l = re.sub(r'(?<=\s)\\tanh(?=\s)', r'\\mathrm { t a n h }', l)
+    l = re.sub(r'(?<=\s)\\cosh(?=\s)', r'\\mathrm { c o s h }', l)
+    l = re.sub(r'(?<=\s)\\sinh(?=\s)', r'\\mathrm { s i n h }', l)
+    l = re.sub(r'(?<=\s)\\coth(?=\s)', r'\\mathrm { c o t h }', l)
+    l = re.sub(r'(?<=\s)\\arcsin(?=\s)', r'\\mathrm { a r c s i n }', l)
+    l = re.sub(r'(?<=\s)\\arccos(?=\s)', r'\\mathrm { a r c c o s }', l)
+    l = re.sub(r'(?<=\s)\\arctan(?=\s)', r'\\mathrm { a r c t a n }', l)
+    
+    # ** token such as \string xxx should be one token
+    pattern = r'\\string [^ ]+ '
+    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+" ")
+    
     # ** token such as \big( should be one token
     pattern = r'\\[Bb]ig[g]?[glrm]? [(){}|\[\]] '
     old_token = re.findall(pattern, l, re.DOTALL)
@@ -235,12 +280,12 @@ def normalize_latex(l, rm_trail=False):
     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 [\\=\<\>][^ ]+ '
+    # ** \rule{1pt}{2pt} lines, shoudle be combined as one token and do not render
+    pattern = r'\\rule {[ .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+" ")
+        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]+}'

modules/visual_matcher.py

@@ -42,41 +42,67 @@ def norm_coords(x, left, right):
 
 def norm_same_token(token):
     special_map = {
+        "\\dot": ".",
+        "\\Dot": ".",
         "\\cdot": ".",
+        "\\cdotp": ".",
+        "\\ldotp": ".",
         "\\mid": "|",
-        "\\to": "\\rightarrow",
+        "\\rightarrow": "\\to",
         "\\top": "T",
         "\\Tilde": "\\tilde",
-        "\\cdots": "\\dots",
         "\\prime": "'",
         "\\ast": "*",
         "\\left<": "\\langle",
-        "\\right>": "\\rangle"
+        "\\right>": "\\rangle",
+        "\\lbrace": "\{",
+        "\\rbrace": "\}",
+        "\\lbrack": "[",
+        "\\rbrack": "]",
+        "\\blackslash": "/",
+        "\\slash": "/",
+        "\\leq": "\\le",
+        "\\geq": "\\ge",
+        "\\neq": "\\ne",
+        "\\Vert": "\\|",
+        "\\lVert": "\\|",
+        "\\rVert": "\\|",
+        "\\vert": "|",
+        "\\lvert": "|",
+        "\\rvert": "|",
+        "\\colon": ":",
+        "\\Ddot": "\\ddot",
+        "\\Bar": "\\bar",
+        "\\Vec": "\\vec",
+        "\\parallel": "\\|",
+        "\\dag": "\\dagger",
+        "\\ddag": "\\ddagger",
+        "\\textlangle": "<",
+        "\\textrangle": ">",
+        "\\textgreater": ">",
+        "\\textless": "<",
+        "\\textbackslash": "n",
+        "\\textunderscore": "_",
+        "\\=": "_",
+        "\\neg": "\\lnot",
+        "\\neq": "\\not=",
     }
-    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 "arrow" not in token and "<" not in token and ">" not in token and "harpoon" not in token:
+            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 in special_map.keys():
+        token = special_map[token]
     if token.startswith('\\wide'):
         return token.replace("wide", "")
     if token.startswith('\\var'):
-        return token.replace("\\var", "")
+        return token.replace("var", "")
+    if token.startswith('\\string'):
+        return token.replace("\\string", "")
     return token
 
 
@@ -91,18 +117,6 @@ class HungarianMatcher:
         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 = {}
@@ -143,7 +157,7 @@ class HungarianMatcher:
         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
+        token_cost[np.logical_and(token_cost==1, norm_token_cost==0)] = 0.005
         return token_cost.T
         
         
@@ -155,10 +169,12 @@ class HungarianMatcher:
             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):
+    def order2array(self, box_list, max_token_lens=None):
+        if not max_token_lens:
+            max_token_lens = len(box_list)
         order_array = []
         for idx, box in enumerate(box_list):
-            order_array.append([idx / len(box_list)])
+            order_array.append([idx / max_token_lens])
         return np.array(order_array)
     
     def calculate_l1_cost(self, gt_array, pred_array):
@@ -170,8 +186,10 @@ class HungarianMatcher:
         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)
+
+        max_token_lens = max(len(box_gt), len(box_pred))
+        gt_order_array = self.order2array(box_gt, max_token_lens)
+        pred_order_array = self.order2array(box_pred, max_token_lens)
 
         token_cost = self.calculate_token_cost(box_gt, box_pred)
         position_cost = self.calculate_l1_cost(gt_box_array, pred_box_array)