ocr_image_text.swift 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. import Foundation
  2. import Vision
  3. import AppKit
  4. struct OcrLine: Encodable {
  5. let text: String
  6. let x: Double
  7. let y: Double
  8. let width: Double
  9. let height: Double
  10. }
  11. let args = CommandLine.arguments
  12. guard args.count >= 2 else {
  13. fputs("usage: swift ocr_image_text.swift <image>\\n", stderr)
  14. exit(2)
  15. }
  16. let imageURL = URL(fileURLWithPath: args[1])
  17. guard let image = NSImage(contentsOf: imageURL),
  18. let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
  19. fputs("failed to load image: \\(args[1])\\n", stderr)
  20. exit(1)
  21. }
  22. let request = VNRecognizeTextRequest()
  23. request.recognitionLevel = .accurate
  24. request.usesLanguageCorrection = true
  25. request.recognitionLanguages = ["zh-Hans", "en-US"]
  26. request.minimumTextHeight = 0.006
  27. let handler = VNImageRequestHandler(cgImage: cgImage, options: [:])
  28. try handler.perform([request])
  29. let observations = request.results ?? []
  30. let lines = observations.compactMap { observation -> OcrLine? in
  31. guard let candidate = observation.topCandidates(1).first else { return nil }
  32. let box = observation.boundingBox
  33. return OcrLine(
  34. text: candidate.string,
  35. x: Double(box.minX),
  36. y: Double(box.minY),
  37. width: Double(box.width),
  38. height: Double(box.height)
  39. )
  40. }.sorted {
  41. let y0 = Int($0.y * 1000)
  42. let y1 = Int($1.y * 1000)
  43. if y0 != y1 { return y0 > y1 }
  44. return $0.x < $1.x
  45. }
  46. let data = try JSONEncoder().encode(lines)
  47. FileHandle.standardOutput.write(data)