wehub-resource-sync
117 lines
3.9 KiB
Swift
117 lines
3.9 KiB
Swift
import AppKit
|
|
import Foundation
|
|
import MacToolsPluginKit
|
|
import Vision
|
|
|
|
private final class TranslatorOCRBundleToken {}
|
|
|
|
@MainActor
|
|
protocol OCRTextRecognizing {
|
|
func recognizeText(in image: NSImage) async throws -> OCRTextRecognitionResult
|
|
}
|
|
|
|
enum OCRTextRecognitionError: Error, Equatable, LocalizedError, Sendable {
|
|
case invalidImage
|
|
case emptyResult
|
|
case requestFailed(String)
|
|
|
|
var errorDescription: String? {
|
|
switch self {
|
|
case .invalidImage:
|
|
return PluginLocalization(bundle: Bundle(for: TranslatorOCRBundleToken.self)).string(
|
|
"ocr.error.invalidImage",
|
|
defaultValue: "无法读取截图。"
|
|
)
|
|
case .emptyResult:
|
|
return PluginLocalization(bundle: Bundle(for: TranslatorOCRBundleToken.self)).string(
|
|
"ocr.error.emptyResult",
|
|
defaultValue: "截图中没有识别到文字。"
|
|
)
|
|
case let .requestFailed(message):
|
|
return message
|
|
}
|
|
}
|
|
}
|
|
|
|
struct VisionOCRTextRecognizer: OCRTextRecognizing {
|
|
func recognizeText(in image: NSImage) async throws -> OCRTextRecognitionResult {
|
|
guard let cgImage = image.cgImageForOCR else {
|
|
throw OCRTextRecognitionError.invalidImage
|
|
}
|
|
|
|
let lines = try await recognizeLines(
|
|
in: cgImage,
|
|
recognitionLanguages: nil,
|
|
automaticallyDetectsLanguage: true
|
|
)
|
|
|
|
let fallbackLines: [OCRRecognizedLine]
|
|
if lines.isEmpty {
|
|
fallbackLines = try await recognizeLines(
|
|
in: cgImage,
|
|
recognitionLanguages: ["ja-JP"],
|
|
automaticallyDetectsLanguage: false
|
|
)
|
|
} else {
|
|
fallbackLines = lines
|
|
}
|
|
|
|
let text = OCRTextMerge.merge(fallbackLines)
|
|
guard !text.isEmpty else {
|
|
throw OCRTextRecognitionError.emptyResult
|
|
}
|
|
|
|
return OCRTextRecognitionResult(text: text, lines: fallbackLines)
|
|
}
|
|
|
|
private func recognizeLines(
|
|
in cgImage: CGImage,
|
|
recognitionLanguages: [String]?,
|
|
automaticallyDetectsLanguage: Bool
|
|
) async throws -> [OCRRecognizedLine] {
|
|
try await withCheckedThrowingContinuation { continuation in
|
|
DispatchQueue.global(qos: .userInitiated).async {
|
|
let request = VNRecognizeTextRequest()
|
|
request.recognitionLevel = .accurate
|
|
request.usesLanguageCorrection = true
|
|
|
|
if #available(macOS 13.0, *) {
|
|
request.automaticallyDetectsLanguage = automaticallyDetectsLanguage
|
|
}
|
|
|
|
if let recognitionLanguages {
|
|
request.recognitionLanguages = recognitionLanguages
|
|
}
|
|
|
|
do {
|
|
try VNImageRequestHandler(cgImage: cgImage).perform([request])
|
|
let recognizedLines = (request.results ?? []).compactMap { observation -> OCRRecognizedLine? in
|
|
guard let candidate = observation.topCandidates(1).first else {
|
|
return nil
|
|
}
|
|
return OCRRecognizedLine(
|
|
text: candidate.string,
|
|
boundingBox: observation.boundingBox,
|
|
confidence: candidate.confidence
|
|
)
|
|
}
|
|
continuation.resume(returning: recognizedLines)
|
|
} catch {
|
|
continuation.resume(throwing: OCRTextRecognitionError.requestFailed(error.localizedDescription))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private extension NSImage {
|
|
var cgImageForOCR: CGImage? {
|
|
guard size.width > 0, size.height > 0 else {
|
|
return nil
|
|
}
|
|
|
|
var proposedRect = CGRect(origin: .zero, size: size)
|
|
return cgImage(forProposedRect: &proposedRect, context: nil, hints: nil)
|
|
}
|
|
}
|