KHandy/khandy/image/crop_or_pad.py

import numbers

import khandy
import numpy as np


def crop_or_pad(image, x_min, y_min, x_max, y_max, border_value=0):
    """
    See Also:
        translate_image
        
    References:
        tf.image.resize_image_with_crop_or_pad
    """
    assert khandy.is_numpy_image(image)
    assert isinstance(x_min, numbers.Integral) and isinstance(y_min, numbers.Integral)
    assert isinstance(x_max, numbers.Integral) and isinstance(y_max, numbers.Integral)
    assert (x_min <= x_max) and (y_min <= y_max)
    
    src_height, src_width = image.shape[:2]
    dst_height, dst_width = y_max - y_min + 1, x_max - x_min + 1
    channels = 1 if image.ndim == 2 else image.shape[2]
    
    if image.ndim == 2: 
        dst_image_shape = (dst_height, dst_width)
    else:
        dst_image_shape = (dst_height, dst_width, channels)

    if isinstance(border_value, numbers.Real):
        dst_image = np.full(dst_image_shape, border_value, dtype=image.dtype)
    elif isinstance(border_value, tuple):
        assert len(border_value) == channels, \
            'Expected the num of elements in tuple equals the channels' \
            'of input image. Found {} vs {}'.format(
                len(border_value), channels)
        if channels == 1:
            dst_image = np.full(dst_image_shape, border_value[0], dtype=image.dtype)
        else:
            border_value = np.asarray(border_value, dtype=image.dtype)
            dst_image = np.empty(dst_image_shape, dtype=image.dtype)
            dst_image[:] = border_value
    else:
        raise ValueError(
            'Invalid type {} for `border_value`.'.format(type(border_value)))

    src_x_begin = max(x_min, 0)
    src_x_end   = min(x_max + 1, src_width)
    dst_x_begin = src_x_begin - x_min
    dst_x_end   = src_x_end - x_min

    src_y_begin = max(y_min, 0)
    src_y_end   = min(y_max + 1, src_height)
    dst_y_begin = src_y_begin - y_min
    dst_y_end   = src_y_end - y_min
    
    if (src_x_begin >= src_x_end) or (src_y_begin >= src_y_end):
        return dst_image
    dst_image[dst_y_begin: dst_y_end, dst_x_begin: dst_x_end, ...] = \
        image[src_y_begin: src_y_end, src_x_begin: src_x_end, ...]
    return dst_image
    
    
def crop_or_pad_coords(boxes, image_width, image_height):
    """
    References:
        `mmcv.impad`
        `pad` in https://github.com/kpzhang93/MTCNN_face_detection_alignment
        `MtcnnDetector.pad` in https://github.com/AITTSMD/MTCNN-Tensorflow
    """
    x_mins = boxes[:, 0]
    y_mins = boxes[:, 1]
    x_maxs = boxes[:, 2]
    y_maxs = boxes[:, 3]
    dst_widths = x_maxs - x_mins + 1
    dst_heights = y_maxs - y_mins + 1
    
    src_x_begin = np.maximum(x_mins, 0)
    src_x_end   = np.minimum(x_maxs + 1, image_width)
    dst_x_begin = src_x_begin - x_mins
    dst_x_end   = src_x_end - x_mins
    
    src_y_begin = np.maximum(y_mins, 0)
    src_y_end   = np.minimum(y_maxs + 1, image_height)
    dst_y_begin = src_y_begin - y_mins
    dst_y_end   = src_y_end - y_mins

    coords = np.stack([dst_y_begin, dst_y_end, dst_x_begin, dst_x_end, 
                       src_y_begin, src_y_end, src_x_begin, src_x_end, 
                       dst_heights, dst_widths], axis=0)
    return coords
    
    
def center_crop(image, dst_width, dst_height, strict=True):
    """
    strict: 
        when True, raise error if src size is less than dst size. 
        when False, remain unchanged if src size is less than dst size, otherwise center crop.
    """
    assert khandy.is_numpy_image(image)
    assert isinstance(dst_width, numbers.Integral) and isinstance(dst_height, numbers.Integral)
    src_height, src_width = image.shape[:2]
    if strict:
        assert (src_height >= dst_height) and (src_width >= dst_width)

    crop_top = max((src_height - dst_height) // 2, 0)
    crop_left = max((src_width - dst_width) // 2, 0)
    cropped = image[crop_top: dst_height + crop_top, 
                    crop_left: dst_width + crop_left, ...]
    return cropped


def center_pad(image, dst_width, dst_height, strict=True):
    """
    strict: 
        when True, raise error if src size is greater than dst size. 
        when False, remain unchanged if src size is greater than dst size, otherwise center pad.
    """
    assert khandy.is_numpy_image(image)
    assert isinstance(dst_width, numbers.Integral) and isinstance(dst_height, numbers.Integral)
    
    src_height, src_width = image.shape[:2]
    if strict:
        assert (src_height <= dst_height) and (src_width <= dst_width)
    
    padding_x = max(dst_width - src_width, 0)
    padding_y = max(dst_height - src_height, 0)
    padding_top = padding_y // 2
    padding_left = padding_x // 2
    if image.ndim == 2:
        padding = ((padding_top, padding_y - padding_top), 
                   (padding_left, padding_x - padding_left))
    else:
        padding = ((padding_top, padding_y - padding_top), 
                   (padding_left, padding_x - padding_left), (0, 0))
    return np.pad(image, padding, 'constant')