| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091 |
- import AppKit
- import Foundation
- import Vision
- struct OCRLine: Codable {
- let text: String
- let x: Double
- let y: Double
- let width: Double
- let height: Double
- }
- struct OCRPage: Codable {
- let grade: String
- let file: String
- let path: String
- let lines: [OCRLine]
- }
- let args = CommandLine.arguments
- guard args.count == 3 else {
- fputs("usage: research_ocr_2026_fall <source-root> <output-json>\n", stderr)
- exit(2)
- }
- let sourceRoot = URL(fileURLWithPath: args[1], isDirectory: true)
- let outputURL = URL(fileURLWithPath: args[2])
- let manager = FileManager.default
- let allowedExtensions = Set(["png", "jpg", "jpeg"])
- let gradeURLs = try manager.contentsOfDirectory(
- at: sourceRoot,
- includingPropertiesForKeys: nil,
- options: [.skipsHiddenFiles]
- ).filter { $0.hasDirectoryPath }.sorted { $0.lastPathComponent < $1.lastPathComponent }
- var pages: [OCRPage] = []
- for gradeURL in gradeURLs {
- let imageURLs = try manager.contentsOfDirectory(
- at: gradeURL,
- includingPropertiesForKeys: nil,
- options: [.skipsHiddenFiles]
- ).filter { allowedExtensions.contains($0.pathExtension.lowercased()) }
- .sorted { $0.lastPathComponent < $1.lastPathComponent }
- for (index, imageURL) in imageURLs.enumerated() {
- guard let image = NSImage(contentsOf: imageURL),
- let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
- fputs("failed to load \(imageURL.path)\n", stderr)
- exit(1)
- }
- let request = VNRecognizeTextRequest()
- request.recognitionLevel = .accurate
- request.usesLanguageCorrection = true
- request.recognitionLanguages = ["zh-Hans", "en-US"]
- request.minimumTextHeight = 0.006
- try VNImageRequestHandler(cgImage: cgImage, options: [:]).perform([request])
- let lines = (request.results ?? []).compactMap { observation -> OCRLine? in
- guard let candidate = observation.topCandidates(1).first else { return nil }
- let box = observation.boundingBox
- return OCRLine(
- text: candidate.string,
- x: Double(box.minX),
- y: Double(box.minY),
- width: Double(box.width),
- height: Double(box.height)
- )
- }.sorted {
- let leftY = Int($0.y * 1000)
- let rightY = Int($1.y * 1000)
- if leftY != rightY { return leftY > rightY }
- return $0.x < $1.x
- }
- pages.append(OCRPage(
- grade: gradeURL.lastPathComponent,
- file: imageURL.lastPathComponent,
- path: imageURL.path,
- lines: lines
- ))
- fputs("[\(gradeURL.lastPathComponent)] \(index + 1)/\(imageURLs.count) \(imageURL.lastPathComponent)\n", stderr)
- }
- }
- let encoder = JSONEncoder()
- encoder.outputFormatting = [.prettyPrinted, .sortedKeys, .withoutEscapingSlashes]
- try encoder.encode(pages).write(to: outputURL, options: .atomic)
- fputs("wrote \(pages.count) pages to \(outputURL.path)\n", stderr)
|