from typing import IO, ClassVar
from PIL import Image
from PIL.ImageFile import ImageFile, PyDecoder, PyEncoder, _safe_read, _Tile
from PIL.ImageFile import _save as image_save
from ...viam_rgba import RGBA_FORMAT_LABEL, RGBA_HEADER_LENGTH, RGBA_MAGIC_NUMBER
def _accept(prefix: bytes):
return prefix[:4] == RGBA_MAGIC_NUMBER
[docs]
class RGBAEncoder(PyEncoder):
ENCODER_NAME = RGBA_FORMAT_LABEL
_pushes_fd = True
[docs]
def encode(self, bufsize: int) -> tuple[int, int, bytes]:
data_arr = bytearray()
assert self.im is not None
width, height = self.im.size
for y in range(height):
for x in range(width):
data_arr.extend(self.im.getpixel((x, y)))
data = bytes(data_arr)
return len(data), 0, data
def _save_rgba(img: Image.Image, fp: IO[bytes], filename: str | bytes):
width, height = img.size
fp.write(RGBA_MAGIC_NUMBER)
fp.write(width.to_bytes(4, byteorder="big"))
fp.write(height.to_bytes(4, byteorder="big"))
tile = _Tile(codec_name=RGBAEncoder.ENCODER_NAME, extents=(0, 0, width, height), offset=0, args=("RGBA", 0, 1))
image_save(img, fp, [tile])
[docs]
class RGBAImage(ImageFile):
format: ClassVar[str] = RGBA_FORMAT_LABEL # pyright: ignore [reportIncompatibleVariableOverride]
format_description = "Viam's Raw RGBA Format"
def _open(self):
assert self.fp is not None
header = self.fp.read(RGBA_HEADER_LENGTH)
width = int.from_bytes(header[4:8], "big")
height = int.from_bytes(header[8:12], "big")
self._size = width, height
if hasattr(self, "_mode"):
self._mode = "RGBA"
else:
self.mode = "RGBA" # type: ignore -- (njooma) newer versions of PIL hide this behind _mode, which is why we check
# data descriptor
tile = _Tile(codec_name=RGBAEncoder.ENCODER_NAME, extents=(0, 0, width, height), offset=RGBA_HEADER_LENGTH, args=(self.mode, 0, 1))
self.tile = [tile]
[docs]
class RGBADecoder(PyDecoder):
_pulls_fd = True
[docs]
def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]:
assert self.im is not None
assert self.fd is not None
width, height = self.im.size
self.set_as_raw(_safe_read(self.fd, width * height * 4))
return -1, 0
# Register Viam's custom RGBA format so that it can be read using Image.open
Image.register_open(RGBAImage.format, RGBAImage, _accept)
Image.register_save(RGBAImage.format, _save_rgba)
Image.register_mime(RGBAImage.format, "image/vnd.viam.rgba")
Image.register_extension(RGBAImage.format, ".viam.rgba")
Image.register_encoder(RGBAEncoder.ENCODER_NAME, RGBAEncoder)
Image.register_decoder(RGBA_FORMAT_LABEL, RGBADecoder)