cv-analysis-service/cv_analysis/utils/image_operations.py

from typing import Tuple

import cv2 as cv
import numpy as np
from PIL import ImageOps, Image
from loguru import logger

from cv_analysis.utils.conversion import normalize_image_format_to_pil

Color = Tuple[int, int, int]


def blur(image: np.ndarray):
    return cv.blur(image, (3, 3))


def sharpen(image: np.ndarray):
    return cv.filter2D(image, -1, np.array([[-1, -1, -1], [-1, 6, -1], [-1, -1, -1]]))


def overlay(images, mode=np.sum):
    assert mode in [np.sum, np.max]
    images = np.stack(list(images))
    image = mode(images, axis=0)
    image = (image / image.max() * 255).astype(np.uint8)
    return image


def tint_image(src, color="#FFFFFF"):
    src.load()
    r, g, b, alpha = src.split()
    gray = ImageOps.grayscale(src)
    result = ImageOps.colorize(gray, (0, 0, 0), color)
    result.putalpha(alpha)
    return result


def color_shift_array(image: np.ndarray, color: Color):
    """Creates a 3-tensor from a 2-tensor by stacking the 2-tensor three times weighted by the color tuple."""
    assert image.ndim == 3
    assert image.shape[-1] == 3
    assert isinstance(color, tuple)
    assert max(color) <= 255
    assert image.max() <= 255

    color = np.array(color)
    weights = color / color.sum() / 10
    assert max(weights) <= 1

    colored = (image * weights).astype(np.uint8)

    assert colored.shape == image.shape

    return colored


def superimpose(
    base_image: Image,
    image_to_superimpose: Image,
    crop_to_content=True,
    pad=True,
) -> Image:
    """Superimposes an image with transparency onto another image.

    Args:
        base_image: The page image.
        image_to_superimpose: The texture image.
        crop_to_content: If True, the texture will be cropped to content (i.e. the bounding box of all non-transparent
            parts of the texture image).
        pad: If True, the texture will be padded to the size of the page.

    Returns:
        Image where the texture is superimposed onto the page.
    """
    base_image = normalize_image_format_to_pil(base_image)
    image_to_superimpose = normalize_image_format_to_pil(image_to_superimpose)

    if crop_to_content:
        image_to_superimpose = image_to_superimpose.crop(image_to_superimpose.getbbox())

    if base_image.size != image_to_superimpose.size:
        logger.trace(f"Size of page and texture do not match: {base_image.size} != {image_to_superimpose.size}")
        if pad:
            logger.trace(f"Padding texture before pasting to fit size {base_image.size}")
            image_to_superimpose = pad_image_to_size(image_to_superimpose, base_image.size)
        else:
            logger.trace(f"Resizing texture before pasting to fit size {base_image.size}")
            image_to_superimpose = image_to_superimpose.resize(base_image.size)

    assert base_image.size == image_to_superimpose.size
    assert image_to_superimpose.mode == "RGBA"

    base_image.paste(image_to_superimpose, (0, 0), image_to_superimpose)
    return base_image


def pad_image_to_size(image: Image, size: Tuple[int, int]) -> Image:
    """Pads an image to a given size."""
    if image.size == size:
        return image

    if image.size[0] > size[0] or image.size[1] > size[1]:
        raise ValueError(f"Image size {image.size} is larger than target size {size}.")

    padded = Image.new(image.mode, size, color=255)

    pasting_coords = compute_pasting_coordinates(image, padded)
    assert image.mode == "RGBA"
    padded.paste(image, pasting_coords)
    return padded


def compute_pasting_coordinates(smaller: Image, larger: Image.Image):
    """Computes the coordinates for centrally pasting a smaller image onto a larger image."""
    return abs(larger.width - smaller.width) // 2, abs(larger.height - smaller.height) // 2