Image to ndarray; Image to bytes buffer; Image to buffer pickle; Pickle to Image. Pickle to ndarray

compressed_image.py

from __future__ import annotations

import io
from functools import lru_cache
from PIL import Image
import numpy as np
import pickle
import os


class CompressedImage:
    format = "PNG"

    def __init__(self, image_array: np.ndarray, format=None):
        self._image_array = image_array
        if format:
            self.format = format

    @property
    def image_array(self) -> np.ndarray:
        return self._image_array

    @property
    @lru_cache()
    def image(self) -> Image:
        return Image.fromarray(self.image_array)

    @property
    @lru_cache()
    def image_buffer(self) -> io.BytesIO:
        _io = io.BytesIO()
        self.image.save(_io, format=self.format)
        return _io

    @property
    @lru_cache()
    def image_buffer_pickle(self, *args, **kwargs) -> bytes:
        return pickle.dumps(self.image_buffer, *args, **kwargs)

    @staticmethod
    def from_pickle_data(pickle_data: bytes) -> CompressedImage:
        _io = pickle.loads(pickle_data)
        image: Image = Image.open(_io)
        image_array = np.asarray(image)
        return CompressedImage(image_array, format=image.format)


def test_compressed_image():
    img = np.empty((256, 256, 3), dtype='uint8')

    ci = CompressedImage(img)
    assert ci.image_array.sum() == img.sum()

    new_ci = CompressedImage.from_pickle_data(ci.image_buffer_pickle)
    assert new_ci.image_array.sum() == ci.image_array.sum()

    file_name = 'temp.pick'
    with open(file_name, 'wb') as f:
        f.write(new_ci.image_buffer_pickle)

    with open(file_name, 'rb') as f:
        new_new_ci = CompressedImage.from_pickle_data(f.read())
    assert new_new_ci.image_array.sum() == new_ci.image_array.sum()

    os.remove(file_name)
    assert not os.path.isfile(file_name)