stablediffusion-infinity/PyPatchMatch/patch_match.py

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : patch_match.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 01/09/2020
#
# Distributed under terms of the MIT license.

import ctypes
import os.path as osp
from typing import Optional, Union

import numpy as np
from PIL import Image


__all__ = ['set_random_seed', 'set_verbose', 'inpaint', 'inpaint_regularity']


class CShapeT(ctypes.Structure):
    _fields_ = [
        ('width', ctypes.c_int),
        ('height', ctypes.c_int),
        ('channels', ctypes.c_int),
    ]


class CMatT(ctypes.Structure):
    _fields_ = [
        ('data_ptr', ctypes.c_void_p),
        ('shape', CShapeT),
        ('dtype', ctypes.c_int)
    ]


PMLIB = ctypes.CDLL(osp.join(osp.dirname(__file__), 'libpatchmatch.so'))

PMLIB.PM_set_random_seed.argtypes = [ctypes.c_uint]
PMLIB.PM_set_verbose.argtypes = [ctypes.c_int]
PMLIB.PM_free_pymat.argtypes = [CMatT]
PMLIB.PM_inpaint.argtypes = [CMatT, CMatT, ctypes.c_int]
PMLIB.PM_inpaint.restype = CMatT
PMLIB.PM_inpaint_regularity.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
PMLIB.PM_inpaint_regularity.restype = CMatT
PMLIB.PM_inpaint2.argtypes = [CMatT, CMatT, CMatT, ctypes.c_int]
PMLIB.PM_inpaint2.restype = CMatT
PMLIB.PM_inpaint2_regularity.argtypes = [CMatT, CMatT, CMatT, CMatT, ctypes.c_int, ctypes.c_float]
PMLIB.PM_inpaint2_regularity.restype = CMatT


def set_random_seed(seed: int):
    PMLIB.PM_set_random_seed(ctypes.c_uint(seed))


def set_verbose(verbose: bool):
    PMLIB.PM_set_verbose(ctypes.c_int(verbose))


def inpaint(
    image: Union[np.ndarray, Image.Image],
    mask: Optional[Union[np.ndarray, Image.Image]] = None,
    *,
    global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
    patch_size: int = 15
) -> np.ndarray:
    """
    PatchMatch based inpainting proposed in:

        PatchMatch : A Randomized Correspondence Algorithm for Structural Image Editing
        C.Barnes, E.Shechtman, A.Finkelstein and Dan B.Goldman
        SIGGRAPH 2009

    Args:
        image (Union[np.ndarray, Image.Image]): the input image, should be 3-channel RGB/BGR.
        mask (Union[np.array, Image.Image], optional): the mask of the hole(s) to be filled, should be 1-channel.
        If not provided (None), the algorithm will treat all purely white pixels as the holes (255, 255, 255).
        global_mask (Union[np.array, Image.Image], optional): the target mask of the output image.
        patch_size (int): the patch size for the inpainting algorithm.

    Return:
        result (np.ndarray): the repaired image, of the same size as the input image.
    """

    if isinstance(image, Image.Image):
        image = np.array(image)
    image = np.ascontiguousarray(image)
    assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'

    if mask is None:
        mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
        mask = np.ascontiguousarray(mask)
    else:
        mask = _canonize_mask_array(mask)

    if global_mask is None:
        ret_pymat = PMLIB.PM_inpaint(np_to_pymat(image), np_to_pymat(mask), ctypes.c_int(patch_size))
    else:
        global_mask = _canonize_mask_array(global_mask)
        ret_pymat = PMLIB.PM_inpaint2(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(global_mask), ctypes.c_int(patch_size))

    ret_npmat = pymat_to_np(ret_pymat)
    PMLIB.PM_free_pymat(ret_pymat)

    return ret_npmat


def inpaint_regularity(
    image: Union[np.ndarray, Image.Image],
    mask: Optional[Union[np.ndarray, Image.Image]],
    ijmap: np.ndarray,
    *,
    global_mask: Optional[Union[np.ndarray, Image.Image]] = None,
    patch_size: int = 15, guide_weight: float = 0.25
) -> np.ndarray:
    if isinstance(image, Image.Image):
        image = np.array(image)
    image = np.ascontiguousarray(image)

    assert isinstance(ijmap, np.ndarray) and ijmap.ndim == 3 and ijmap.shape[2] == 3 and ijmap.dtype == 'float32'
    ijmap = np.ascontiguousarray(ijmap)

    assert image.ndim == 3 and image.shape[2] == 3 and image.dtype == 'uint8'
    if mask is None:
        mask = (image == (255, 255, 255)).all(axis=2, keepdims=True).astype('uint8')
        mask = np.ascontiguousarray(mask)
    else:
        mask = _canonize_mask_array(mask)


    if global_mask is None:
        ret_pymat = PMLIB.PM_inpaint_regularity(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(ijmap), ctypes.c_int(patch_size), ctypes.c_float(guide_weight))
    else:
        global_mask = _canonize_mask_array(global_mask)
        ret_pymat = PMLIB.PM_inpaint2_regularity(np_to_pymat(image), np_to_pymat(mask), np_to_pymat(global_mask), np_to_pymat(ijmap), ctypes.c_int(patch_size), ctypes.c_float(guide_weight))

    ret_npmat = pymat_to_np(ret_pymat)
    PMLIB.PM_free_pymat(ret_pymat)

    return ret_npmat


def _canonize_mask_array(mask):
    if isinstance(mask, Image.Image):
        mask = np.array(mask)
    if mask.ndim == 2 and mask.dtype == 'uint8':
        mask = mask[..., np.newaxis]
    assert mask.ndim == 3 and mask.shape[2] == 1 and mask.dtype == 'uint8'
    return np.ascontiguousarray(mask)


dtype_pymat_to_ctypes = [
    ctypes.c_uint8,
    ctypes.c_int8,
    ctypes.c_uint16,
    ctypes.c_int16,
    ctypes.c_int32,
    ctypes.c_float,
    ctypes.c_double,
]


dtype_np_to_pymat = {
    'uint8': 0,
    'int8': 1,
    'uint16': 2,
    'int16': 3,
    'int32': 4,
    'float32': 5,
    'float64': 6,
}


def np_to_pymat(npmat):
    assert npmat.ndim == 3
    return CMatT(
        ctypes.cast(npmat.ctypes.data, ctypes.c_void_p),
        CShapeT(npmat.shape[1], npmat.shape[0], npmat.shape[2]),
        dtype_np_to_pymat[str(npmat.dtype)]
    )


def pymat_to_np(pymat):
    npmat = np.ctypeslib.as_array(
        ctypes.cast(pymat.data_ptr, ctypes.POINTER(dtype_pymat_to_ctypes[pymat.dtype])),
        (pymat.shape.height, pymat.shape.width, pymat.shape.channels)
    )
    ret = np.empty(npmat.shape, npmat.dtype)
    ret[:] = npmat
    return ret