128 lines
3.7 KiB
Python
128 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
SWIFT_CODE = r'''
|
|
import Foundation
|
|
import Vision
|
|
import AppKit
|
|
|
|
let path = CommandLine.arguments[1]
|
|
let url = URL(fileURLWithPath: path)
|
|
guard let img = NSImage(contentsOf: url),
|
|
let tiff = img.tiffRepresentation,
|
|
let bitmap = NSBitmapImageRep(data: tiff),
|
|
let cg = bitmap.cgImage else {
|
|
fputs("image load failed\n", stderr)
|
|
exit(3)
|
|
}
|
|
|
|
struct Line {
|
|
let text: String
|
|
let x: CGFloat
|
|
let y: CGFloat
|
|
}
|
|
|
|
var lines: [Line] = []
|
|
var requestError: Error?
|
|
let request = VNRecognizeTextRequest { req, err in
|
|
if let err = err {
|
|
requestError = err
|
|
return
|
|
}
|
|
let observations = (req.results as? [VNRecognizedTextObservation]) ?? []
|
|
lines = observations.compactMap { observation in
|
|
guard let candidate = observation.topCandidates(1).first else { return nil }
|
|
let text = candidate.string.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
guard !text.isEmpty else { return nil }
|
|
return Line(text: text, x: observation.boundingBox.minX, y: observation.boundingBox.maxY)
|
|
}
|
|
}
|
|
request.recognitionLevel = .accurate
|
|
request.usesLanguageCorrection = true
|
|
request.recognitionLanguages = ["zh-Hans", "en-US"]
|
|
|
|
let handler = VNImageRequestHandler(cgImage: cg, options: [:])
|
|
do {
|
|
try handler.perform([request])
|
|
} catch {
|
|
fputs("\(error.localizedDescription)\n", stderr)
|
|
exit(4)
|
|
}
|
|
if let requestError = requestError {
|
|
fputs("\(requestError.localizedDescription)\n", stderr)
|
|
exit(5)
|
|
}
|
|
|
|
let sorted = lines.sorted {
|
|
if abs($0.y - $1.y) > 0.012 {
|
|
return $0.y > $1.y
|
|
}
|
|
return $0.x < $1.x
|
|
}
|
|
for line in sorted {
|
|
print(line.text)
|
|
}
|
|
'''
|
|
|
|
|
|
def page_number(path: Path) -> int:
|
|
match = re.search(r"(\d+)", path.stem)
|
|
return int(match.group(1)) if match else 0
|
|
|
|
|
|
def run_ocr(image: Path, module_cache: Path) -> tuple[str, str | None]:
|
|
module_cache.mkdir(parents=True, exist_ok=True)
|
|
env = os.environ.copy()
|
|
env["CLANG_MODULE_CACHE_PATH"] = str(module_cache)
|
|
proc = subprocess.run(
|
|
["swift", "-e", SWIFT_CODE, str(image)],
|
|
text=True,
|
|
capture_output=True,
|
|
check=False,
|
|
env=env,
|
|
)
|
|
if proc.returncode != 0:
|
|
return "", proc.stderr.strip() or f"swift exited {proc.returncode}"
|
|
return proc.stdout.strip(), None
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("output", type=Path)
|
|
parser.add_argument("images", nargs="+", type=Path)
|
|
args = parser.parse_args()
|
|
|
|
images = sorted(args.images, key=lambda p: (page_number(p), str(p)))
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
module_cache = args.output.parent.parent / "tmp" / "swift_module_cache"
|
|
|
|
with args.output.open("w", encoding="utf-8") as out:
|
|
out.write("# GB/T 33923-2017 OCR 文本版\n\n")
|
|
out.write("> 说明：此文本由 PDF 页面图像 OCR 生成，用于阅读、搜索和复制。公式、表格、上下标和个别专业符号可能需要对照原 PDF 校核。\n\n")
|
|
for index, image in enumerate(images, 1):
|
|
page = page_number(image)
|
|
print(f"OCR page {index}/{len(images)}: {image.name}", file=sys.stderr, flush=True)
|
|
text, error = run_ocr(image, module_cache)
|
|
out.write(f"\n\n## 第 {page} 页\n\n")
|
|
if error:
|
|
out.write(f"[OCR 失败：{error}]\n")
|
|
elif text:
|
|
out.write(text)
|
|
out.write("\n")
|
|
else:
|
|
out.write("[OCR 未识别出文本]\n")
|
|
out.flush()
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|