8

CnSTD V1.2.1 发布,增加数学公式检测和版面分析预训练模型

 1 year ago
source link: https://breezedeus.github.io/2023/02/01/breezedeus-cnstd-v121.html
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

CnSTD V1.2.1 发布

CnSTD 之前的版本,定位都是做文字检测,新版 V1.2.1 加入了数学公式检测Mathematical Formula Detection,简称MFD)和 版面分析Layout Analysis)模型,并提供预训练好的模型可直接使用。这两个模型底层架构用的都是YOLOv7的tiny模型,速度还是挺快的,但效果上并不算出众。

CnSTD V1.2.1其实是为了我另一个开源工具 Pix2Text (P2T) 的新版定制的。旧的P2T只能识别要么是文字,要么是数学公式的图片。依托于 CnSTD V1.2.1,新版的P2T能够识别既包含文字又包含数学公式的混合图片,比如这个图片:

混合图片

P2T 可以识别为:

这样就有

$$
{\cal L}_{t}=\mathbb{E}_{x_{0},\epsilon}\biggl[{\frac{(1-\alpha_{t})^{2}}{2\alpha_{t}(1-\tilde{\alpha}_{t})|\mathbf{E}_\theta||_{2}^{2}}}||\epsilon_{t}-\epsilon_{\theta}({\sqrt{\tilde{\alpha}_{t}}}\mathbf{x}_{0}+{\sqrt{1-\tilde{\alpha}_{t}}}\epsilon_{t},t)|^{2}\biggr]
$$

Noise Prediction $\mathbb{C}_{\theta}$ 的输入和输出长度都与 $\mathbf{x}_{t}$ 大小一致(即图片的size),比较大,所以可以对∈。使用U-Net架构,降低计算量。

我们会争取尽快把P2T新版开源出来。P2T是个Python包,对于不熟悉Python的朋友还是太不友好,我们近期也会提供P2T的网页版,直接把图片丢进网页就能输出P2T的解析结果。

网页版会提供一些免费名额供有需要的朋友使用,优先在校的学生( MathPix 每月要5美元,对在校生来说还是蛮贵的)。

有兴趣的朋友也可以留言说明,以便我们更好地预估网页版使用人数,提前准备好服务器。

也欢迎扫码加入微信交流群,了解P2T、CnSTD、CnOCR等工具的最新进展:

微信群二维码

接下来介绍 CnSTD V1.2.1 的新功能,更详细的说明请转到项目主页 CnSTD

数学公式检测(MFD)

MFD 模型检测图片中包含的数学公式,其中行内的公式检测为 embedding 类别,独立行的公式检测为 isolated。模型训练使用了英文 IBEM 和中文 CnMFD_Dataset 两个数据集。

中文MFD效果
中文MFD效果
英文MFD效果

版面分析(Layout Analysis)

版面分析模型识别图片中的不同排版元素。模型训练使用的是 CDLA 数据集。可识别以下10中版面元素:

正文 标题 图片 图片标题 表格 表格标题 页眉 页脚 注释 公式
Text Title Figure Figure caption Table Table caption Header Footer Reference Equation
版面分析效果

数学公式检测(MFD)与 版面分析(Layout Analysis)

数学公式检测(MFD)与 版面分析(Layout Analysis)都是检测图片中感兴趣的元素,它们使用的都是基于YOLOv7的检测架构,在CnSTD都来源于相同的类 LayoutAnalyzer,差别只是训练模型使用的数据不同。

这两个模型的训练代码在 yolov7 中(Forked from WongKinYiu/yolov7,感谢原作者。)

LayoutAnalyzer 的初始化函数如下:

class LayoutAnalyzer(object):
    def __init__(
        self,
        model_name: str = 'mfd',  # 'layout' or 'mfd'
        *,
        model_type: str = 'yolov7_tiny',
        model_backend: str = 'pytorch',
        model_fp: Optional[str] = None,
        root: Union[str, Path] = data_dir(),
        device: str = 'cpu',
        **kwargs,
    ):

