96 lines
3.2 KiB
Swift
96 lines
3.2 KiB
Swift
import AppKit
|
|
import Foundation
|
|
import Vision
|
|
|
|
struct OCRLine {
|
|
let text: String
|
|
let x: CGFloat
|
|
let y: CGFloat
|
|
}
|
|
|
|
func pageNumber(from path: String) -> Int {
|
|
let name = URL(fileURLWithPath: path).deletingPathExtension().lastPathComponent
|
|
let digits = name.replacingOccurrences(of: "[^0-9]", with: "", options: .regularExpression)
|
|
return Int(digits) ?? 0
|
|
}
|
|
|
|
func loadCGImage(_ path: String) -> CGImage? {
|
|
let url = URL(fileURLWithPath: path)
|
|
guard let image = NSImage(contentsOf: url),
|
|
let tiff = image.tiffRepresentation,
|
|
let bitmap = NSBitmapImageRep(data: tiff) else {
|
|
return nil
|
|
}
|
|
return bitmap.cgImage
|
|
}
|
|
|
|
func recognize(_ path: String) throws -> [OCRLine] {
|
|
guard let cgImage = loadCGImage(path) else {
|
|
throw NSError(domain: "VisionOCR", code: 1, userInfo: [NSLocalizedDescriptionKey: "Could not load image: \(path)"])
|
|
}
|
|
|
|
var lines: [OCRLine] = []
|
|
var requestError: Error?
|
|
let request = VNRecognizeTextRequest { request, error in
|
|
if let error = error {
|
|
requestError = error
|
|
return
|
|
}
|
|
let observations = (request.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 OCRLine(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: cgImage, options: [:])
|
|
try handler.perform([request])
|
|
if let requestError = requestError {
|
|
throw requestError
|
|
}
|
|
|
|
return lines.sorted {
|
|
if abs($0.y - $1.y) > 0.012 {
|
|
return $0.y > $1.y
|
|
}
|
|
return $0.x < $1.x
|
|
}
|
|
}
|
|
|
|
let args = Array(CommandLine.arguments.dropFirst())
|
|
guard args.count >= 2 else {
|
|
FileHandle.standardError.write(Data("Usage: vision_ocr.swift OUTPUT.md IMAGE...\n".utf8))
|
|
exit(2)
|
|
}
|
|
|
|
let output = args[0]
|
|
let imagePaths = args.dropFirst().sorted {
|
|
let lhs = pageNumber(from: $0)
|
|
let rhs = pageNumber(from: $1)
|
|
return lhs == rhs ? $0 < $1 : lhs < rhs
|
|
}
|
|
|
|
var result = "# GB/T 33923-2017 OCR 文本版\n\n"
|
|
result += "> 说明：此文本由 PDF 页面图像 OCR 生成，用于阅读、搜索和复制。公式、表格、上下标和个别专业符号可能需要对照原 PDF 校核。\n\n"
|
|
|
|
for (index, path) in imagePaths.enumerated() {
|
|
let page = pageNumber(from: path)
|
|
FileHandle.standardError.write(Data("OCR page \(index + 1)/\(imagePaths.count): \(URL(fileURLWithPath: path).lastPathComponent)\n".utf8))
|
|
do {
|
|
let lines = try recognize(path)
|
|
result += "\n\n## 第 \(page) 页\n\n"
|
|
result += lines.map(\.text).joined(separator: "\n")
|
|
result += "\n"
|
|
} catch {
|
|
result += "\n\n## 第 \(page) 页\n\n[OCR 失败：\(error.localizedDescription)]\n"
|
|
}
|
|
}
|
|
|
|
try result.write(toFile: output, atomically: true, encoding: .utf8)
|