research_ocr_2026_fall.swift 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. import AppKit
  2. import Foundation
  3. import Vision
  4. struct OCRLine: Codable {
  5. let text: String
  6. let x: Double
  7. let y: Double
  8. let width: Double
  9. let height: Double
  10. }
  11. struct OCRPage: Codable {
  12. let grade: String
  13. let file: String
  14. let path: String
  15. let lines: [OCRLine]
  16. }
  17. let args = CommandLine.arguments
  18. guard args.count == 3 else {
  19. fputs("usage: research_ocr_2026_fall <source-root> <output-json>\n", stderr)
  20. exit(2)
  21. }
  22. let sourceRoot = URL(fileURLWithPath: args[1], isDirectory: true)
  23. let outputURL = URL(fileURLWithPath: args[2])
  24. let manager = FileManager.default
  25. let allowedExtensions = Set(["png", "jpg", "jpeg"])
  26. let gradeURLs = try manager.contentsOfDirectory(
  27. at: sourceRoot,
  28. includingPropertiesForKeys: nil,
  29. options: [.skipsHiddenFiles]
  30. ).filter { $0.hasDirectoryPath }.sorted { $0.lastPathComponent < $1.lastPathComponent }
  31. var pages: [OCRPage] = []
  32. for gradeURL in gradeURLs {
  33. let imageURLs = try manager.contentsOfDirectory(
  34. at: gradeURL,
  35. includingPropertiesForKeys: nil,
  36. options: [.skipsHiddenFiles]
  37. ).filter { allowedExtensions.contains($0.pathExtension.lowercased()) }
  38. .sorted { $0.lastPathComponent < $1.lastPathComponent }
  39. for (index, imageURL) in imageURLs.enumerated() {
  40. guard let image = NSImage(contentsOf: imageURL),
  41. let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
  42. fputs("failed to load \(imageURL.path)\n", stderr)
  43. exit(1)
  44. }
  45. let request = VNRecognizeTextRequest()
  46. request.recognitionLevel = .accurate
  47. request.usesLanguageCorrection = true
  48. request.recognitionLanguages = ["zh-Hans", "en-US"]
  49. request.minimumTextHeight = 0.006
  50. try VNImageRequestHandler(cgImage: cgImage, options: [:]).perform([request])
  51. let lines = (request.results ?? []).compactMap { observation -> OCRLine? in
  52. guard let candidate = observation.topCandidates(1).first else { return nil }
  53. let box = observation.boundingBox
  54. return OCRLine(
  55. text: candidate.string,
  56. x: Double(box.minX),
  57. y: Double(box.minY),
  58. width: Double(box.width),
  59. height: Double(box.height)
  60. )
  61. }.sorted {
  62. let leftY = Int($0.y * 1000)
  63. let rightY = Int($1.y * 1000)
  64. if leftY != rightY { return leftY > rightY }
  65. return $0.x < $1.x
  66. }
  67. pages.append(OCRPage(
  68. grade: gradeURL.lastPathComponent,
  69. file: imageURL.lastPathComponent,
  70. path: imageURL.path,
  71. lines: lines
  72. ))
  73. fputs("[\(gradeURL.lastPathComponent)] \(index + 1)/\(imageURLs.count) \(imageURL.lastPathComponent)\n", stderr)
  74. }
  75. }
  76. let encoder = JSONEncoder()
  77. encoder.outputFormatting = [.prettyPrinted, .sortedKeys, .withoutEscapingSlashes]
  78. try encoder.encode(pages).write(to: outputURL, options: .atomic)
  79. fputs("wrote \(pages.count) pages to \(outputURL.path)\n", stderr)