blender-portable-repo/extensions/user_default/memsaver_personal/derivative_generator.py

# copyright (c) 2018- polygoniq xyz s.r.o.

import os
import sys
import logging

logger = logging.getLogger(f"polygoniq.{__name__}")


GENERATORS = []


try:
    # OpenImageIO is installed by default in Blender 3.5+
    import OpenImageIO as oiio

    def generate_derivative_OIIO(
        original_abs_path: str, derivative_path: str, side_size: int
    ) -> bool:
        _, original_ext = os.path.splitext(original_abs_path)
        _, derivative_ext = os.path.splitext(derivative_path)
        assert original_ext == derivative_ext

        # TODO: More extensions?
        # TODO: Update faq.md in docs when adding new extensions
        if original_ext.lower() not in {
            ".bmp",
            ".exr",
            ".jpg",
            ".jpeg",
            ".jpe",
            ".jif",
            ".jfif",
            ".png",
            ".tga",
            ".hdr",
            ".tiff",
            ".tif",
            ".webp",
        }:
            logger.warning(
                f"Can't generate derivative of {original_abs_path} because its extension "
                f"{original_ext} is not supported by OpenImageIO."
            )
            return False

        try:
            original_image = oiio.ImageInput.open(original_abs_path)
            if original_image is None:
                raise RuntimeError(
                    f"OIIO returned None when trying to load {original_abs_path}. "
                    f"OIIO error: {oiio.geterror()}"
                )
            original_spec = original_image.spec()
        except:
            logger.exception(
                f"Uncaught exception while loading {original_abs_path} image with OpenImageIO"
            )
            return False

        original_width = original_spec.width
        original_height = original_spec.height
        if side_size >= original_width and side_size >= original_height:
            logger.warning(
                f"Refused to generate derivative of {original_abs_path} of side size {side_size} "
                f"because the original size {original_width}x{original_height} is larger or equal!"
            )
            return False

        # Compute the new dimensions while maintaining aspect ratio
        if original_width >= original_height:
            new_width = side_size
            new_height = int(side_size * original_height / original_width)
        else:
            new_height = side_size
            new_width = int(side_size * original_width / original_height)

        try:
            derivative_spec = oiio.ImageSpec(
                new_width, new_height, original_spec.nchannels, original_spec.format
            )
            icc_profile = original_spec.get_string_attribute("ICCProfile")
            if icc_profile != "":
                derivative_spec.attribute("ICCProfile", oiio.TypeDesc.TypeString, icc_profile)
            original_buf = oiio.ImageBuf(original_abs_path)
            derivative_buf = oiio.ImageBuf(derivative_spec)
            oiio.ImageBufAlgo.resize(derivative_buf, original_buf, filtername="lanczos3")
            derivative_buf.write(derivative_path)
        except:
            logger.exception(
                f"Uncaught exception while generating derivative for {original_abs_path} "
                f"with OpenImageIO. OIIO error: {oiio.geterror()}"
            )
            return False

        logger.info(
            f"Generated derivative of size {new_width}x{new_height} from original "
            f"{original_abs_path} using OpenImageIO"
        )
        return True

    GENERATORS.append(generate_derivative_OIIO)
    logger.info("OpenImageIO successfully imported and will be used for memsaver.")
except ImportError:
    logger.info("OpenImageIO could not be imported, we can't use it for memsaver.")


if len(GENERATORS) == 0:  # We will only try to use PIL/Pillow if OpenImageIO is not present
    try:
        # Install modules which are not in Blender python by default
        # https://conference.blender.org/2022/presentations/1405/
        # or change to this: https://blender.stackexchange.com/questions/168448/bundling-python-library-with-addon
        try:
            import PIL.Image
        except ModuleNotFoundError as ex:
            import subprocess

            python_exe = sys.executable
            args = [python_exe, "-m", "ensurepip", "--upgrade", "--default-pip"]
            if subprocess.call(args=args) != 0:
                raise RuntimeError("Couldn't ensure pip in Blender's python!")
            args = [python_exe, "-m", "pip", "install", "--upgrade", "Pillow"]
            if subprocess.call(args=args) != 0:
                raise RuntimeError("Couldn't install Pillow module in Blender's python!")
            import PIL.Image

        def generate_derivative_PIL(
            original_abs_path: str, derivative_path: str, side_size: int
        ) -> bool:
            _, original_ext = os.path.splitext(original_abs_path)
            _, derivative_ext = os.path.splitext(derivative_path)
            assert original_ext == derivative_ext

            # TODO: More extensions?
            # TODO: Update faq.md in docs when adding new extensions
            if original_ext.lower() not in {
                ".bmp",
                ".jpg",
                ".jpeg",
                ".jpe",
                ".jif",
                ".jfif",
                ".png",
                ".tga",
                ".tiff",
                ".tif",
                ".webp",
            }:
                logger.warning(
                    f"Can't generate derivative of {original_abs_path} because its extension "
                    f"{original_ext} is not supported by PIL."
                )
                return False

            try:
                original_image = PIL.Image.open(original_abs_path)
            except:
                logger.exception(
                    f"Uncaught exception while loading {original_abs_path} image with PIL"
                )
                return False

            original_width, original_height = original_image.size
            if side_size >= original_width and side_size >= original_height:
                logger.warning(
                    f"Refused to generate derivative of {original_abs_path} of side size {side_size} "
                    f"because the original size {original_width}x{original_height} is larger or equal!"
                )
                return False

            derivative_image = original_image.copy()
            # thumbnail() creates image no larger than side_size while keeping original aspect ratio
            derivative_image.thumbnail((side_size, side_size), PIL.Image.Resampling.LANCZOS)
            derivative_image.save(
                derivative_path, icc_profile=original_image.info.get("icc_profile", b"")
            )
            logger.info(
                f"Generated derivative of size {side_size} from original {original_abs_path} using PIL"
            )
            return True

        GENERATORS.append(generate_derivative_PIL)
        logger.info("PIL/Pillow successfully imported and will be used for memsaver.")
    except ImportError:
        logger.error("PIL/Pillow could not be imported, we can't use it for memsaver.")


if len(GENERATORS) == 0:
    logger.error("No generators available, memsaver won't be able to create derivative images!")


def generate_derivative(original_abs_path: str, derivative_path: str, side_size: int) -> bool:
    assert side_size > 0
    logger.debug(
        f"Asked to generate derivative of size {side_size} from original {original_abs_path}"
    )

    for generator in GENERATORS:
        if generator(original_abs_path, derivative_path, side_size):
            return True

    # either it's not supported or its side size is already smaller or equal
    return False


def is_generator_available() -> bool:
    return len(GENERATORS) > 0