telea_inpaint.py

import numpy as np
from tqdm.auto import tqdm


def make_circle_kernel(radius):
    """

    test:
    
    for radius in [5]:
        knl, wgt = make_circle_kernel(radius)
        knl2 = np.zeros_like(knl)
        cv2.circle(knl2, (radius, radius), radius, 1, -1)
        assert np.all(knl == knl2), (knl, knl2)
    
    with np.printoptions(precision=1, suppress=True):
        print(wgt)
    
    with fourplots() as axs:
        axs[0].imshow(knl)
        axs[1].imshow(wgt)
    
    """
    rad_sq = radius**2
    knl = np.zeros((2 * radius + 1, 2 * radius + 1), "uint8")
    wgt = np.zeros((2 * radius + 1, 2 * radius + 1), "float32")
    for y in range(knl.shape[0]):
        for x in range(knl.shape[1]):
            dist_sq = (x - radius) ** 2 + (y - radius) ** 2
            knl[y, x] = int(dist_sq <= rad_sq) # 0 or 1
            wgt[y, x] = (radius - np.sqrt(dist_sq)) * knl[y, x]
    return knl, wgt


def get_ngbs(bitmaps, i, j, radius):
    # h, w = bitmap.shape[:2]
    # assert (radius <= i <= h - radius - 1) and (radius <= j <= w - radius - 1)
    y1, y2 = i - radius, i + radius + 1
    x1, x2 = j - radius, j + radius + 1
    return [bitmap[y1:y2, x1:x2] for bitmap in bitmaps]


def telea_inpaint(img, mask, radius=None, *, _debug=None):
    """
    Perform TELEA inpainting on an image using the provided mask.

    Args:
        img: Input image as a NumPy array.
        mask: Inpainting mask as a NumPy array (same size as img).
              Non-zero pixels indicate the area to be inpainted.
        radius: Radius of the circular neighborhood considered by the algorithm.

    Returns:
        Inpainted image as a NumPy array.


    Usage:
        # Load the input image
        img = cv2.imread('input.jpg')

        # Create the inpainting mask
        mask = np.zeros(img.shape[:2], dtype=np.uint8)
        mask[100:200, 100:200] = 255  # Example mask region

        # Perform TELEA inpainting
        inpainted = telea_inpaint(img, mask, radius=3)

        # Save the inpainted image
        cv2.imwrite('output.jpg', inpainted)

    """
    h, w = img.shape[:2]
    radius = radius or -int(-min(h, w) / 40)
    mask = (mask != 0).astype("uint8") # strictly 0 or 1

    # pre-compute distance
    knl, wgt = make_circle_kernel(radius)

    # # Find the non-zero pixels in the mask
    # mask_indices = np.argwhere(mask > 0)

    # # # pre-check bounds
    # assert all(
    #     (radius <= i <= h - radius - 1) and (radius <= j <= w - radius - 1)
    #     for i, j in mask_indices
    # ), dict(i=i, j=j, h=h, w=w, radius=radius)
    
    # # sort mask_indices by valid ngb non-masked pixels
    # mask_ngb_pixels = [
    #     np.sum(get_ngbs([mask], i, j, radius)[0] * knl) for i, j in mask_indices
    # ]
    # mask_order = sorted(range(len(mask_ngb_pixels)), key=mask_ngb_pixels.__getitem__)
    # mask_indices = mask_indices[mask_order]

    # Create a copy of the input image
    inpainted = img.copy()
    new_mask = mask.copy()

    for _ in range(9999):

        if new_mask.sum() == 0:
            break


        # Find the non-zero pixels in the mask
        mask_indices = np.argwhere(new_mask > 0)

        # # pre-check bounds
        assert all(
            (radius <= i <= h - radius - 1) and (radius <= j <= w - radius - 1)
            for i, j in mask_indices
        ), dict(i=i, j=j, h=h, w=w, radius=radius)
        
        # sort mask_indices by valid ngb non-masked pixels
        mask_ngb_pixels = [
            np.sum(get_ngbs([new_mask], i, j, radius)[0] * knl) for i, j in mask_indices
        ]
        mask_order = sorted(range(len(mask_ngb_pixels)), key=mask_ngb_pixels.__getitem__)
        mask_indices = mask_indices[mask_order]


        # Iterate over each pixel to be inpainted
        for i, j in mask_indices:  # tqdm(mask_indices):

            # Extract the neighborhood around the pixel
            ngb, ngb_mask = get_ngbs([inpainted, new_mask], i, j, radius)

            # select where mask[y,x] == 0
            weights = wgt * np.uint32(ngb_mask == 0)

            # Normalize the weights
            if np.any(weights > 0):
                weights /= np.sum(weights)

            # Compute the weighted sum of the neighborhood pixels
            # Assign the inpainted value to the pixel
            inpainted[i, j] = np.sum(ngb * weights)
            if np.any(ngb_mask == 0):
                new_mask[i, j] = 0

    if isinstance(_debug, dict):
        _debug.udpate(locals())
        del _debug["_debug"]

    return inpainted