Split project into frontend and backend
This commit is contained in:
@@ -0,0 +1,95 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user