e904b667c6
Build/Publish Develop Docs / deploy (push) Failing after 1s
PaddleOCR Code Style Check / check-code-style (push) Failing after 1s
PaddleOCR PR Tests GPU / detect-changes (push) Failing after 1s
PaddleOCR PR Tests / detect-changes (push) Failing after 1s
PaddleOCR PR Tests GPU / test-pr-gpu (push) Has been cancelled
PaddleOCR PR Tests / test-pr (push) Has been cancelled
PaddleOCR PR Tests GPU / test-pr-gpu-impl (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.13) (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.8) (push) Has been cancelled
PaddleOCR PR Tests / test-pr-python (3.9) (push) Has been cancelled
381 lines
13 KiB
Swift
381 lines
13 KiB
Swift
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
||
//
|
||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||
// you may not use this file except in compliance with the License.
|
||
// You may obtain a copy of the License at
|
||
//
|
||
// http://www.apache.org/licenses/LICENSE-2.0
|
||
//
|
||
// Unless required by applicable law or agreed to in writing, software
|
||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
// See the License for the specific language governing permissions and
|
||
// limitations under the License.
|
||
|
||
import CoreGraphics
|
||
import Foundation
|
||
|
||
// MARK: - Errors
|
||
|
||
enum RecPreprocessorError: LocalizedError {
|
||
case invalidImage
|
||
case configMissing(String)
|
||
case pixelExtractionFailed
|
||
|
||
var errorDescription: String? {
|
||
switch self {
|
||
case .invalidImage:
|
||
return "Invalid input image: width or height is zero"
|
||
case .configMissing(let detail):
|
||
return "Required recognition preprocessing config missing: \(detail)"
|
||
case .pixelExtractionFailed:
|
||
return "Failed to extract pixel data from image"
|
||
}
|
||
}
|
||
}
|
||
|
||
// MARK: - Preprocessing Result
|
||
|
||
/// The output of recognition preprocessing: a float32 tensor and metadata.
|
||
struct RecPreprocessResult {
|
||
/// Flat CHW array of shape [1, imgC, imgH, targetW].
|
||
let tensorData: [Float]
|
||
/// Tensor dimensions: [1, imgC, imgH, targetW].
|
||
let tensorShape: [Int]
|
||
/// Original image dimensions before any resizing.
|
||
let originalSize: (width: Int, height: Int)
|
||
/// Actual content width before zero-padding (the resized image width).
|
||
let resizedWidth: Int
|
||
}
|
||
|
||
/// Batched recognition input: N × C × H × W with a common width after padding.
|
||
struct RecBatchPreprocessResult {
|
||
let tensorData: [Float]
|
||
let tensorShape: [Int]
|
||
}
|
||
|
||
// MARK: - RecPreprocessor
|
||
|
||
/// Recognition resize + normalize + pad: dynamic width, fixed height, scale to [-1, 1],
|
||
/// CHW layout, right-pad to target width.
|
||
///
|
||
/// Recognition resize, per-channel normalization, HWC→CHW, and width padding.
|
||
/// `DecodeImage.img_mode` is applied via `InferenceConfig.decodeImageChannelOrder`. Resize uses OpenCV `INTER_LINEAR`.
|
||
///
|
||
/// Set environment variable `PADDLEOCR_REC_ALIGN_WIDTH_TO_8` to `1`, `true`, or `yes` to round the
|
||
/// canvas width up to a multiple of 8 (off by default).
|
||
struct RecPreprocessor {
|
||
/// Channel count from `RecResizeImg.image_shape` (3).
|
||
private let imgC: Int
|
||
/// Target height from config RecResizeImg.image_shape (e.g., 48).
|
||
private let imgH: Int
|
||
/// Default width from config RecResizeImg.image_shape (e.g., 320).
|
||
private let imgW: Int
|
||
/// Absolute maximum width cap to prevent excessive memory usage.
|
||
private let maxImgW: Int = 3200
|
||
/// Channel order from `DecodeImage.img_mode` in the model config file.
|
||
private let channelOrder: InferenceImageChannelOrder
|
||
|
||
/// Creates a RecPreprocessor by extracting image_shape from the RecResizeImg transform op.
|
||
///
|
||
/// - Parameter config: A parsed InferenceConfig from the recognition model's config file.
|
||
/// - Throws: `RecPreprocessorError.configMissing` if RecResizeImg with image_shape is absent.
|
||
init(config: InferenceConfig) throws {
|
||
var foundImageShape: [Int]?
|
||
|
||
for op in config.preProcess.transformOps {
|
||
switch op {
|
||
case .decodeImage, .detResizeForTest, .normalizeImage, .toCHWImage, .unknown:
|
||
break
|
||
case .recResizeImg(let imageShape):
|
||
foundImageShape = imageShape
|
||
}
|
||
}
|
||
|
||
guard let imageShape = foundImageShape, imageShape.count >= 3 else {
|
||
throw RecPreprocessorError.configMissing("RecResizeImg with image_shape")
|
||
}
|
||
|
||
self.imgC = imageShape[0]
|
||
self.imgH = imageShape[1]
|
||
self.imgW = imageShape[2]
|
||
self.channelOrder = config.decodeImageChannelOrder
|
||
}
|
||
|
||
/// Runs recognition preprocessing on a `CGImage`.
|
||
///
|
||
/// Steps: extract pixels → compute dimensions → resize → normalize → HWC-to-CHW → pad
|
||
///
|
||
/// - Parameter image: The input image (typically a cropped text region).
|
||
/// - Returns: A `RecPreprocessResult` with the float32 CHW tensor and metadata.
|
||
func preprocess(_ image: CGImage) throws -> RecPreprocessResult {
|
||
let originalW = image.width
|
||
let originalH = image.height
|
||
|
||
guard originalW > 0, originalH > 0 else {
|
||
throw RecPreprocessorError.invalidImage
|
||
}
|
||
|
||
let pixelBytes = try extractHWCPixels(from: image, width: originalW, height: originalH, order: channelOrder)
|
||
let (paddedData, canvasW, resizedW) = try lineResizeNormalizeCHW(
|
||
pixelBytes: pixelBytes,
|
||
originalW: originalW,
|
||
originalH: originalH,
|
||
alignCanvasTo8: Self.isRecognitionCanvasWidthAlignedTo8Enabled()
|
||
)
|
||
|
||
return RecPreprocessResult(
|
||
tensorData: paddedData,
|
||
tensorShape: [1, imgC, imgH, canvasW],
|
||
originalSize: (width: originalW, height: originalH),
|
||
resizedWidth: resizedW
|
||
)
|
||
}
|
||
|
||
/// Builds a batched tensor: each line is resized and padded to its canvas width, then all rows
|
||
/// are padded on the right to the maximum canvas width in the batch (shape `[N, C, H, W]`).
|
||
func preprocessBatch(_ images: [CGImage]) throws -> RecBatchPreprocessResult {
|
||
guard !images.isEmpty else {
|
||
throw RecPreprocessorError.invalidImage
|
||
}
|
||
|
||
let align8 = Self.isRecognitionCanvasWidthAlignedTo8Enabled()
|
||
var rowCHW: [[Float]] = []
|
||
var canvasWidths: [Int] = []
|
||
|
||
for image in images {
|
||
let originalW = image.width
|
||
let originalH = image.height
|
||
guard originalW > 0, originalH > 0 else {
|
||
throw RecPreprocessorError.invalidImage
|
||
}
|
||
let pixelBytes = try extractHWCPixels(from: image, width: originalW, height: originalH, order: channelOrder)
|
||
let (chw, canvasW, _) = try lineResizeNormalizeCHW(
|
||
pixelBytes: pixelBytes,
|
||
originalW: originalW,
|
||
originalH: originalH,
|
||
alignCanvasTo8: align8
|
||
)
|
||
rowCHW.append(chw)
|
||
canvasWidths.append(canvasW)
|
||
}
|
||
|
||
guard let maxCanvasW = canvasWidths.max() else {
|
||
throw RecPreprocessorError.invalidImage
|
||
}
|
||
|
||
let n = rowCHW.count
|
||
let channelSize = imgH * maxCanvasW
|
||
var flat = [Float](repeating: 0, count: n * imgC * channelSize)
|
||
|
||
for (b, row) in rowCHW.enumerated() {
|
||
let w = canvasWidths[b]
|
||
let paddedRow: [Float]
|
||
if w == maxCanvasW {
|
||
paddedRow = row
|
||
} else {
|
||
let contentW = w
|
||
paddedRow = padToTargetWidth(
|
||
row,
|
||
contentWidth: contentW,
|
||
targetWidth: maxCanvasW,
|
||
height: imgH,
|
||
channels: imgC
|
||
)
|
||
}
|
||
let batchOffset = b * imgC * channelSize
|
||
for i in 0..<paddedRow.count {
|
||
flat[batchOffset + i] = paddedRow[i]
|
||
}
|
||
}
|
||
|
||
return RecBatchPreprocessResult(
|
||
tensorData: flat,
|
||
tensorShape: [n, imgC, imgH, maxCanvasW]
|
||
)
|
||
}
|
||
|
||
/// When `PADDLEOCR_REC_ALIGN_WIDTH_TO_8` is set to a truthy value, the canvas width is rounded up to a multiple of 8.
|
||
private static func isRecognitionCanvasWidthAlignedTo8Enabled() -> Bool {
|
||
let v = ProcessInfo.processInfo.environment["PADDLEOCR_REC_ALIGN_WIDTH_TO_8"]?.lowercased()
|
||
return v == "1" || v == "true" || v == "yes"
|
||
}
|
||
|
||
/// One text line: HWC uint8 → resize → normalize → CHW, padded to the computed canvas width.
|
||
private func lineResizeNormalizeCHW(
|
||
pixelBytes: [UInt8],
|
||
originalW: Int,
|
||
originalH: Int,
|
||
alignCanvasTo8: Bool
|
||
) throws -> (chw: [Float], canvasW: Int, resizedW: Int) {
|
||
let w = originalW
|
||
let h = originalH
|
||
let whRatio = Float(w) / Float(h)
|
||
let defaultRatio = Float(imgW) / Float(imgH)
|
||
let maxWhRatio = max(defaultRatio, whRatio)
|
||
var canvasW = Int(Float(imgH) * maxWhRatio)
|
||
let resizedW: Int
|
||
|
||
if canvasW > maxImgW {
|
||
canvasW = maxImgW
|
||
resizedW = maxImgW
|
||
} else {
|
||
let ratio = Float(w) / Float(h)
|
||
if Int(ceil(Float(imgH) * ratio)) > canvasW {
|
||
resizedW = canvasW
|
||
} else {
|
||
resizedW = Int(ceil(Float(imgH) * ratio))
|
||
}
|
||
}
|
||
|
||
if alignCanvasTo8 {
|
||
canvasW = min(max(roundUpToMultipleOf(canvasW, 8), resizedW), maxImgW)
|
||
}
|
||
|
||
let resizedPixels = resizeWithOpenCV(pixelBytes, srcW: w, srcH: h, dstW: resizedW, dstH: imgH)
|
||
let normalizedHWC = normalizePixels(resizedPixels, width: resizedW, height: imgH)
|
||
let chwData = hwcToCHW(normalizedHWC, width: resizedW, height: imgH)
|
||
let paddedData = padToTargetWidth(
|
||
chwData,
|
||
contentWidth: resizedW,
|
||
targetWidth: canvasW,
|
||
height: imgH,
|
||
channels: imgC
|
||
)
|
||
return (paddedData, canvasW, resizedW)
|
||
}
|
||
|
||
// MARK: - Step 0: Pixel extraction (BGR / RGB)
|
||
|
||
private func extractHWCPixels(
|
||
from image: CGImage,
|
||
width: Int,
|
||
height: Int,
|
||
order: InferenceImageChannelOrder
|
||
) throws -> [UInt8] {
|
||
let bytesPerPixel = 4
|
||
let bytesPerRow = width * bytesPerPixel
|
||
var rgbaData = [UInt8](repeating: 0, count: height * bytesPerRow)
|
||
|
||
guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB),
|
||
let context = CGContext(
|
||
data: &rgbaData,
|
||
width: width,
|
||
height: height,
|
||
bitsPerComponent: 8,
|
||
bytesPerRow: bytesPerRow,
|
||
space: colorSpace,
|
||
bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue
|
||
) else {
|
||
throw RecPreprocessorError.pixelExtractionFailed
|
||
}
|
||
|
||
context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
|
||
|
||
var out = [UInt8](repeating: 0, count: height * width * 3)
|
||
for i in 0..<(height * width) {
|
||
let r = rgbaData[i * 4]
|
||
let g = rgbaData[i * 4 + 1]
|
||
let b = rgbaData[i * 4 + 2]
|
||
switch order {
|
||
case .bgr:
|
||
out[i * 3] = b
|
||
out[i * 3 + 1] = g
|
||
out[i * 3 + 2] = r
|
||
case .rgb:
|
||
out[i * 3] = r
|
||
out[i * 3 + 1] = g
|
||
out[i * 3 + 2] = b
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
// MARK: - Resize (OpenCV)
|
||
|
||
private func resizeWithOpenCV(
|
||
_ pixels: [UInt8], srcW: Int, srcH: Int, dstW: Int, dstH: Int
|
||
) -> [UInt8] {
|
||
if srcW == dstW && srcH == dstH {
|
||
return pixels
|
||
}
|
||
let out = PDBOpenCVImageBridge.resizeRGBU8(
|
||
Data(pixels), srcWidth: srcW, srcHeight: srcH, dstWidth: dstW, dstHeight: dstH
|
||
)
|
||
let expected = dstW * dstH * 3
|
||
guard out.count == expected else {
|
||
return [UInt8](repeating: 0, count: expected)
|
||
}
|
||
return [UInt8](out)
|
||
}
|
||
|
||
// MARK: - Normalize (recognition)
|
||
|
||
/// Maps pixel values from [0, 255] to [-1, 1] via `pixel / 127.5 - 1.0`.
|
||
private func normalizePixels(_ pixels: [UInt8], width: Int, height: Int) -> [Float] {
|
||
let count = width * height * 3
|
||
var result = [Float](repeating: 0, count: count)
|
||
|
||
for i in 0..<count {
|
||
result[i] = Float(pixels[i]) / 127.5 - 1.0
|
||
}
|
||
|
||
return result
|
||
}
|
||
|
||
// MARK: - HWC → CHW
|
||
|
||
private func hwcToCHW(_ hwcData: [Float], width: Int, height: Int) -> [Float] {
|
||
let channelSize = height * width
|
||
var chwData = [Float](repeating: 0, count: 3 * channelSize)
|
||
|
||
for y in 0..<height {
|
||
for x in 0..<width {
|
||
let hwcIdx = (y * width + x) * 3
|
||
let pixelOffset = y * width + x
|
||
chwData[0 * channelSize + pixelOffset] = hwcData[hwcIdx]
|
||
chwData[1 * channelSize + pixelOffset] = hwcData[hwcIdx + 1]
|
||
chwData[2 * channelSize + pixelOffset] = hwcData[hwcIdx + 2]
|
||
}
|
||
}
|
||
|
||
return chwData
|
||
}
|
||
|
||
// MARK: - Zero-Padding
|
||
|
||
/// Right-pads a CHW tensor from contentWidth to targetWidth with zeros.
|
||
private func padToTargetWidth(
|
||
_ chwData: [Float],
|
||
contentWidth: Int,
|
||
targetWidth: Int,
|
||
height: Int,
|
||
channels: Int
|
||
) -> [Float] {
|
||
if contentWidth == targetWidth {
|
||
return chwData
|
||
}
|
||
|
||
let targetChannelSize = height * targetWidth
|
||
let sourceChannelSize = height * contentWidth
|
||
var padded = [Float](repeating: 0, count: channels * targetChannelSize)
|
||
|
||
for c in 0..<channels {
|
||
for y in 0..<height {
|
||
let srcOffset = c * sourceChannelSize + y * contentWidth
|
||
let dstOffset = c * targetChannelSize + y * targetWidth
|
||
for x in 0..<contentWidth {
|
||
padded[dstOffset + x] = chwData[srcOffset + x]
|
||
}
|
||
}
|
||
}
|
||
|
||
return padded
|
||
}
|
||
}
|
||
|
||
private func roundUpToMultipleOf(_ value: Int, _ multiple: Int) -> Int {
|
||
guard multiple > 0 else { return value }
|
||
return (value + multiple - 1) / multiple * multiple
|
||
}
|