From 58b2fa9e36463fdaaf741e9b1c75a08880d02401 Mon Sep 17 00:00:00 2001 From: Steve Tietze Date: Thu, 18 Dec 2025 12:33:04 +0700 Subject: [PATCH] numpy for images --- kie_api/images.py | 13 ++++++++----- kie_api/jobs.py | 0 2 files changed, 8 insertions(+), 5 deletions(-) create mode 100644 kie_api/jobs.py diff --git a/kie_api/images.py b/kie_api/images.py index 527f9b2..eb1859a 100644 --- a/kie_api/images.py +++ b/kie_api/images.py @@ -1,6 +1,7 @@ from io import BytesIO import torch +import numpy as np from PIL import Image @@ -15,10 +16,12 @@ def _image_bytes_to_tensor(image_bytes: bytes) -> torch.Tensor: try: with Image.open(BytesIO(image_bytes)) as img: rgb_image = img.convert("RGB") - rgb_bytes = rgb_image.tobytes() - - tensor = torch.frombuffer(rgb_bytes, dtype=torch.uint8).clone() - tensor = tensor.view(rgb_image.height, rgb_image.width, 3).float() / 255.0 + + # FIX: Use numpy to create the array. This creates a writable buffer + # automatically, resolving the PyTorch warning. + tensor = torch.from_numpy(np.array(rgb_image)) + + tensor = tensor.float() / 255.0 return tensor.unsqueeze(0) except Exception as exc: - raise RuntimeError("Failed to decode result image.") from exc + raise RuntimeError("Failed to decode result image.") from exc \ No newline at end of file diff --git a/kie_api/jobs.py b/kie_api/jobs.py new file mode 100644 index 0000000..e69de29