其中的参数含义如下:

  • model_name: 字符串类型,表示模型类型。可选值:’mfd’ 表示数学公式检测;’layout’ 表示版面分析。默认值:’mfd’

  • model_type: 字符串类型,表示模型类型。当前仅支持 ‘yolov7_tiny’;默认值:’yolov7_tiny’

  • model_backend: 字符串类型,表示backend。当前仅支持: ‘pytorch’;默认值:’pytorch’

  • model_fp: 字符串类型,表示模型文件的路径。默认值:None,表示使用默认的文件路径

  • root: 字符串或Path类型,表示模型文件所在的根目录。
    • Linux/Mac下默认值为 ~/.cnstd,表示模型文件所处文件夹类似 ~/.cnstd/1.2/analysis
    • Windows下默认值为 C:/Users/<username>/AppData/Roaming/cnstd
  • device: 字符串类型,表示运行模型的设备,可选值:’cpu’ 或 ‘gpu’;默认值:’cpu’

  • **kwargs: 额外的参数。

函数输出结果为一个list,其中每个元素表示识别出的版面中的一个元素,包含以下信息:

  • type: 版面元素对应的类型;可选值来自:self.categories ;
  • box: 版面元素对应的矩形框;np.ndarray, shape: (4, 2),对应 box 4个点的坐标值 (x, y) ;
  • score: 得分,越高表示越可信 。

类函数LayoutAnalyzer.analyze()

对指定图片(列表)进行版面分析。

def analyze(
    self,
    img_list: Union[
        str,
        Path,
        Image.Image,
        np.ndarray,
        List[Union[str, Path, Image.Image, np.ndarray]],
    ],
    resized_shape: Union[int, Tuple[int, int]] = 700,
    box_margin: int = 2,
    conf_threshold: float = 0.25,
    iou_threshold: float = 0.45,
) -> Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]:

函数说明

函数输入参数包括:

  • img_list (str or list): 待识别图片或图片列表;如果是 np.ndarray,则应该是shape为 [H, W, 3] 的 RGB 格式数组
  • resized_shape (int or tuple): (H, W); 把图片resize到此大小再做分析;默认值为 700
  • box_margin (int): 对识别出的内容框往外扩展的像素大小;默认值为 2
  • conf_threshold (float): 分数阈值;默认值为 0.25
  • iou_threshold (float): IOU阈值;默认值为 0.45
  • **kwargs: 额外的参数。
from cnstd import LayoutAnalyzer
img_fp = 'examples/mfd/zh5.jpg'
analyzer = LayoutAnalyzer('mfd')
out = analyzer.analyze(img_fp, resized_shape=700)
print(out)

cnstd 包含了几个命令行工具,安装 cnstd 后即可使用。

MFD or Layout Analysis 预测单个文件

使用命令 cnstd analyze 获得单个文件的 MFD 或者 Layout Analysis 结果,以下是使用说明:

(venv) ➜  cnstd git:(master) ✗ cnstd analyze -h
Usage: cnstd analyze [OPTIONS]

  对给定图片进行 MFD 或者 版面分析。

Options:
  -m, --model-name [mfd|layout]   模型类型。`mfd` 表示数学公式检测,`layout`
                                  表示版面分析;默认为:`mfd`
  -t, --model-type TEXT           模型类型。当前仅支持 `yolov7_tiny`
  -b, --model-backend [pytorch|onnx]
                                  模型后端架构。当前仅支持 `pytorch`
  -p, --model-fp TEXT             使用训练好的模型。默认为 `None`,表示使用系统自带的预训练模型
  --device TEXT                   cuda device, i.e. 0 or 0,1,2,3 or cpu
  -i, --img-fp TEXT               待分析的图片路径
  -o, --output-fp TEXT            分析结果输出的图片路径。默认为 `None`,会存储在当前文件夹,文件名称为输入文件名称
                                  前面增加`out-`;如输入文件名为 `img.jpg`, 输出文件名即为 `out-
                                  img.jpg`
  --resized-shape INTEGER         分析时把图片resize到此大小再进行。默认为 `700`
  --conf-thresh FLOAT             Confidence Threshold。默认值为 `0.25`
  --iou-thresh FLOAT              IOU threshold for NMS。默认值为 `0.45`
  -h, --help                      Show this message and exit.

例如可以使用以下命令对图片 examples/mfd/zh.jpg 进行 MFD,并把检测结果存放在文件 out-zh.jpg 中:

(venv) ➜  cnstd analyze -m mfd --conf-thresh 0.25 --resized-shape 800 -i examples/mfd/zh.jpg -o out-zh.jpg

具体使用也可参考文件 Makefile


官方代码库:https://github.com/breezedeus/cnstd




About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK