fix: define rgb_bytes before torch.frombuffer in image decode
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user