Rustで学術論文からテキストを抽出する #10 - Util系関数およびテストたち

:::note
## Summary
- 細かいUtilの紹介です

GiHub -> https://github.com/akitenkrad/rsrpp
crates.io -> https://crates.io/crates/rsrpp

## ToDo
- [x] pdftotextで論文から単語単位のテキストと位置情報を取得する (`Word`,`Line`,`Block`,`Page`)
- [x] テキストの属性 (本文, タイトル, 脚注, etc.) を判定する
    - [x] テキストが含まれるエリアを抽出する
        - [x] 2段組みを扱えるようにする
    - [x] セクションのタイトルを識別する
- [x] 図表に含まれるテキストを除外する
    - [x] 表を除外する
        - [x] PDFを画像に変換
        - [x] 画像処理で表の位置を特定
## 今日のファイル
```text
rsrpp
├── Cargo.toml
├── rsrpp
│   ├── Cargo.toml
│   └── src
│       ├── lib.rs
│       └── parser
│           ├── mod.rs <-- today
│           ├── structs.rs <-- today
│           └── tests.rs <-- today
└── rsrpp-cli
    ├── Cargo.toml
    └── src
        └── main.rs
```
:::

## 前回までのあらすじ

[前回：Rustで学術論文からテキストを抽出する #9](https://qiita.com/akitenkrad/items/4acca22b04c277baf1e8)

前回までで，PDFをパースしてJSONに変換するプログラムが大体完成しました．
今回は，これまでのところであまり説明してこなかったUtil系の処理やテストに関する部分を紹介しようと思います．

```mermaid
graph LR
    A[要件定義]

    subgraph Popplerの観察
    B[Popperの観察]
    B1[pdftotext]
    B2[pdfinfo]
    B3[pdftohtml]
    B4[pdftocairo]
    end

    subgraph テキストの属性判定
    C[テキストの属性判定]
    C1[テキストエリアの抽出]
    end

    subgraph 図表の処理
    D[図表の処理]
    subgraph 表の検出

    D1[表の検出]
    end
    end
 
    subgraph 今日はここ
    E[JSON変換]
    end
    
    A-->B
    B1-->C
    B3-->C
    C1-->D
    D-->E
    B4-->D1
```

## ParserConfig

既に登場していますが，今回実装したパーサの設定やファイルパスなどのメタ情報を保持する構造体です．
`ParserConfig`には，`clean_files`という関数を実装しています．
今回はPopplerのツール群をフル活用しているので，PDFから変換したHTML，XML，画像ファイルなどを`/tmp`に一時保存するように実装しています．
`/tmp`とはいえ，中間ファイルをそのまま残しておくのは気分が良くないので，PDFのパースが完了したらファイルを消せるようにしておきました．
立つ鳥跡を濁さず．

また，[第5回](https://qiita.com/akitenkrad/items/0a4c8ddcabaa4e55c7a5)でセクションタイトルもこの構造体に格納されます．

```rust: rsrpp > rsrpp > src > parser > structs.rs
#[derive(Debug, Clone, PartialEq)]
pub struct ParserConfig {
    pub pdf_path: String,
    pub pdf_text_path: String,
    pub pdf_figures: HashMap<PageNumber, String>,
    pub pdf_xml_path: String,
    pub sections: Vec<(PageNumber, String)>,
    pub pdf_info: HashMap<String, String>,
}

impl ParserConfig {
    pub fn new() -> ParserConfig {
        let mut rng = rand::thread_rng();
        let random_value = rng.gen_range(10000..99999);
        let mut pdf_path = String::new();
        pdf_path.push_str("/tmp/pdf_");
        pdf_path.push_str(&random_value.to_string());
        pdf_path.push_str(".pdf");

        let pdf_figures = HashMap::new();
        let pdf_html_path = pdf_path.clone().replace(".pdf", ".text.html");
        let pdf_raw_html_path = pdf_path.clone().replace(".pdf", ".xml");
        let sections = Vec::new();
        ParserConfig {
            pdf_path: pdf_path,
            pdf_text_path: pdf_html_path,
            pdf_figures: pdf_figures,
            pdf_xml_path: pdf_raw_html_path,
            sections: sections,
            pdf_info: HashMap::new(),
        }
    }

    pub fn pdf_width(&self) -> i32 {
        return self.pdf_info.get("page_width").unwrap().parse::<i32>().unwrap();
    }

    pub fn pdf_height(&self) -> i32 {
        return self.pdf_info.get("page_height").unwrap().parse::<i32>().unwrap();
    }

    pub fn clean_files(&self) -> Result<()> {
        if Path::new(&self.pdf_path).exists() {
            std::fs::remove_file(&self.pdf_path)?;
        }
        if Path::new(&self.pdf_text_path).exists() {
            std::fs::remove_file(&self.pdf_text_path)?;
        }
        if Path::new(&self.pdf_xml_path).exists() {
            std::fs::remove_file(&self.pdf_xml_path)?;
        }
        for figure in self.pdf_figures.values() {
            if Path::new(figure).exists() {
                std::fs::remove_file(figure)?;
            }
        }
        return Ok(());
    }
}
```

## Popplerツール群のラッパー

Popperのコマンドを実行するにあたって，1コマンド1関数で`ParserConfig`から情報をもらって実行するような設計にしました．

### pdfinfo

`pdfinfo`コマンドはこれまで出てこなかったかもしれませんが，PDFの幅と高さを取得するために使っています．
ちなみに，Rustでシェルコマンドを実行する際には`std::process::Command`を使うことができます．

```rust: rsrpp > rsrpp > src > parser > mod.rs
fn get_pdf_info(config: &mut ParserConfig) -> Result<()> {
    let res =
        Command::new("pdfinfo").args(&[config.pdf_path.clone()]).stdout(Stdio::piped()).output();
    let text = String::from_utf8(res?.stdout)?;
    for line in text.split("\n") {
        let parts: Vec<&str> = line.split(":").collect();
        if parts.len() < 2 {
            continue;
        }
        let key = parts[0].trim().to_string().to_lowercase().replace(" ", "_");
        let value = parts[1].trim().to_string();

        if key == "page_size" {
            let regex = regex::Regex::new(r"(\d+) x (\d+)")?;
            let caps = regex.captures(&value).unwrap();
            config.pdf_info.insert("page_width".to_string(), caps[1].to_string());
            config.pdf_info.insert("page_height".to_string(), caps[2].to_string());
        }
        config.pdf_info.insert(key, value);
    }
    return Ok(());
}
```

### pdftocairo

こちらはPDFを画像に変換するコマンドです．
以前紹介した通り，変換後のHTMLと画像の縮尺を合わせるために `-r 72`を指定しています．
変換後の画像は`/tmp`に格納し，画像のパスのリストを`ParserConfig`に保持するようにしています．

```rust: rsrpp > rsrpp > src > parser > mod.rs
fn save_pdf_as_figures(config: &mut ParserConfig) -> Result<()> {
    let pdf_path = Path::new(config.pdf_path.as_str());
    let dst_path = pdf_path.parent().unwrap().join(pdf_path.file_stem().unwrap().to_str().unwrap());

    // save pdf as jpeg files
    let res = Command::new("pdftocairo")
        .args(&[
            "-jpeg".to_string(),
            "-r".to_string(),
            "72".to_string(),
            pdf_path.to_str().unwrap().to_string(),
            dst_path.to_str().unwrap().to_string(),
        ])
        .stdout(Stdio::piped())
        .output();
    if let Err(e) = res {
        return Err(Error::msg(format!("Error: {}", e)));
    }

    // get all jpeg files
    let glob_query = dst_path.file_name().unwrap().to_str().unwrap().to_string() + "*.jpg";
    let glob_query = dst_path.parent().unwrap().join(glob_query);
    for entry in glob(glob_query.to_str().unwrap())? {
        match entry {
            Ok(path) => {
                let page_number: PageNumber = path
                    .file_stem()
                    .unwrap()
                    .to_str()
                    .unwrap()
                    .split("-")
                    .last()
                    .unwrap()
                    .parse::<i8>()?;
                config.pdf_figures.insert(page_number, path.to_str().unwrap().to_string());
            }
            Err(e) => return Err(Error::msg(format!("Error: {}", e))),
        }
    }

    return Ok(());
}
```

### pdftohtml

PDFをHTMLに変換します．出力をXMLにしているので関数名は`_as_xml`ですが．
この関数の重要な役割はフォントサイズに基づいたセクションタイトルの抽出です．

```rust: rsrpp > rsrpp > src > parser > mod.rs
fn save_pdf_as_xml(config: &mut ParserConfig) -> Result<()> {
    let xml_path = Path::new(&config.pdf_xml_path);

    Command::new("pdftohtml")
        .args(&[
            "-c".to_string(),
            "-s".to_string(),
            "-dataurls".to_string(),
            "-xml".to_string(),
            "-zoom".to_string(),
            "1.0".to_string(),
            config.pdf_path.as_str().to_string(),
            xml_path.to_str().unwrap().to_string(),
        ])
        .stdout(Stdio::piped())
        .output()?;

    // get title font size
    let mut font_number = 0;
    let xml_text = std::fs::read_to_string(xml_path)?;
    let mut reader = quick_xml::Reader::from_str(&xml_text);
    reader.config_mut().trim_text(true);
    loop {
        match reader.read_event() {
            Ok(Event::Start(e)) => {
                if e.name().as_ref() == b"text" {
                    for attr in e.attributes() {
                        let attr = attr?;
                        if attr.key.as_ref() == b"font" {
                            font_number = String::from_utf8_lossy(attr.value.as_ref())
                                .parse::<i32>()
                                .unwrap();
                        }
                    }
                }
            }
            Ok(Event::Text(e)) => {
                if String::from_utf8_lossy(e.as_ref()).to_lowercase() == "abstract" {
                    break;
                }
            }
            Err(_e) => {
                break;
            }
            _ => {}
        }
    }

    // get sections
    let mut page_number = 0;
    let mut is_title = false;
    let regex_is_number = regex::Regex::new(r"^\d+$").unwrap();
    let regex_trim_number = regex::Regex::new(r"\d\.").unwrap();
    let mut reader = quick_xml::Reader::from_str(&xml_text);
    reader.config_mut().trim_text(true);
    loop {
        match reader.read_event() {
            Ok(Event::Start(e)) => {
                if e.name().as_ref() == b"page" {
                    for attr in e.attributes() {
                        let attr = attr?;
                        if attr.key.as_ref() == b"number" {
                            page_number =
                                String::from_utf8_lossy(attr.value.as_ref()).parse::<i8>().unwrap();
                        }
                    }
                } else if e.name().as_ref() == b"text" {
                    let _font_number = String::from_utf8_lossy(
                        e.attributes()
                            .find(|attr| attr.clone().unwrap().key.as_ref() == b"font")
                            .unwrap()
                            .unwrap()
                            .value
                            .as_ref(),
                    )
                    .parse::<i32>()
                    .unwrap();

                    if font_number == _font_number {
                        is_title = true;
                    } else {
                        is_title = false;
                    }
                    continue;
                }
            }
            Ok(Event::Text(e)) => {
                let text = String::from_utf8_lossy(e.as_ref());
                if regex_is_number.is_match(&text) {
                    continue;
                }
                let text = regex_trim_number.replace(&text, "").to_string().trim().to_string();
                if is_title {
                    config.sections.push((page_number, text.to_string()));
                    if text.to_lowercase() == "references" {
                        break;
                    }
                }
            }
            Err(_e) => {
                break;
            }
            _ => {}
        }
    }

    return Ok(());
}
```

### pdftotext

最後が，今回実装したPDFパーサの中核になるコマンドです．
PDFを構造化されたHTMLファイルとして保存します．

```rust: rsrpp > rsrpp > src > parser > mod.rs
fn save_pdf_as_text(config: &mut ParserConfig) -> Result<()> {
    let html_path = Path::new(config.pdf_text_path.as_str());

    // parse pdf into html
    let _ = Command::new("pdftotext")
        .args(&[
            "-nopgbrk".to_string(),
            "-htmlmeta".to_string(),
            "-bbox-layout".to_string(),
            "-r".to_string(),
            "72".to_string(),
            config.pdf_path.as_str().to_string(),
            html_path.to_str().unwrap().to_string(),
        ])
        .stdout(Stdio::piped())
        .output()?;

    return Ok(());
}
```

### save_pdf()

上述のコマンドをまとめて発行してファイルを保存する関数です．
引数は元のPDFファイルなのですが，ローカルのファイルかURLを指定できるようにしています．
URLの場合は一度`/tmp`にオリジナルのPDFとしてファイルを保存してから変換しています．

```rust: rsrpp > rsrpp > src > parser > mod.rs
async fn save_pdf(path_or_url: &str, config: &mut ParserConfig) -> Result<()> {
    let save_path = config.pdf_path.as_str();
    if path_or_url.starts_with("http") {
        let res = request::get(path_or_url).await;
        let bytes = res?.bytes().await;
        let out = File::create(save_path);
        std::io::copy(&mut bytes?.as_ref(), &mut out?)?;
    } else {
        let path = Path::new(path_or_url);
        let _ = std::fs::copy(path.as_os_str(), save_path);
    }

    // get pdf info
    get_pdf_info(config).unwrap();

    // save pdf as jpeg files
    save_pdf_as_figures(config).unwrap();

    // save pdf as html
    save_pdf_as_xml(config).unwrap();

    // save pdf as text
    save_pdf_as_text(config).unwrap();

    return Ok(());
}
```

## テスト

これまでは全く触れてこなかったですが，主要な関数に関しては簡単にテストを書いています．
Rustはコードとテストの距離が非常に近いので，実装しながらテストを書くのがとてもやりやすいです．

### PDF保存のテスト

単体テストをどの粒度で実装するかは非常に難しい問題だと思うのですが，個人的には[単体テストの考え方/使い方](https://www.amazon.co.jp/%E5%8D%98%E4%BD%93%E3%83%86%E3%82%B9%E3%83%88%E3%81%AE%E8%80%83%E3%81%88%E6%96%B9-%E4%BD%BF%E3%81%84%E6%96%B9-Vladimir-Khorikov-ebook/dp/B0BLTG8Z9K)を参考にして，機能としてひとまとまりになっているモジュールをテスト対象にしています．

全ては載せられないので，いくつかサンプルをば．

```rust: rsrpp > rsrpp > src > parser > tests.rs
#[tokio::test]
async fn test_save_pdf_1() {
    let mut config = ParserConfig::new();
    let url = "https://arxiv.org/pdf/1706.03762";
    // let url = "https://arxiv.org/pdf/2308.10379";
    save_pdf(url, &mut config).await.unwrap();

    assert!(Path::new(&config.pdf_path).exists());

    for (_, path) in config.pdf_figures.iter() {
        println!("path: {}", path);
        assert!(Path::new(path).exists());
    }

    assert_eq!(config.sections[0], (1, "Abstract".to_string()));
    assert_eq!(config.sections[1], (2, "Introduction".to_string()));
    assert_eq!(config.sections[2], (2, "Background".to_string()));
    assert_eq!(config.sections[3], (2, "Model Architecture".to_string()));
    assert_eq!(config.sections[4], (6, "Why Self-Attention".to_string()));
    assert_eq!(config.sections[5], (7, "Training".to_string()));
    assert_eq!(config.sections[6], (8, "Results".to_string()));
    assert_eq!(config.sections[7], (10, "Conclusion".to_string()));
    assert_eq!(config.sections[8], (10, "References".to_string()));

    for (page, section) in config.sections.iter() {
        println!("page: {}, section: {}", page, section);
    }

    let _ = config.clean_files();
}
```

こちらはURLまたはローカルファイルをインプットにしたときのテスト．

```rust: rsrpp > rsrpp > src > parser > tests.rs
#[tokio::test]
async fn test_pdf2html_url() {
    let mut config = ParserConfig::new();
    let url = "https://arxiv.org/pdf/1706.03762";
    let res = pdf2html(url, &mut config).await;
    let html = res.unwrap();
    assert!(html.html().contains("arXiv:1706.03762"));
    let _ = config.clean_files();
}

#[tokio::test]
async fn test_pdf2html_file() {
    let mut config = ParserConfig::new();
    let url = "https://arxiv.org/pdf/1706.03762";
    let response = request::get(url).await.unwrap();
    let bytes = response.bytes().await.unwrap();
    let path = "/tmp/test.pdf";
    let mut file = File::create(path).unwrap();
    std::io::copy(&mut bytes.as_ref(), &mut file).unwrap();

    let res = pdf2html("/tmp/test.pdf", &mut config).await;
    let html = res.unwrap();
    assert!(html.html().contains("arXiv:1706.03762"));

    let _ = config.clean_files();
}
```

## 次回

次回はPDFパーサを実行する関数本体を実装します．
これで，PDFをパースしてJSONに落とす機能の実装はひとまず完成です．

[次回：Rustで学術論文からテキストを抽出する #11](https://qiita.com/akitenkrad/items/1718aa3096137ef76fec)