criminisi.pyx

cpdef _inpaint_criminisi(cnp.float_t [:, ::1] painted,
                         cnp.int8_t [:, ::1] mask,
                         cnp.uint8_t window,
                         cnp.float_t ssd_thresh):

    cdef:
        cnp.float_t [:, ::1] source_image, confidence, gauss_mask, dx, dy
        cnp.float_t [:, ::1] smooth_painted, template, valid_mask
        cnp.int16_t [:, ::1] fill_front
        cnp.int8_t [:, ::1] nx, ny, mask_template
        cnp.uint8_t [:, ::1] mask_ubyte
        Py_ssize_t i_max, j_max, i_m, j_m
        cnp.uint8_t offset
        cnp.float_t sigma

    offset = window // 2
    # inner = (slice(offset, -offset), slice(offset, -offset))
    source_image = painted.base[offset:-offset, offset:-offset].copy()

    t_row, t_col = np.mgrid[-offset:offset + 1, -offset:offset + 1]
    sigma = window / 3
    gauss_mask = _gaussian(sigma, (window, window))
    confidence = 1. - mask.base
    mask_ubyte = mask.base.astype(np.uint8)

    while True:
        # Generate the fill_front, boundary of ROI (region to be synthesised)
        fill_front = mask.base - minimum(mask_ubyte.base, disk(1))
        if not fill_front.base.any():
            if mask.base.any():
                # If the remaining region is 1-pixel thick
                fill_front = mask.base
            else:
                break

        smooth_painted = gaussian_filter(painted, sigma=1)
        smooth_painted.base[mask.base == 1] = np.nan

        # Generate the image gradient and normal vector to the boundary
        # Order reversed; image gradient rotated by 90 degree
        dx = sobel(smooth_painted, axis=0)
        dy = -sobel(smooth_painted, axis=1)
        dx.base[np.isnan(dx)] = 0
        dy.base[np.isnan(dy)] = 0
        smooth_painted.base[np.isnan(smooth_painted)] = 0
        ny = sobel(mask, axis=0)
        nx = sobel(mask, axis=1)

        # Priority calculation; pixel for which inpainting is done first
        i_max, j_max = _priority_calc(fill_front, confidence,
                                      dx, dy, nx, ny, window)

        template = painted.base[i_max + t_row, j_max + t_col]
        mask_template = mask.base[i_max + t_row, j_max + t_col]
        valid_mask = gauss_mask * (1 - mask_template.base)

        # Template matching
        i_m, j_m = _sum_sq_diff(source_image,
                                mask[offset:-offset, offset:-offset], template,
                                valid_mask, ssd_thresh, i_max, j_max)

        if i_m != -1 and j_m != -1:
            painted.base[i_max + t_row, j_max + t_col] += (
                mask_template * painted.base[i_m + t_row, j_m + t_col])
            confidence.base[i_max + t_row, j_max + t_col] += (
                mask_template * confidence.base[i_max, j_max])
            mask_ubyte.base[i_max + t_row, j_max + t_col] = 0
            mask.base[i_max + t_row, j_max + t_col] = 0

    return painted[offset:-offset, offset:-offset]