MinerU/magic_pdf/libs/calc_span_stats.py

import os
import csv
import json
import pandas as pd
from pandas import DataFrame as df
from matplotlib import pyplot as plt
from termcolor import cprint

"""
Execute this script in the following way:

1. Make sure there are pdf_dic.json files under the directory code-clean/tmp/unittest/md/, such as the following:

    code-clean/tmp/unittest/md/scihub/scihub_00500000/libgen.scimag00527000-00527999.zip_10.1002/app.25178/pdf_dic.json
    
2. Under the directory code-clean, execute the following command:

    $ python -m libs.calc_span_stats
    
"""


def print_green_on_red(text):
    cprint(text, "green", "on_red", attrs=["bold"], end="\n\n")


def print_green(text):
    print()
    cprint(text, "green", attrs=["bold"], end="\n\n")


def print_red(text):
    print()
    cprint(text, "red", attrs=["bold"], end="\n\n")


def safe_get(dict_obj, key, default):
    val = dict_obj.get(key)
    if val is None:
        return default
    else:
        return val


class SpanStatsCalc:
    """Calculate statistics of span."""

    def draw_charts(self, span_stats: pd.DataFrame, fig_num: int, save_path: str):
        """Draw multiple figures in one figure."""
        # make a canvas
        fig = plt.figure(fig_num, figsize=(20, 20))

        pass

    def calc_stats_per_dict(self, pdf_dict) -> pd.DataFrame:
        """Calculate statistics per pdf_dict."""
        span_stats = pd.DataFrame()

        span_stats = []
        span_id = 0
        for page_id, blocks in pdf_dict.items():
            if page_id.startswith("page_"):
                if "para_blocks" in blocks.keys():
                    for para_block in blocks["para_blocks"]:
                        for line in para_block["lines"]:
                            for span in line["spans"]:
                                span_text = safe_get(span, "text", "")
                                span_font_name = safe_get(span, "font", "")
                                span_font_size = safe_get(span, "size", 0)
                                span_font_color = safe_get(span, "color", "")
                                span_font_flags = safe_get(span, "flags", 0)

                                span_font_flags_decoded = safe_get(span, "decomposed_flags", {})
                                span_is_super_script = safe_get(span_font_flags_decoded, "is_superscript", False)
                                span_is_italic = safe_get(span_font_flags_decoded, "is_italic", False)
                                span_is_serifed = safe_get(span_font_flags_decoded, "is_serifed", False)
                                span_is_sans_serifed = safe_get(span_font_flags_decoded, "is_sans_serifed", False)
                                span_is_monospaced = safe_get(span_font_flags_decoded, "is_monospaced", False)
                                span_is_proportional = safe_get(span_font_flags_decoded, "is_proportional", False)
                                span_is_bold = safe_get(span_font_flags_decoded, "is_bold", False)

                                span_stats.append(
                                    {
                                        "span_id": span_id,  # id of span
                                        "page_id": page_id,  # page number of pdf
                                        "span_text": span_text,  # text of span
                                        "span_font_name": span_font_name,  # font name of span
                                        "span_font_size": span_font_size,  # font size of span
                                        "span_font_color": span_font_color,  # font color of span
                                        "span_font_flags": span_font_flags,  # font flags of span
                                        "span_is_superscript": int(
                                            span_is_super_script
                                        ),  # indicate whether the span is super script or not
                                        "span_is_italic": int(span_is_italic),  # indicate whether the span is italic or not
                                        "span_is_serifed": int(span_is_serifed),  # indicate whether the span is serifed or not
                                        "span_is_sans_serifed": int(
                                            span_is_sans_serifed
                                        ),  # indicate whether the span is sans serifed or not
                                        "span_is_monospaced": int(
                                            span_is_monospaced
                                        ),  # indicate whether the span is monospaced or not
                                        "span_is_proportional": int(
                                            span_is_proportional
                                        ),  # indicate whether the span is proportional or not
                                        "span_is_bold": int(span_is_bold),  # indicate whether the span is bold or not
                                    }
                                )

                                span_id += 1

        span_stats = pd.DataFrame(span_stats)
        # print(span_stats)

        return span_stats


