岚法律办公自动化程序「香港」

# -*- coding: utf-8 -*-
# @Time    : 2025/8/27 18:13
# @Author  : AI悦创
# @FileName: pdf_to_img.py
# @Software: PyCharm
# @Blog    ：https://bornforthis.cn/
# code is far away from bugs with the god animal protecting
#    I love animals. They taste delicious.
from pathlib import Path
from datetime import datetime
import traceback

# ====== 需要你根据实际情况修改的常量（路径直接写在代码中）======
INPUT_DIR = Path(r"./PDF-DATA")  # 放 PDF 的文件夹
OUTPUT_ROOT = Path(r"./RESULT")  # 总输出根目录：RESULT
RECURSIVE = False  # 是否递归遍历子目录
DPI = 600  # “最高清”建议600；可调 300/600/1200
SAVE_PNG = True  # 是否导出 PNG
SAVE_JPG = False  # 是否导出 JPG
JPG_QUALITY = 95  # JPG 质量(1-100)，95 接近无损但体积可控
LOG_FILE = OUTPUT_ROOT / "log.txt"  # 失败日志
# =====================================================

# 依赖 PyMuPDF
try:
    import fitz  # PyMuPDF
except ImportError as e:
    raise SystemExit("未安装 PyMuPDF，请先运行：pip install pymupdf") from e


def strip_pdf_suffixes(name):
    """多次去掉结尾的 .pdf（应对 xxx.pdf.pdf 这种情况）"""
    base = name
    while base.lower().endswith(".pdf"):
        base = base[:-4]
    return base or "untitled"


def log_error(msg):
    LOG_FILE.parent.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(f"[{ts}] {msg}\n")


def convert_pdf_to_images(pdf_path: Path):
    """
    将单个PDF按页导出为 PNG+JPG。
    出错会写日志并继续其它页 / 文件。
    """
    try:
        doc = fitz.open(pdf_path)
    except Exception as e:
        log_error(f"打开失败: {pdf_path} | {e}")
        return

    try:
        base_name = strip_pdf_suffixes(pdf_path.name)
        out_dir = OUTPUT_ROOT / base_name
        out_dir.mkdir(parents=True, exist_ok=True)

        zoom = DPI / 72.0
        mat = fitz.Matrix(zoom, zoom)

        total = len(doc)
        for i in range(total):
            page_no = i + 1  # 1-based
            try:
                page = doc.load_page(i)
                pix = page.get_pixmap(matrix=mat, alpha=False)  # alpha=False 更通用

                if SAVE_PNG:
                    (out_dir / f"{page_no}.png").write_bytes(pix.tobytes("png"))

                if SAVE_JPG:
                    # PyMuPDF 支持 jpg_quality 参数
                    pix.save(str(out_dir / f"{page_no}.jpg"), jpg_quality=JPG_QUALITY)

            except Exception as e_page:
                # 单页失败，记录后继续下一页
                log_error(f"页面失败: {pdf_path} | 页 {page_no}/{total} | {e_page}\n{traceback.format_exc()}")
                continue

    finally:
        try:
            doc.close()
        except Exception:
            pass


def find_all_pdfs(root, recursive=False):
    if root.is_file() and root.suffix.lower() == ".pdf":
        return [root]
    if root.is_dir():
        return sorted(root.rglob("*.pdf") if recursive else root.glob("*.pdf"))
    return []


def main():
    if not INPUT_DIR.exists():
        raise SystemExit(f"输入目录不存在：{INPUT_DIR}")
    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
    # 本次运行写一个分隔头
    log_error("=== 本次转换开始 ===")

    pdfs = find_all_pdfs(INPUT_DIR, recursive=RECURSIVE)
    if not pdfs:
        log_error(f"未找到PDF：{INPUT_DIR}（递归={RECURSIVE}）")
        print("未找到任何PDF。")
        return

    print(f"共找到 {len(pdfs)} 个PDF，开始转换（DPI={DPI}，PNG={SAVE_PNG}，JPG={SAVE_JPG}, Q={JPG_QUALITY})...")
    for idx, pdf in enumerate(pdfs, 1):
        print(f"[{idx}/{len(pdfs)}] {pdf}")
        try:
            convert_pdf_to_images(pdf)
        except Exception as e_file:
            # 理论上 convert 已经自带容错，这里兜底
            log_error(f"文件失败: {pdf} | {e_file}\n{traceback.format_exc()}")
            continue
    log_error("=== 本次转换结束 ===")
    print("全部完成。失败详情见：", LOG_FILE)


if __name__ == '__main__':
    main()