import sys
import os
import re
import json
import time
import shutil
import argparse
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from datetime import datetime
from multiprocessing import Pool
from multiprocessing.dummy import Pool as ThreadPool
from PIL import Image, ImageDraw
from skimage.measure import ransac
from modules.latex2bbox_color import latex2bbox_color
from modules.tokenize_latex.tokenize_latex import tokenize_latex
from modules.visual_matcher import HungarianMatcher, SimpleAffineTransform
def gen_color_list(num=10, gap=15):
num += 1
single_num = 255 // gap + 1
max_num = single_num ** 3
num = min(num, max_num)
color_list = []
for idx in range(num):
R = idx // single_num**2
GB = idx % single_num**2
G = GB // single_num
B = GB % single_num
color_list.append((R*gap, G*gap, B*gap))
return color_list[1:]
def update_inliers(ori_inliers, sub_inliers):
inliers = np.copy(ori_inliers)
sub_idx = -1
for idx in range(len(ori_inliers)):
if ori_inliers[idx] == False:
sub_idx += 1
if sub_inliers[sub_idx] == True:
inliers[idx] = True
return inliers
def reshape_inliers(ori_inliers, sub_inliers):
inliers = np.copy(ori_inliers)
sub_idx = -1
for idx in range(len(ori_inliers)):
if ori_inliers[idx] == False:
sub_idx += 1
if sub_inliers[sub_idx] == True:
inliers[idx] = True
else:
inliers[idx] = False
return inliers
def gen_token_order(box_list):
new_box_list = copy.deepcopy(box_list)
for idx, box in enumerate(new_box_list):
new_box_list[idx]['order'] = idx / len(new_box_list)
return new_box_list
def evaluation(data_root, user_id="test"):
data_root = os.path.join(data_root, user_id)
gt_box_dir = os.path.join(data_root, "gt")
pred_box_dir = os.path.join(data_root, "pred")
match_vis_dir = os.path.join(data_root, "vis_match")
os.makedirs(match_vis_dir, exist_ok=True)
max_iter = 3
min_samples = 3
residual_threshold = 25
max_trials = 50
metrics_per_img = {}
gt_basename_list = [item.split(".")[0] for item in os.listdir(os.path.join(gt_box_dir, 'bbox'))]
for basename in tqdm(gt_basename_list):
gt_valid, pred_valid = True, True
if not os.path.exists(os.path.join(gt_box_dir, 'bbox', basename+".jsonl")):
gt_valid = False
else:
with open(os.path.join(gt_box_dir, 'bbox', basename+".jsonl"), 'r') as f:
box_gt = []
for line in f:
info = json.loads(line)
if info['bbox']:
box_gt.append(info)
if not box_gt:
gt_valid = False
if not gt_valid:
continue
if not os.path.exists(os.path.join(pred_box_dir, 'bbox', basename+".jsonl")):
pred_valid = False
else:
with open(os.path.join(pred_box_dir, 'bbox', basename+".jsonl"), 'r') as f:
box_pred = []
for line in f:
info = json.loads(line)
if info['bbox']:
box_pred.append(info)
if not box_pred:
pred_valid = False
if not pred_valid:
metrics_per_img[basename] = {
"recall": 0,
"precision": 0,
"F1_score": 0,
}
continue
gt_img_path = os.path.join(gt_box_dir, 'vis', basename+"_base.png")
pred_img_path = os.path.join(pred_box_dir, 'vis', basename+"_base.png")
img_gt = Image.open(gt_img_path)
img_pred = Image.open(pred_img_path)
matcher = HungarianMatcher()
matched_idxes = matcher(box_gt, box_pred, img_gt.size, img_pred.size)
src = []
dst = []
for (idx1, idx2) in matched_idxes:
x1min, y1min, x1max, y1max = box_gt[idx1]['bbox']
x2min, y2min, x2max, y2max = box_pred[idx2]['bbox']
x1_c, y1_c = float((x1min+x1max)/2), float((y1min+y1max)/2)
x2_c, y2_c = float((x2min+x2max)/2), float((y2min+y2max)/2)
src.append([y1_c, x1_c])
dst.append([y2_c, x2_c])
src = np.array(src)
dst = np.array(dst)
if src.shape[0] <= min_samples:
inliers = np.array([True for _ in matched_idxes])
else:
inliers = np.array([False for _ in matched_idxes])
for i in range(max_iter):
if src[inliers==False].shape[0] <= min_samples:
break
model, inliers_1 = ransac((src[inliers==False], dst[inliers==False]), SimpleAffineTransform, min_samples=min_samples, residual_threshold=residual_threshold, max_trials=max_trials, random_state=42)
if inliers_1 is not None and inliers_1.any():
inliers = update_inliers(inliers, inliers_1)
else:
break
if len(inliers[inliers==True]) >= len(matched_idxes):
break
for idx, (a,b) in enumerate(matched_idxes):
if inliers[idx] == True and matcher.cost['token'][a, b] == 1:
inliers[idx] = False
final_match_num = len(inliers[inliers==True])
recall = round(final_match_num/(len(box_gt)), 3)
precision = round(final_match_num/(len(box_pred)), 3)
F1_score = round(2*final_match_num/(len(box_gt)+len(box_pred)), 3)
metrics_per_img[basename] = {
"recall": recall,
"precision": precision,
"F1_score": F1_score,
}
if True:
gap = 5
W1, H1 = img_gt.size
W2, H2 = img_pred.size
H = H1 + H2 + gap
W = max(W1, W2)
vis_img = Image.new('RGB', (W, H), (255, 255, 255))
vis_img.paste(img_gt, (0, 0))
vis_img.paste(Image.new('RGB', (W, gap), (120, 120, 120)), (0, H1))
vis_img.paste(img_pred, (0, H1+gap))
match_img = vis_img.copy()
match_draw = ImageDraw.Draw(match_img)
gt_matched_idx = {
a: flag
for (a,b), flag in
zip(matched_idxes, inliers)
}
pred_matched_idx = {
b: flag
for (a,b), flag in
zip(matched_idxes, inliers)
}
for idx, box in enumerate(box_gt):
if idx in gt_matched_idx and gt_matched_idx[idx]==True:
color = "green"
else:
color = "red"
x_min, y_min, x_max, y_max = box['bbox']
match_draw.rectangle([x_min-1, y_min-1, x_max+1, y_max+1], fill=None, outline=color, width=2)
for idx, box in enumerate(box_pred):
if idx in pred_matched_idx and pred_matched_idx[idx]==True:
color = "green"
else:
color = "red"
x_min, y_min, x_max, y_max = box['bbox']
match_draw.rectangle([x_min-1, y_min-1+H1+gap, x_max+1, y_max+1+H1+gap], fill=None, outline=color, width=2)
vis_img.save(os.path.join(match_vis_dir, basename+"_base.png"))
match_img.save(os.path.join(match_vis_dir, basename+".png"))
score_list = [val['F1_score'] for _, val in metrics_per_img.items()]
exp_list = [1 if score==1 else 0 for score in score_list]
metrics_res = {
"mean_score": round(np.mean(score_list), 3),
"exp_rate": round(np.mean(exp_list), 3),
"details": metrics_per_img
}
metric_res_path = os.path.join(data_root, "metrics_res.json")
with open(metric_res_path, "w") as f:
f.write(json.dumps(metrics_res, indent=2))
return metrics_res, metric_res_path, match_vis_dir
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--input', '-i', type=str, default="assets/example/input_example.json")
parser.add_argument('--output', '-o', type=str, default="output")
parser.add_argument('--pools', '-p', type=int, default=240)
args = parser.parse_args()
print(args)
json_input, data_root, pool_num = args.input, args.output, args.pools
temp_dir = os.path.join(data_root, "temp_dir")
exp_name = os.path.basename(json_input).split('.')[0]
with open(json_input, "r") as f:
input_data = json.load(f)
img_ids = []
groundtruths = []
predictions = []
for idx, item in enumerate(input_data):
if "img_id" in item:
img_ids.append(item["img_id"])
else:
img_ids.append(f"sample_{idx}")
groundtruths.append(item['gt'])
predictions.append(item['pred'])
a = time.time()
user_id = exp_name
total_color_list = gen_color_list(num=5800)
data_root = os.path.join(data_root, user_id)
output_dir_info = {}
input_args = []
for subset, latex_list in zip(['gt', 'pred'], [groundtruths, predictions]):
sub_temp_dir = os.path.join(temp_dir, f"{exp_name}_{subset}")
os.makedirs(sub_temp_dir, exist_ok=True)
output_path = os.path.join(data_root, subset)
output_dir_info[output_path] = []
os.makedirs(os.path.join(output_path, 'bbox'), exist_ok=True)
os.makedirs(os.path.join(output_path, 'vis'), exist_ok=True)
for idx, latex in tqdm(enumerate(latex_list), desc=f"collect {subset} latex ..."):
basename = img_ids[idx]
input_arg = latex, basename, output_path, sub_temp_dir, total_color_list
input_args.append(input_arg)
if pool_num > 1:
print(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), "using processpool, pool num:", pool_num, ", job num:", len(input_args))
myP = Pool(args.pools)
for input_arg in input_args:
myP.apply_async(latex2bbox_color, args=(input_arg,))
myP.close()
myP.join()
else:
for input_arg in input_args:
latex2bbox_color(input_arg)
b = time.time()
print(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), "extract bbox done, time cost:", round(b-a, 3), "s")
for subset in ['gt', 'pred']:
shutil.rmtree(os.path.join(temp_dir, f"{exp_name}_{subset}"))
c = time.time()
metrics_res, metric_res_path, match_vis_dir = evaluation(args.output, exp_name)
d = time.time()
print(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), "calculate metrics done, time cost:", round(d-c, 3), "s")
print(f"=> process done, mean f1 score: {metrics_res['mean_score']}.")
print(f"=> more details of metrics are saved in `{metric_res_path}`")
print(f"=> visulization images are saved under `{match_vis_dir}`")