src/com_image.py

from pathlib import Path
from typing import Any, Union

import numpy as np

import cv2

from PIL import Image, ImageEnhance


def load_image(file_name, path_to_images=None, rgb: bool = True):
    path = (
        file_name
        if isinstance(file_name, Path) is True
        else path_to_images.joinpath(file_name)
    )

    try:
        img = cv2.imread(str(path))
        if rgb is True:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    except Exception as e:
        print(file_name)
    return img


def to_pil(image):
    return Image.fromarray(image)


def to_cv2(image):
    return np.array(image)


def enhance_pil_image(
    image, color=1, brightness=1, contrast=1, sharpness=1
) -> Image.Image:
    image = ImageEnhance.Sharpness(
        image=ImageEnhance.Brightness(
            image=ImageEnhance.Contrast(
                image=ImageEnhance.Color(
                    image=(
                        image
                        if isinstance(image, Image.Image) is True
                        else to_pil(image=image)
                    )
                ).enhance(color)
            ).enhance(contrast)
        ).enhance(brightness)
    ).enhance(sharpness)
    return image


def ensure_odd(
    i: int,
    min_val: Union[None, int] = None,
    max_val: Union[None, int] = None,
) -> int:
    """Transforms an odd number into pair number by adding one
    Arguments:
        i {int} -- number
    Returns:
        int -- Odd number
    """
    if (i > 0) and (i % 2 == 0):
        i += 1
    if min_val is not None:
        return max(i, min_val)
    if max_val is not None:
        return min(i, max_val)
    return i


def get_morphology_kernel(size: int, shape: int):
    """Builds morphology kernel
    :param size: kernel size, must be odd number
    :param shape: select shape of kernel
    :return: Morphology kernel
    """
    size = ensure_odd(size)
    return cv2.getStructuringElement(shape, (size, size))


def close(
    image: Any,
    kernel_size: int = 3,
    kernel_shape: int = cv2.MORPH_ELLIPSE,
    rois: tuple = (),
    proc_times: int = 1,
):
    """Morphology - Close wrapper
    Arguments:
        image {numpy array} -- Source image
        kernel_size {int} -- kernel size
        kernel_shape {int} -- cv2 constant
        roi -- Region of Interest
        proc_times {int} -- iterations
    Returns:
        numpy array -- closed image
    """
    morph_kernel = get_morphology_kernel(kernel_size, kernel_shape)
    if rois:
        result = image.copy()
        for roi in rois:
            r = roi.as_rect()
            result[r.top : r.bottom, r.left : r.right] = cv2.morphologyEx(
                result[r.top : r.bottom, r.left : r.right],
                cv2.MORPH_CLOSE,
                morph_kernel,
                iterations=proc_times,
            )
    else:
        result = cv2.morphologyEx(
            image, cv2.MORPH_CLOSE, morph_kernel, iterations=proc_times
        )
    return result


def get_concat_h_multi_resize(im_list, resample=Image.Resampling.BICUBIC):
    min_height = min(im.height for im in im_list)
    im_list_resize = [
        im.resize(
            (int(im.width * min_height / im.height), min_height), resample=resample
        )
        for im in im_list
    ]
    total_width = sum(im.width for im in im_list_resize)
    dst = Image.new("RGB", (total_width, min_height))
    pos_x = 0
    for im in im_list_resize:
        dst.paste(im, (pos_x, 0))
        pos_x += im.width
    return dst


def get_concat_v_multi_resize(im_list, resample=Image.Resampling.BICUBIC):
    min_width = min(im.width for im in im_list)
    im_list_resize = [
        im.resize((min_width, int(im.height * min_width / im.width)), resample=resample)
        for im in im_list
    ]
    total_height = sum(im.height for im in im_list_resize)
    dst = Image.new("RGB", (min_width, total_height))
    pos_y = 0
    for im in im_list_resize:
        dst.paste(im, (0, pos_y))
        pos_y += im.height
    return dst


def get_concat_tile_resize(im_list_2d, resample=Image.Resampling.BICUBIC):
    im_list_v = [
        get_concat_h_multi_resize(im_list_h, resample=resample)
        for im_list_h in im_list_2d
    ]
    return get_concat_v_multi_resize(im_list_v, resample=resample)


def get_tiles(img_list, row_count, resample=Image.Resampling.BICUBIC):
    if isinstance(img_list, np.ndarray) is False:
        img_list = np.asarray(img_list, dtype="object")
    return get_concat_tile_resize(np.split(img_list, row_count), resample)