def __find_pdf_dic_files(
    jf_name="pdf_dic.json",
    base_code_name="code-clean",
    tgt_base_dir_name="tmp",
    unittest_dir_name="unittest",
    md_dir_name="md",
    book_names=[
        "scihub",
    ],  # other possible values: "zlib", "arxiv" and so on
):
    pdf_dict_files = []

    curr_dir = os.path.dirname(__file__)

    for i in range(len(curr_dir)):
        if curr_dir[i : i + len(base_code_name)] == base_code_name:
            base_code_dir_name = curr_dir[: i + len(base_code_name)]
            for book_name in book_names:
                search_dir_relative_name = os.path.join(tgt_base_dir_name, unittest_dir_name, md_dir_name, book_name)
                if os.path.exists(base_code_dir_name):
                    search_dir_name = os.path.join(base_code_dir_name, search_dir_relative_name)
                    for root, dirs, files in os.walk(search_dir_name):
                        for file in files:
                            if file == jf_name:
                                pdf_dict_files.append(os.path.join(root, file))
                break

    return pdf_dict_files


def combine_span_texts(group_df, span_stats):
    combined_span_texts = []
    for _, row in group_df.iterrows():
        curr_span_id = row.name
        curr_span_text = row["span_text"]

        pre_span_id = curr_span_id - 1
        pre_span_text = span_stats.at[pre_span_id, "span_text"] if pre_span_id in span_stats.index else ""

        next_span_id = curr_span_id + 1
        next_span_text = span_stats.at[next_span_id, "span_text"] if next_span_id in span_stats.index else ""

        # pointer_sign is a right arrow if the span is superscript, otherwise it is a down arrow
        pointer_sign = "→ → → "
        combined_text = "\n".join([pointer_sign + pre_span_text, pointer_sign + curr_span_text, pointer_sign + next_span_text])
        combined_span_texts.append(combined_text)

    return "\n\n".join(combined_span_texts)


# pd.set_option("display.max_colwidth", None)  # 设置为 None 来显示完整的文本
pd.set_option("display.max_rows", None)  # 设置为 None 来显示更多的行


def main():
    pdf_dict_files = __find_pdf_dic_files()
    # print(pdf_dict_files)

    span_stats_calc = SpanStatsCalc()

    for pdf_dict_file in pdf_dict_files:
        print("-" * 100)
        print_green_on_red(f"Processing {pdf_dict_file}")

        with open(pdf_dict_file, "r", encoding="utf-8") as f:
            pdf_dict = json.load(f)

            raw_df = span_stats_calc.calc_stats_per_dict(pdf_dict)
            save_path = pdf_dict_file.replace("pdf_dic.json", "span_stats_raw.csv")
            raw_df.to_csv(save_path, index=False)

            filtered_df = raw_df[raw_df["span_is_superscript"] == 1]
            if filtered_df.empty:
                print("No superscript span found!")
                continue

            filtered_grouped_df = filtered_df.groupby(["span_font_name", "span_font_size", "span_font_color"])

            combined_span_texts = filtered_grouped_df.apply(combine_span_texts, span_stats=raw_df)  # type: ignore

            final_df = filtered_grouped_df.size().reset_index(name="count")
            final_df["span_texts"] = combined_span_texts.reset_index(level=[0, 1, 2], drop=True)

            print(final_df)

            final_df["span_texts"] = final_df["span_texts"].apply(lambda x: x.replace("\n", "\r\n"))

            save_path = pdf_dict_file.replace("pdf_dic.json", "span_stats_final.csv")
            # 使用 UTF-8 编码并添加 BOM，确保所有字段被双引号包围
            final_df.to_csv(save_path, index=False, encoding="utf-8-sig", quoting=csv.QUOTE_ALL)

            # 创建一个 2x2 的图表布局
            fig, axs = plt.subplots(2, 2, figsize=(15, 10))

            # 按照 span_font_name 分类作图
            final_df.groupby("span_font_name")["count"].sum().plot(kind="bar", ax=axs[0, 0], title="By Font Name")

            # 按照 span_font_size 分类作图
            final_df.groupby("span_font_size")["count"].sum().plot(kind="bar", ax=axs[0, 1], title="By Font Size")

            # 按照 span_font_color 分类作图
            final_df.groupby("span_font_color")["count"].sum().plot(kind="bar", ax=axs[1, 0], title="By Font Color")

            # 按照 span_font_name、span_font_size 和 span_font_color 共同分类作图
            grouped = final_df.groupby(["span_font_name", "span_font_size", "span_font_color"])
            grouped["count"].sum().unstack().plot(kind="bar", ax=axs[1, 1], title="Combined Grouping")

            # 调整布局
            plt.tight_layout()

            # 显示图表
            # plt.show()

            # 保存图表到 PNG 文件
            save_path = pdf_dict_file.replace("pdf_dic.json", "span_stats_combined.png")
            plt.savefig(save_path)

            # 清除画布
            plt.clf()


if __name__ == "__main__":
    main()