openimageio helpers for jupyter notebook

jupyter_oiio.py

import os
import tempfile

from contextlib import contextmanager
from dataclasses import dataclass

import OpenImageIO as oiio
import einops

DTYPE_TO_OIIO = {
    'byte': oiio.UINT8,
    'float16': oiio.HALF,
    'float32': oiio.FLOAT,
    'float64': oiio.DOUBLE,
    'uint16': oiio.UINT16
}


@dataclass
class TempFileContainer:
    name: str
    content: bytes = None


@contextmanager
def get_from_tempfile(**tempfile_kwargs):
    # Check for RAM-based directories (common on Linux)
    ram_disk_dirs = ['/dev/shm', '/run/shm']
    temp_dir = next((d for d in ram_disk_dirs if os.path.exists(d)), None)

    with tempfile.NamedTemporaryFile(dir=temp_dir, delete_on_close=False, **tempfile_kwargs) as temp_file:
        temp_file.close()
        result = TempFileContainer(name=temp_file.name)
        yield result

        with open(temp_file.name, 'rb') as f:
            result.content = f.read()


def pixels_hwc_to_webp(image_array, format="webp"):
    spec = oiio.ImageSpec(
        image_array.shape[1],  # width
        image_array.shape[0],  # height
        image_array.shape[2],  # WebP only supports RGB or RGBA
        DTYPE_TO_OIIO[str(image_array.dtype)]  # WebP will convert to uint8
    )
    spec.attribute("compression", "lossless")

    with get_from_tempfile(suffix="." + format) as temp_file:
        output = oiio.ImageOutput.create(temp_file.name)
        output.open(temp_file.name, spec)
        try:
            output.write_image(image_array)  # must be in H W C order here
            if output.has_error:
                raise RuntimeError(output.geterror())
        finally:
            output.close()

    return temp_file.content
    

def array_f_to_png(image_array):
    """Write an array as a single-channel PNG."""
    if len(image_array.shape) == 2:
        image_array = einops.rearrange(image_array, "h w -> h w ()")

    return pixels_hwc_to_webp(image_array, format="png")


def oiio_imagebuf_mimebundle(ib):
    if ib.spec().nchannels < 3:
        format = "png"
    else:
        format = "webp"
    with get_from_tempfile(suffix="." + format) as temp_file:
        ib.write(temp_file.name)
        if ib.has_error:
            raise RuntimeError(ib.geterror())

    return {'image/' + format: temp_file.content}


get_ipython().display_formatter.mimebundle_formatter.for_type(oiio.ImageBuf, oiio_imagebuf_mimebundle)