OneClip 原生 PP-OCRv5 ONNX 方案解析¶

let env = try ORTEnv(loggingLevel: .warning)
let options = try ORTSessionOptions()
let detSession = try ORTSession(env: env, modelPath: detURL.path, sessionOptions: options)
let recSession = try ORTSession(env: env, modelPath: recURL.path, sessionOptions: options)
let clsSession = try ORTSession(env: env, modelPath: clsURL.path, sessionOptions: options)

for timestep in 0..<timesteps {
    // 取当前时间步最优类别
    guard bestIndex != 0, bestIndex != previousIndex, bestIndex < characters.count else { continue }
    text += characters[bestIndex]
    confidenceSum += bestScore
    characterCount += 1
}

let digest = SHA256.hash(data: data).map { String(format: "%02x", $0) }.joined()
guard digest.caseInsensitiveCompare(expectedSHA256) == .orderedSame else {
    throw OCRPluginError.downloadFailed("\(displayName) SHA-256 校验失败")
}

import OnnxRuntimeBindings
import CoreGraphics

// 你的模型目录（det/rec/cls 三个 onnx + 字典都在这里）
let modelsDir = try FileManager.default
    .url(for: .applicationSupportDirectory, in: .userDomainMask, appropriateFor: nil, create: true)
    .appendingPathComponent("ocr", isDirectory: true)

let runtime = NativePPOCRRuntime()
let image: CGImage = /* 你的截图 / 剪贴板图片 / 文件图片 */

// 识别，返回合并后的全部文字
let text = try await runtime.recognize(
    image,
    modelsDirectory: modelsDir,
    modelVersion: .v6Small   // 或 .v5Mobile
)
print(text)

// 若要逐行坐标与置信度，用 recognizeResult：
let result = try await runtime.recognizeResult(
    image,
    modelsDirectory: modelsDir,
    modelVersion: .v6Small
)
for line in result.lines {
    print(line.text, line.score, line.box)
}

struct SessionInfo {
    let session: ORTSession
    let inputName: String
    let outputNames: Set<String>
}

func loadSession(_ env: ORTEnv, _ url: URL) throws -> SessionInfo {
    let session = try ORTSession(env: env, modelPath: url.path, sessionOptions: ORTSessionOptions())
    let inputName = try session.inputNames().first ?? "input"
    let outputNames = try session.outputNames()          // 返回 Set<String>
    return SessionInfo(session: session, inputName: inputName, outputNames: outputNames)
}

func run(_ info: SessionInfo, input: [Float], shape: [Int]) throws -> [Float] {
    let data = NSMutableData(bytes: input, length: input.count * MemoryLayout<Float>.stride)
    let tensor = try ORTValue(tensorData: data, elementType: .float, shape: shape.map(NSNumber(value:)))
    let outputs = try info.session.run(
        withInputs: [info.inputName: tensor],
        outputNames: info.outputNames,
        runOptions: nil
    )
    let value = try outputs.values.first!
    let tensorData = try value.tensorData() as Data
    return tensorData.withUnsafeBytes { Array($0.bindMemory(to: Float.self)) }
}

func makeDetectionInput(_ image: CGImage) -> (data: [Float], height: Int, width: Int) {
    let minSide = 736, maxSide = 2048
    let ow = image.width, oh = image.height
    var scale: CGFloat = 1
    if min(ow, oh) < minSide { scale = CGFloat(minSide) / CGFloat(min(ow, oh)) }
    if CGFloat(max(ow, oh)) * scale > CGFloat(maxSide) { scale = CGFloat(maxSide) / CGFloat(max(ow, oh)) }
    let width = max(32, Int((CGFloat(ow) * scale / 32).rounded()) * 32)
    let height = max(32, Int((CGFloat(oh) * scale / 32).rounded()) * 32)
    let pixels = rgbaPixels(image, width: width, height: height)   // [r,g,b,a] * N
    var chw = [Float](repeating: 0, count: width * height * 3)
    for i in 0..<width*height {
        chw[i]                   = Float(pixels[i*4+2]) / 127.5 - 1   // B
        chw[width*height + i]    = Float(pixels[i*4+1]) / 127.5 - 1   // G
        chw[width*height*2 + i]  = Float(pixels[i*4])   / 127.5 - 1   // R
    }
    return (chw, height, width)
}

func detectionBoxes(
    probabilityMap: [Float], mapWidth: Int, mapHeight: Int, unclipRatio: CGFloat
) -> [[CGPoint]] {
    let threshold: Float = 0.3
    let boxThreshold: Float = 0.5
    var visited = [Bool](repeating: false, count: mapWidth * mapHeight)
    let neighbors = [(-1,-1),(0,-1),(1,-1),(-1,0),(1,0),(-1,1),(0,1),(1,1)]
    var boxes: [[CGPoint]] = []

    for seed in 0..<mapWidth*mapHeight {
        if visited[seed] || probabilityMap[seed] <= threshold { continue }
        visited[seed] = true
        var queue = [seed], head = 0, pixelCount = 1, boundary: [CGPoint] = []
        while head < queue.count {
            let cur = queue[head]; head += 1
            let x = cur % mapWidth, y = cur / mapWidth
            var isEdge = false
            for (dx, dy) in neighbors {
                let nx = x+dx, ny = y+dy
                guard nx >= 0, nx < mapWidth, ny >= 0, ny < mapHeight else { isEdge = true; continue }
                let next = ny*mapWidth + nx
                if probabilityMap[next] <= threshold { isEdge = true; continue }
                if !visited[next] { visited[next] = true; queue.append(next); pixelCount += 1 }
            }
            if isEdge { boundary.append(CGPoint(x: CGFloat(x)+0.5, y: CGFloat(y)+0.5)) }
        }
        guard pixelCount >= 8, boundary.count >= 3 else { continue }
        let hull = convexHull(boundary)
        guard hull.count >= 3 else { continue }
        let score = polygonScore(hull, probabilityMap: probabilityMap, width: mapWidth, height: mapHeight)
        guard score >= boxThreshold else { continue }             // 框内概率均值需达标
        let expanded = unclip(hull, ratio: unclipRatio)           // 见 ④
        if let quad = minimumAreaRectangle(for: expanded) { boxes.append(quad) }
    }
    return boxes
}

func polygonScore(_ polygon: [CGPoint], probabilityMap: [Float], width: Int, height: Int) -> Float {
    guard let first = polygon.first else { return 0 }
    let minX = max(0, Int(floor(polygon.dropFirst().reduce(first.x) { min($0, $1.x) })))
    let maxX = min(width-1, Int(ceil(polygon.dropFirst().reduce(first.x) { max($0, $1.x) })))
    let minY = max(0, Int(floor(polygon.dropFirst().reduce(first.y) { min($0, $1.y) })))
    let maxY = min(height-1, Int(ceil(polygon.dropFirst().reduce(first.y) { max($0, $1.y) })))
    guard minX <= maxX, minY <= maxY else { return 0 }
    var total: Float = 0, count = 0
    for y in minY...maxY {
        for x in minX...maxX where pointInPolygon(CGPoint(x: CGFloat(x)+0.5, y: CGFloat(y)+0.5), polygon: polygon) {
            total += probabilityMap[y*width + x]; count += 1
        }
    }
    return count == 0 ? 0 : total / Float(count)
}

func pointInPolygon(_ point: CGPoint, polygon: [CGPoint]) -> Bool {
    var inside = false, prev = polygon.count - 1
    for cur in polygon.indices {
        let a = polygon[cur], b = polygon[prev]
        if (a.y > point.y) != (b.y > point.y) {
            let crossX = (b.x - a.x) * (point.y - a.y) / (b.y - a.y) + a.x
            if point.x < crossX { inside.toggle() }
        }
        prev = cur
    }
    return inside
}

func unclip(_ polygon: [CGPoint], ratio: CGFloat) -> [CGPoint] {
    guard polygon.count >= 3 else { return polygon }
    var area: CGFloat = 0, perimeter: CGFloat = 0
    for i in polygon.indices {
        let n = polygon[(i+1) % polygon.count]
        area += polygon[i].x*n.y - n.x*polygon[i].y
        perimeter += hypot(n.x-polygon[i].x, n.y-polygon[i].y)
    }
    area /= 2
    guard perimeter > 0 else { return polygon }
    let distance = abs(area) * ratio / perimeter
    let points = area < 0 ? Array(polygon.reversed()) : polygon
    let lines = points.indices.map { i -> (point: CGPoint, direction: CGPoint) in
        let s = points[i], e = points[(i+1)%points.count]
        let dir = CGPoint(x: e.x-s.x, y: e.y-s.y)
        let len = max(hypot(dir.x, dir.y), 0.0001)
        let out = CGPoint(x: dir.y/len, y: -dir.x/len)
        return (CGPoint(x: s.x+out.x*distance, y: s.y+out.y*distance), dir)
    }
    return lines.indices.map { i in
        let p = lines[(i+lines.count-1) % lines.count]
        let c = lines[i]
        return lineIntersection(p, c) ?? c.point
    }
}

