fix: define rgb_bytes before torch.frombuffer in image decode

This commit is contained in:
Steve Tietze
2025-12-17 22:09:11 +07:00
parent e13a1cfb92
commit 89e5e91f40
+7 -1
View File
@@ -268,8 +268,14 @@ 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
tensor = tensor.view(
rgb_image.height, rgb_image.width, 3
).float() / 255.0
return tensor.unsqueeze(0)
except Exception as exc:
raise RuntimeError("Failed to decode result image.") from exc