Source code for viam.media.utils.pil.viam_rgba_plugin

from typing import ClassVar, Tuple

from PIL import Image
from PIL.ImageFile import ImageFile, PyDecoder, PyEncoder, _safe_read  # type: ignore -- (njooma) this exists, manually checked
from PIL.ImageFile import _save as image_save  # type: ignore -- (njooma) this exists, manually checked

from ...viam_rgba import RGBA_FORMAT_LABEL, RGBA_HEADER_LENGTH, RGBA_MAGIC_NUMBER


def _accept(prefix: str):
    return prefix[:4] == RGBA_MAGIC_NUMBER


[docs]class RGBAEncoder(PyEncoder): ENCODER_NAME = RGBA_FORMAT_LABEL _pushes_fd = True
[docs] def encode(self, bufsize): # pyright: ignore [reportIncompatibleMethodOverride] data_arr = bytearray() 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, fp, filename): 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")) image_save(img, fp, [(RGBAEncoder.ENCODER_NAME, (0, 0, width, height), 0, ("RGBA", 0, 1))])
[docs]class RGBAImage(ImageFile): format: ClassVar[str] = RGBA_FORMAT_LABEL # pyright: ignore [reportIncompatibleVariableOverride] format_description = "Viam's Raw RGBA Format" def _open(self): 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 self.tile = [(RGBAEncoder.ENCODER_NAME, (0, 0, width, height), RGBA_HEADER_LENGTH, (self.mode, 0, 1))]
[docs]class RGBADecoder(PyDecoder): _pulls_fd = True
[docs] def decode(self, buffer) -> Tuple[int, int]: 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)