func lineIntersection(_ lhs: (point: CGPoint, direction: CGPoint),
                      _ rhs: (point: CGPoint, direction: CGPoint)) -> CGPoint? {
    let denom = lhs.direction.x*rhs.direction.y - lhs.direction.y*rhs.direction.x
    guard abs(denom) > 0.0001 else { return nil }
    let delta = CGPoint(x: rhs.point.x-lhs.point.x, y: rhs.point.y-lhs.point.y)
    let t = (delta.x*rhs.direction.y - delta.y*rhs.direction.x) / denom
    return CGPoint(x: lhs.point.x + lhs.direction.x*t, y: lhs.point.y + lhs.direction.y*t)
}

func minimumAreaRectangle(for points: [CGPoint]) -> [CGPoint]? {
    guard points.count >= 3 else { return nil }
    var bestArea = CGFloat.greatestFiniteMagnitude, bestCorners: [CGPoint]?
    for i in points.indices {
        let n = points[(i+1) % points.count]
        let a = atan2(n.y-points[i].y, n.x-points[i].x)
        let cs = cos(a), sn = sin(a)
        var minX = CGFloat.infinity, maxX = -CGFloat.infinity, minY = CGFloat.infinity, maxY = -CGFloat.infinity
        for p in points {
            let x = p.x*cs + p.y*sn, y = -p.x*sn + p.y*cs
            minX = min(minX, x); maxX = max(maxX, x); minY = min(minY, y); maxY = max(maxY, y)
        }
        let area = (maxX-minX)*(maxY-minY)
        guard area < bestArea else { continue }
        bestArea = area
        bestCorners = [CGPoint(x:minX,y:minY), CGPoint(x:maxX,y:minY),
                       CGPoint(x:maxX,y:maxY), CGPoint(x:minX,y:maxY)]
            .map { CGPoint(x: $0.x*cs - $0.y*sn, y: $0.x*sn + $0.y*cs) }
    }
    return bestCorners
}

func makeRecognitionInput(_ image: CGImage) -> (data: [Float], shape: [Int]) {
    let targetHeight = 48
    let w = max(1, min(3200, Int(ceil(CGFloat(targetHeight) * CGFloat(image.width) / CGFloat(image.height)))))
    let canvas = max(320, w)
    let pixels = rgbaPixels(image, width: w, height: targetHeight)
    var chw = [Float](repeating: 0, count: targetHeight * canvas * 3)
    for y in 0..<targetHeight {
        for x in 0..<w {
            let s = (y*w + x)*4, d = y*canvas + x
            chw[d]                       = Float(pixels[s+2]) / 127.5 - 1
            chw[targetHeight*canvas + d] = Float(pixels[s+1]) / 127.5 - 1
            chw[targetHeight*canvas*2 + d] = Float(pixels[s]) / 127.5 - 1
        }
    }
    return (chw, [1, 3, targetHeight, canvas])
}

func ctcDecode(_ output: [Float], timesteps: Int, classes: Int, characters: [String]) -> (text: String, score: Float) {
    var text = "", previousIndex = -1, sum: Float = 0, count = 0
    for t in 0..<timesteps {
        let offset = t * classes
        var bestIndex = 0, bestScore = output[offset]
        for i in 1..<classes where output[offset + i] > bestScore {
            bestScore = output[offset + i]; bestIndex = i
        }
        defer { previousIndex = bestIndex }
        guard bestIndex != 0, bestIndex != previousIndex, bestIndex < characters.count else { continue }
        text += characters[bestIndex]
        sum += bestScore; count += 1
    }
    return (text, count == 0 ? 0 : sum / Float(count))
}

func loadCharacters(from url: URL) -> [String] {
    let content = try! String(contentsOf: url, encoding: .utf8)
        .replacingOccurrences(of: "\r\n", with: "\n")
    return content.components(separatedBy: "\n").filter { !$0.isEmpty }
}