feat: support decoding images from memory addresses (issue #738)

python/rapidocr/main.py

@@ -17,7 +17,7 @@ from .ch_ppocr_det import TextDetector, TextDetOutput
 from .ch_ppocr_rec import TextRecInput, TextRecognizer, TextRecOutput
 from .cli import check_install, check_required_files, generate_cfg
 from .utils.download_models import download_models
-from .utils.load_image import LoadImage
+from .utils.load_image import InputType, LoadImage
 from .utils.log import logger
 from .utils.output import RapidOCROutput
 from .utils.parse_parameters import ParseParams
@@ -91,7 +91,7 @@ class RapidOCR:
 
     def __call__(
         self,
-        img_content: Union[str, np.ndarray, bytes, Path],
+        img_content: InputType,
         use_det: Optional[bool] = None,
         use_cls: Optional[bool] = None,
         use_rec: Optional[bool] = None,

python/rapidocr/utils/load_image.py

@@ -1,6 +1,7 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
+from ctypes import string_at
 from io import BytesIO
 from pathlib import Path
 from typing import Any, Union
@@ -13,7 +14,15 @@ from PIL import Image, ImageOps, UnidentifiedImageError
 from .utils import is_url
 
 root_dir = Path(__file__).resolve().parent
-InputType = Union[str, np.ndarray, bytes, Path, Image.Image]
+
+
+class MemoryImage:
+    def __init__(self, address: int, length: int):
+        self.address = address
+        self.length = length
+
+
+InputType = Union[str, np.ndarray, bytes, Path, Image.Image, MemoryImage]
 
 
 class LoadImage:
@@ -57,6 +66,12 @@ class LoadImage:
         if isinstance(img, Image.Image):
             return self.img_to_ndarray(img)
 
+        if isinstance(img, MemoryImage):
+            try:
+                return self.load_memory_image(img)
+            except Exception as e:
+                raise LoadImageError("MemoryImage loading failed.") from e
+
         raise LoadImageError(f"{type(img)} is not supported!")
 
     @staticmethod
@@ -159,6 +174,19 @@ class LoadImage:
 
         return cv2.cvtColor(blended, cv2.COLOR_RGB2BGR)
 
+    def load_memory_image(self, img: MemoryImage) -> np.ndarray:
+        if img.address <= 0 or img.length <= 0:
+            raise LoadImageError("invalid memory address or length")
+
+        encoded = string_at(img.address, img.length)
+        arr = np.frombuffer(encoded, dtype=np.uint8)
+        result = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+
+        if result is None:
+            raise LoadImageError("cannot decode image from memory")
+
+        return result
+
 
 class LoadImageError(Exception):
     pass

python/tests/test_input.py

@@ -1,6 +1,7 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
+import ctypes
 import sys
 from pathlib import Path
 
@@ -21,6 +22,7 @@ from rapidocr import (
     RapidOCR,
 )
 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"
@@ -120,6 +122,51 @@ def test_input_bytes(engine):
     assert result.txts[0] == "正品促销"
 
 
+def test_input_memory_image(engine):
+    """An encoded image can be passed by address and byte length."""
+    encoded = img_path.read_bytes()
+    buffer = ctypes.create_string_buffer(encoded)
+    memory_img = MemoryImage(ctypes.addressof(buffer), len(encoded))
+
+    result = engine(memory_img)
+
+    assert len(result) == 18
+    assert result.txts[0] == "正品促销"
+
+
+def test_input_memory_image_matches_bytes(engine):
+    encoded = img_path.read_bytes()
+    buffer = ctypes.create_string_buffer(encoded)
+    memory_img = MemoryImage(ctypes.addressof(buffer), len(encoded))
+
+    expected = engine(encoded)
+    actual = engine(memory_img)
+
+    assert actual.txts == expected.txts
+    assert actual.scores == expected.scores
+    np.testing.assert_allclose(actual.boxes, expected.boxes)
+
+
+@pytest.mark.parametrize(
+    "address,length",
+    [(0, 1), (-1, 1), (1, 0), (1, -1)],
+)
+def test_input_memory_image_invalid_address_or_length(address, length):
+    loader = LoadImage()
+
+    with pytest.raises(LoadImageError, match="MemoryImage loading failed"):
+        loader(MemoryImage(address, length))
+
+
+def test_input_memory_image_invalid_encoded_data():
+    encoded = b"not an image"
+    buffer = ctypes.create_string_buffer(encoded)
+    loader = LoadImage()
+
+    with pytest.raises(LoadImageError, match="MemoryImage loading failed"):
+        loader(MemoryImage(ctypes.addressof(buffer), len(encoded)))
+
+
 def test_input_path(engine):
     result = engine(img_path)
     assert len(result) == 18