Document AI yomitoku for python で png画像を html化

oss版と商用版があるようですが、oss版で十分すごいです

https://github.com/kotaro-kinoshita/yomitoku

# -*- coding: utf-8 -*-
import yomitoku
import glob
import logging
import os
import torch
from tensorflow.python.client import device_lib;

logger = None
device = None
src_png_dir = "ns_png"
src_html_dir = "ns_html"


def main():
    # GPU接続確認 https://end0tknr.hateblo.jp/entry/20240814/1723617124
    chk_gpu()

    # logger初期化 - yomitokuが既にloggerを設定している為
    loggers = init_loggers()
    global logger
    logger = loggers[0]

    # --- GPU設定 ---
    global device
    device = "cuda" # cpu or cuda

    # --- PDFのmarkdown化 ---
    conv_png_to_html(src_png_dir, device)

def chk_gpu():
    print( "torch.cuda.is_available() :", torch.cuda.is_available() )
    print( "torch.version.cuda :", torch.version.cuda )
    print( "torch.cuda.device_count() :", torch.cuda.device_count() )
    print( "torch.cuda.get_device_name() :", torch.cuda.get_device_name() )
    print( "device_lib.list_local_devices() :", device_lib.list_local_devices() )
    

def conv_png_to_html(src_png_dir, device):
    doc_analyzer = yomitoku.DocumentAnalyzer(device=device )
    ret_base_dir = src_html_dir

    for i, src_png_path in enumerate(glob.glob(src_png_dir + "/**/*.png",
                                               recursive=True)):
        print(i, src_png_path)
        (src_dir, src_filename) = os.path.split(src_png_path)

        md_dir = src_dir.replace(src_png_dir, src_html_dir)
        if not os.path.exists(md_dir):
            os.makedirs(md_dir)

        basename = os.path.splitext(src_filename)[0]
        html_path = src_html_dir + "/" + basename + ".html"

        imgs = yomitoku.data.functions.load_image(src_png_path)
        for j, img in enumerate(imgs):
            results, ocr_vis, layout_vis = doc_analyzer(img)
            results.to_html(html_path, img=img )
    return ret_base_dir

def init_loggers():
    """ yomitokuが内部に持つloggerを設定変更 """
    ret_datas = []
    for logger_name in ['yomitoku', 'yomitoku.base']:
        tmp_logger = logging.getLogger(logger_name)
        tmp_logger.setLevel(logging.INFO) # 出力levelをINFOへ

        for handler in tmp_logger.handlers:
            if isinstance(handler, logging.StreamHandler):
                tmp_logger.removeHandler(handler)

        fh = logging.FileHandler(logger_name + ".log", encoding='utf-8')
        formatter = logging.Formatter(
            '%(asctime)s %(name)s %(levelname)s %(message)s')
        fh.setFormatter(formatter)

        tmp_logger.addHandler(fh)
        ret_datas.append(tmp_logger)

    return ret_datas

if __name__ == '__main__':
    main()