diff --git a/python/rapidocr/utils/load_image.py b/python/rapidocr/utils/load_image.py index 08844bd..802c93a 100644 --- a/python/rapidocr/utils/load_image.py +++ b/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 diff --git a/python/tests/test_input.py b/python/tests/test_input.py index de7bf7a..e5948ba 100644 --- a/python/tests/test_input.py +++ b/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