208 lines
7.7 KiB
Python
208 lines
7.7 KiB
Python
|
import torch
|
|||
|
|
from torch.utils.data import DataLoader, Dataset
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
from mineru.utils.boxbase import calculate_iou
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MathDataset(Dataset):
|
|||
|
|
def __init__(self, image_paths, transform=None):
|
|||
|
|
self.image_paths = image_paths
|
|||
|
|
self.transform = transform
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return len(self.image_paths)
|
|||
|
|
|
|||
|
|
def __getitem__(self, idx):
|
|||
|
|
raw_image = self.image_paths[idx]
|
|||
|
|
if self.transform:
|
|||
|
|
image = self.transform(raw_image)
|
|||
|
|
return image
|
|||
|
|
|
|||
|
|
|
|||
|
|
class UnimernetModel(object):
|
|||
|
|
def __init__(self, weight_dir, _device_="cpu"):
|
|||
|
|
from .unimernet_hf import UnimernetModel
|
|||
|
|
if _device_.startswith("mps") or _device_.startswith("npu"):
|
|||
|
|
self.model = UnimernetModel.from_pretrained(weight_dir, attn_implementation="eager")
|
|||
|
|
else:
|
|||
|
|
self.model = UnimernetModel.from_pretrained(weight_dir)
|
|||
|
|
self.device = _device_
|
|||
|
|
self.model.to(_device_)
|
|||
|
|
if not _device_.startswith("cpu"):
|
|||
|
|
self.model = self.model.to(dtype=torch.float16)
|
|||
|
|
self.model.eval()
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _filter_boxes_by_iou(xyxy, conf, cla, iou_threshold=0.8):
|
|||
|
|
"""过滤IOU超过阈值的重叠框，保留置信度较高的框。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
xyxy: 框坐标张量，shape为(N, 4)
|
|||
|
|
conf: 置信度张量，shape为(N,)
|
|||
|
|
cla: 类别张量，shape为(N,)
|
|||
|
|
iou_threshold: IOU阈值，默认0.9
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
过滤后的xyxy, conf, cla张量
|
|||
|
|
"""
|
|||
|
|
if len(xyxy) == 0:
|
|||
|
|
return xyxy, conf, cla
|
|||
|
|
|
|||
|
|
# 转换为CPU进行处理
|
|||
|
|
xyxy_cpu = xyxy.cpu()
|
|||
|
|
conf_cpu = conf.cpu()
|
|||
|
|
|
|||
|
|
n = len(xyxy_cpu)
|
|||
|
|
keep = [True] * n
|
|||
|
|
|
|||
|
|
for i in range(n):
|
|||
|
|
if not keep[i]:
|
|||
|
|
continue
|
|||
|
|
bbox1 = xyxy_cpu[i].tolist()
|
|||
|
|
for j in range(i + 1, n):
|
|||
|
|
if not keep[j]:
|
|||
|
|
continue
|
|||
|
|
bbox2 = xyxy_cpu[j].tolist()
|
|||
|
|
iou = calculate_iou(bbox1, bbox2)
|
|||
|
|
if iou > iou_threshold:
|
|||
|
|
# 保留置信度较高的框
|
|||
|
|
if conf_cpu[i] >= conf_cpu[j]:
|
|||
|
|
keep[j] = False
|
|||
|
|
else:
|
|||
|
|
keep[i] = False
|
|||
|
|
break # i被删除，跳出内循环
|
|||
|
|
|
|||
|
|
keep_indices = [i for i in range(n) if keep[i]]
|
|||
|
|
if len(keep_indices) == n:
|
|||
|
|
return xyxy, conf, cla
|
|||
|
|
|
|||
|
|
keep_indices = torch.tensor(keep_indices, dtype=torch.long)
|
|||
|
|
return xyxy[keep_indices], conf[keep_indices], cla[keep_indices]
|
|||
|
|
|
|||
|
|
def predict(self, mfd_res, image):
|
|||
|
|
formula_list = []
|
|||
|
|
mf_image_list = []
|
|||
|
|
|
|||
|
|
# 对检测框进行IOU去重，保留置信度较高的框
|
|||
|
|
xyxy_filtered, conf_filtered, cla_filtered = self._filter_boxes_by_iou(
|
|||
|
|
mfd_res.boxes.xyxy, mfd_res.boxes.conf, mfd_res.boxes.cls
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
for xyxy, conf, cla in zip(
|
|||
|
|
xyxy_filtered.cpu(), conf_filtered.cpu(), cla_filtered.cpu()
|
|||
|
|
):
|
|||
|
|
xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
|
|||
|
|
new_item = {
|
|||
|
|
"category_id": 13 + int(cla.item()),
|
|||
|
|
"poly": [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
|
|||
|
|
"score": round(float(conf.item()), 2),
|
|||
|
|
"latex": "",
|
|||
|
|
}
|
|||
|
|
formula_list.append(new_item)
|
|||
|
|
bbox_img = image[ymin:ymax, xmin:xmax]
|
|||
|
|
mf_image_list.append(bbox_img)
|
|||
|
|
|
|||
|
|
dataset = MathDataset(mf_image_list, transform=self.model.transform)
|
|||
|
|
dataloader = DataLoader(dataset, batch_size=32, num_workers=0)
|
|||
|
|
mfr_res = []
|
|||
|
|
for mf_img in dataloader:
|
|||
|
|
mf_img = mf_img.to(dtype=self.model.dtype)
|
|||
|
|
mf_img = mf_img.to(self.device)
|
|||
|
|
with torch.no_grad():
|
|||
|
|
output = self.model.generate({"image": mf_img})
|
|||
|
|
mfr_res.extend(output["fixed_str"])
|
|||
|
|
for res, latex in zip(formula_list, mfr_res):
|
|||
|
|
res["latex"] = latex
|
|||
|
|
return formula_list
|
|||
|
|
|
|||
|
|
def batch_predict(
|
|||
|
|
self,
|
|||
|
|
images_mfd_res: list,
|
|||
|
|
images: list,
|
|||
|
|
batch_size: int = 64,
|
|||
|
|
interline_enable: bool = True,
|
|||
|
|
) -> list:
|
|||
|
|
images_formula_list = []
|
|||
|
|
mf_image_list = []
|
|||
|
|
backfill_list = []
|
|||
|
|
image_info = [] # Store (area, original_index, image) tuples
|
|||
|
|
|
|||
|
|
# Collect images with their original indices
|
|||
|
|
for image_index in range(len(images_mfd_res)):
|
|||
|
|
mfd_res = images_mfd_res[image_index]
|
|||
|
|
image = images[image_index]
|
|||
|
|
formula_list = []
|
|||
|
|
|
|||
|
|
# 对检测框进行IOU去重，保留置信度较高的框
|
|||
|
|
xyxy_filtered, conf_filtered, cla_filtered = self._filter_boxes_by_iou(
|
|||
|
|
mfd_res.boxes.xyxy, mfd_res.boxes.conf, mfd_res.boxes.cls
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
for idx, (xyxy, conf, cla) in enumerate(zip(
|
|||
|
|
xyxy_filtered, conf_filtered, cla_filtered
|
|||
|
|
)):
|
|||
|
|
if not interline_enable and cla.item() == 1:
|
|||
|
|
continue # Skip interline regions if not enabled
|
|||
|
|
xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
|
|||
|
|
new_item = {
|
|||
|
|
"category_id": 13 + int(cla.item()),
|
|||
|
|
"poly": [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
|
|||
|
|
"score": round(float(conf.item()), 2),
|
|||
|
|
"latex": "",
|
|||
|
|
}
|
|||
|
|
formula_list.append(new_item)
|
|||
|
|
bbox_img = image[ymin:ymax, xmin:xmax]
|
|||
|
|
area = (xmax - xmin) * (ymax - ymin)
|
|||
|
|
|
|||
|
|
curr_idx = len(mf_image_list)
|
|||
|
|
image_info.append((area, curr_idx, bbox_img))
|
|||
|
|
mf_image_list.append(bbox_img)
|
|||
|
|
|
|||
|
|
images_formula_list.append(formula_list)
|
|||
|
|
backfill_list += formula_list
|
|||
|
|
|
|||
|
|
# Stable sort by area
|
|||
|
|
image_info.sort(key=lambda x: x[0]) # sort by area
|
|||
|
|
sorted_indices = [x[1] for x in image_info]
|
|||
|
|
sorted_images = [x[2] for x in image_info]
|
|||
|
|
|
|||
|
|
# Create mapping for results
|
|||
|
|
index_mapping = {new_idx: old_idx for new_idx, old_idx in enumerate(sorted_indices)}
|
|||
|
|
|
|||
|
|
# Create dataset with sorted images
|
|||
|
|
dataset = MathDataset(sorted_images, transform=self.model.transform)
|
|||
|
|
|
|||
|
|
# 如果batch_size > len(sorted_images)，则设置为不超过len(sorted_images)的2的幂
|
|||
|
|
batch_size = min(batch_size, max(1, 2 ** (len(sorted_images).bit_length() - 1))) if sorted_images else 1
|
|||
|
|
|
|||
|
|
dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=0)
|
|||
|
|
|
|||
|
|
# Process batches and store results
|
|||
|
|
mfr_res = []
|
|||
|
|
# for mf_img in dataloader:
|
|||
|
|
|
|||
|
|
with tqdm(total=len(sorted_images), desc="MFR Predict") as pbar:
|
|||
|
|
for index, mf_img in enumerate(dataloader):
|
|||
|
|
mf_img = mf_img.to(dtype=self.model.dtype)
|
|||
|
|
mf_img = mf_img.to(self.device)
|
|||
|
|
with torch.no_grad():
|
|||
|
|
output = self.model.generate({"image": mf_img}, batch_size=batch_size)
|
|||
|
|
mfr_res.extend(output["fixed_str"])
|
|||
|
|
|
|||
|
|
# 更新进度条，每次增加batch_size，但要注意最后一个batch可能不足batch_size
|
|||
|
|
current_batch_size = min(batch_size, len(sorted_images) - index * batch_size)
|
|||
|
|
pbar.update(current_batch_size)
|
|||
|
|
|
|||
|
|
# Restore original order
|
|||
|
|
unsorted_results = [""] * len(mfr_res)
|
|||
|
|
for new_idx, latex in enumerate(mfr_res):
|
|||
|
|
original_idx = index_mapping[new_idx]
|
|||
|
|
unsorted_results[original_idx] = latex
|
|||
|
|
|
|||
|
|
# Fill results back
|
|||
|
|
for res, latex in zip(backfill_list, unsorted_results):
|
|||
|
|
res["latex"] = latex
|
|||
|
|
|
|||
|
|
return images_formula_list
|