feat: support loading raw memory image buffers with format and stride metadata

python/rapidocr/utils/load_image.py

@@ -4,7 +4,7 @@
 from ctypes import string_at
 from io import BytesIO
 from pathlib import Path
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import cv2
 import numpy as np
@@ -22,7 +22,31 @@ class MemoryImage:
         self.length = length
 
 
-InputType = Union[str, np.ndarray, bytes, Path, Image.Image, MemoryImage]
+class RawMemoryImage:
+    """Metadata needed to read an uncompressed pixel buffer by address."""
+
+    def __init__(
+        self,
+        address: int,
+        length: int,
+        width: int,
+        height: int,
+        pixel_format: str = "BGRA",
+        stride: Optional[int] = None,
+        bottom_up: bool = False,
+    ):
+        self.address = address
+        self.length = length
+        self.width = width
+        self.height = height
+        self.pixel_format = pixel_format
+        self.stride = stride
+        self.bottom_up = bottom_up
+
+
+InputType = Union[
+    str, np.ndarray, bytes, Path, Image.Image, MemoryImage, RawMemoryImage
+]
 
 
 class LoadImage:
@@ -72,6 +96,12 @@ class LoadImage:
             except Exception as e:
                 raise LoadImageError("MemoryImage loading failed.") from e
 
+        if isinstance(img, RawMemoryImage):
+            try:
+                return self.load_raw_memory_image(img)
+            except Exception as e:
+                raise LoadImageError(f"RawMemoryImage loading failed: {e}") from e
+
         raise LoadImageError(f"{type(img)} is not supported!")
 
     @staticmethod
@@ -187,6 +217,69 @@ class LoadImage:
 
         return result
 
+    @staticmethod
+    def load_raw_memory_image(img: RawMemoryImage) -> np.ndarray:
+        if img.address <= 0 or img.length <= 0:
+            raise LoadImageError("invalid memory address or length")
+        if img.width <= 0 or img.height <= 0:
+            raise LoadImageError("width and height must be positive")
+
+        pixel_format = str(img.pixel_format).upper()
+        channels_by_format = {
+            "GRAY": 1,
+            "GRAY8": 1,
+            "BGR": 3,
+            "BGR24": 3,
+            "RGB": 3,
+            "RGB24": 3,
+            "BGRA": 4,
+            "BGRA32": 4,
+            "BGRX": 4,
+            "RGBA": 4,
+            "RGBA32": 4,
+            "RGBX": 4,
+        }
+        channels = channels_by_format.get(pixel_format)
+        if channels is None:
+            supported = ", ".join(channels_by_format)
+            raise LoadImageError(
+                f"unsupported pixel format {img.pixel_format!r}; supported: {supported}"
+            )
+
+        row_bytes = img.width * channels
+        stride = row_bytes if img.stride is None else img.stride
+        if stride < row_bytes:
+            raise LoadImageError(
+                f"stride ({stride}) is smaller than row size ({row_bytes})"
+            )
+
+        required_length = stride * img.height
+        if img.length < required_length:
+            raise LoadImageError(
+                f"buffer is too small: need {required_length} bytes, got {img.length}"
+            )
+
+        raw = string_at(img.address, required_length)
+        rows = np.frombuffer(raw, dtype=np.uint8).reshape(img.height, stride)
+        pixels = rows[:, :row_bytes]
+        if channels == 1:
+            pixels = pixels.reshape(img.height, img.width)
+        else:
+            pixels = pixels.reshape(img.height, img.width, channels)
+
+        if img.bottom_up:
+            pixels = pixels[::-1]
+
+        if pixel_format in {"GRAY", "GRAY8"}:
+            return cv2.cvtColor(pixels, cv2.COLOR_GRAY2BGR)
+        if pixel_format in {"RGB", "RGB24"}:
+            return cv2.cvtColor(pixels, cv2.COLOR_RGB2BGR)
+        if pixel_format in {"BGRA", "BGRA32", "BGRX"}:
+            return cv2.cvtColor(pixels, cv2.COLOR_BGRA2BGR)
+        if pixel_format in {"RGBA", "RGBA32", "RGBX"}:
+            return cv2.cvtColor(pixels, cv2.COLOR_RGBA2BGR)
+        return np.ascontiguousarray(pixels)
+
 
 class LoadImageError(Exception):
     pass

python/tests/test_input.py

@@ -21,8 +21,8 @@ from rapidocr import (
     OCRVersion,
     RapidOCR,
 )
