|
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}`") |