OCRパイプラインの構造

iOSでONNX Runtimeを使ったOCRパイプラインは、大まかに以下の流れで動作します。

  1. 画像全体に対する文字領域検出(Detection)
  2. 検出された各領域に対する文字認識(Recognition)
  3. 読み順の推定と結合

この中で、検出は画像全体に対して1回だけ実行されます。一方、認識は検出された領域の数だけ繰り返されます。領域数が多い画像では、認識処理が全体の処理時間の大部分を占めることになります。

逐次処理の問題

認識処理を単純なforループで逐次実行すると、処理時間は領域数にほぼ比例します。

for det in detections {
    let rect = CGRect(x: max(0, det.box[0]), y: max(0, det.box[1]),
                      width: max(1, det.box[2] - det.box[0]),
                      height: max(1, det.box[3] - det.box[1]))
    if let cropped = image.cropping(to: rect) {
        _ = try recognizer.recognize(image: cropped)
    }
}

たとえば98領域が検出された画像の場合、1領域あたり約80msとすると合計で約8秒かかります。

withThrowingTaskGroupによる並列化

Swift ConcurrencyのwithThrowingTaskGroupを使うと、各領域の認識処理を並列に実行できます。

let results = try await withThrowingTaskGroup(of: (Int, String).self) { group in
    for i in 0..<detections.count {
        let box = detections[i].box
        let cropRect = CGRect(
            x: max(0, box[0]),
            y: max(0, box[1]),
            width: max(1, box[2] - box[0]),
            height: max(1, box[3] - box[1])
        )
        guard let cropped = image.cropping(to: cropRect) else { continue }
        group.addTask {
            let text = try recognizer.recognize(image: cropped)
            return (i, text)
        }
    }
    var texts: [(Int, String)] = []
    for try await result in group {
        texts.append(result)
    }
    return texts
}
for (i, text) in results {
    recognized[i].text = text
}

各タスクはインデックスとともに結果を返すため、完了順序に関係なく元の順序を保持できます。

ベンチマーク結果

iOSシミュレータ上で、逐次処理と並列処理の実行時間を比較しました。

古典籍(源氏物語、6642x4990px、21領域)

出典: 東京大学デジタルアーカイブ 源氏物語

東京大学デジタルアーカイブ 源氏物語

方式処理時間
逐次4.55s
並列3.24s
高速化率1.4x

近代活字(NDLデジタルコレクション、6890x4706px、98領域)

出典: 国立国会図書館デジタルコレクション

国立国会図書館デジタルコレクション

方式処理時間
逐次7.59s
並列1.13s
高速化率6.7x

高速化率が異なる理由

21領域の場合と98領域の場合で高速化率に差が出ています。

Swift Concurrencyの協調スレッドプール(Cooperative Thread Pool)は、同時実行数をCPUコア数に自動制限します。ONNX Runtimeのセッションは内部でロックを取るため、同時に複数のタスクが推論を実行しようとするとロック競合が発生します。

領域数が少ない場合(21領域)、ロック競合のオーバーヘッドが相対的に大きくなり、高速化率は1.4倍にとどまります。領域数が多い場合(98領域)、利用可能なコアを効率的に使い切ることができ、6.7倍の高速化が得られています。

なお、withThrowingTaskGroupはSwiftランタイムが並行度を管理するため、スレッド爆発のリスクはありません。

補足

上記の計測はすべてiOSシミュレータ上で行ったものです。実機での性能は異なる可能性があります。