+from rapidocr.utils.load_image import LoadImage, MemoryImage, RawMemoryImage
 from rapidocr.utils.parse_parameters import ParseParams
-from rapidocr.utils.load_image import LoadImage, MemoryImage
 
 test_dir = root_dir / "tests" / "test_files"
 img_path = test_dir / "ch_en_num.jpg"
@@ -167,6 +167,101 @@ def test_input_memory_image_invalid_encoded_data():
         loader(MemoryImage(ctypes.addressof(buffer), len(encoded)))
 
 
+def make_raw_memory_image(img, pixel_format, stride=None, bottom_up=False):
+    height, width = img.shape[:2]
+    row_bytes = img.shape[1] * (1 if img.ndim == 2 else img.shape[2])
+    stride = row_bytes if stride is None else stride
+    rows = img[::-1] if bottom_up else img
+    raw = bytearray(stride * height)
+    for row_index, row in enumerate(rows):
+        row_data = row.tobytes()
+        raw[row_index * stride : row_index * stride + row_bytes] = row_data
+    buffer = ctypes.create_string_buffer(bytes(raw))
+    memory_img = RawMemoryImage(
+        ctypes.addressof(buffer),
+        len(raw),
+        width,
+        height,
+        pixel_format=pixel_format,
+        stride=stride,
+        bottom_up=bottom_up,
+    )
+    return memory_img, buffer
+
+
+def test_input_raw_memory_image(engine):
+    img = cv2.imread(str(img_path))
+    memory_img, buffer = make_raw_memory_image(img, "BGR")
+
+    result = engine(memory_img)
+
+    assert buffer is not None  # Keep the backing memory alive through inference.
+    assert len(result) == 18
+    assert result.txts[0] == "正品促销"
+
+
+@pytest.mark.parametrize(
+    "pixel_format,source",
+    [
+        ("BGR", np.array([[[1, 2, 3], [4, 5, 6]]], dtype=np.uint8)),
+        ("RGB", np.array([[[3, 2, 1], [6, 5, 4]]], dtype=np.uint8)),
+        ("BGRA", np.array([[[1, 2, 3, 0], [4, 5, 6, 255]]], dtype=np.uint8)),
+        ("BGRX", np.array([[[1, 2, 3, 8], [4, 5, 6, 9]]], dtype=np.uint8)),
+        ("RGBA", np.array([[[3, 2, 1, 0], [6, 5, 4, 255]]], dtype=np.uint8)),
+        ("RGBX", np.array([[[3, 2, 1, 8], [6, 5, 4, 9]]], dtype=np.uint8)),
+    ],
+)
+def test_raw_memory_image_converts_to_bgr(pixel_format, source):
+    memory_img, buffer = make_raw_memory_image(source, pixel_format)
+
+    result = LoadImage()(memory_img)
+
+    assert buffer is not None
+    np.testing.assert_array_equal(
+        result, np.array([[[1, 2, 3], [4, 5, 6]]], dtype=np.uint8)
+    )
+
+
+def test_raw_memory_image_supports_gray_stride_and_bottom_up():
+    gray = np.array([[10, 20], [30, 40]], dtype=np.uint8)
+    memory_img, buffer = make_raw_memory_image(
+        gray, "GRAY8", stride=4, bottom_up=True
+    )
+
+    result = LoadImage()(memory_img)
+
+    assert buffer is not None
+    expected = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+    np.testing.assert_array_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "kwargs,error",
+    [
+        ({"address": 0}, "invalid memory address"),
+        ({"width": 0}, "width and height must be positive"),
+        ({"pixel_format": "YUV"}, "unsupported pixel format"),
+        ({"stride": 5}, "stride .* is smaller than row size"),
+        ({"length": 8}, "buffer is too small"),
+    ],
+)
+def test_raw_memory_image_rejects_invalid_metadata(kwargs, error):
+    raw = bytes(12)
+    buffer = ctypes.create_string_buffer(raw)
+    params = {
+        "address": ctypes.addressof(buffer),
+        "length": len(raw),
+        "width": 2,
+        "height": 2,
+        "pixel_format": "BGR",
+        "stride": 6,
+    }
+    params.update(kwargs)
+
+    with pytest.raises(LoadImageError, match=error):
+        LoadImage()(RawMemoryImage(**params))
+
+
 def test_input_path(engine):
     result = engine(img_path)
     assert len(result) == 18