utils.py

import cv2
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image


def fill_corruption_with_avg(img, mask, step_mask=None, ksize = 3):
    """
    Method to fill a corrupted area of a 2D image with an approximation computed
    starting from a neighbor of every corrupted pixel.

    This is meant to be used as a step of a recursive procedure.

    img: input greyscale 2D image (0 is white, 1 is black)
    mask: corruption characteristic function (0 is image, 1 is corruption)
    step_mask: boundary of the corrupted area. If None, the whole mask is used
    ksize: kernel size for neighbors

    @returns: A filled 2D image
    """

    #If the boundary mask is not provided then the whole mask is used as boundary
    if step_mask is None:
        step_mask = mask

    I = img.copy()

    #Santity check on the types
    sm = step_mask.astype(np.bool_)

    #I[m] = 0
    kern = np.ones((ksize, ksize), dtype=np.float32)

    #Count the number of pixels contributing to the sum
    avg_not_mask = cv2.filter2D((1-mask).astype(np.float32), -1, kern)

    #Sum all pixels within a kernel of size ksize
    blurred = cv2.filter2D((I).astype(np.float32), -1, kern)

    #For every point in the step mask (boundary), substitute the pixel with the average
    I[sm]=blurred[sm]/avg_not_mask[sm]
    return I

def fill_corruption_with_bfs_avg(img, mask, ksize = 3, advanced_average=False, blurred_image=None):
    """
    Method to fill a corrupted area of a 2D image with an approximation computed
    starting from a neighbor of every corrupted pixel. The procedure is done
    recursively considering at each step the boundary of the corrupted area in
    a BFS (Breadth-First-Search) style

    img: input greyscale 2D image (0 is white, 1 is black)
    mask: corruption characteristic function (0 is image, 1 is corruption)
    ksize: kernel size for neighbors
    advanced_average: False if the fill is done according to the image, True is Advanced Average
    blurred_version: if it is advanced average, then a blurred version should be provided

    @returns: A filled 2D image
    """

    #Sanity check
    if advanced_average and blurred_image is None:
        raise Exception("Cannot compute advance average without blurred_version")

    #We expect I[m] = 0 on the mask
    I = img.copy()
    m = mask.copy()

    flag = True
    while flag:
        flag = False
        step_mask = np.zeros(m.shape, dtype=np.bool_)
        for y in range(I.shape[0]):
            for x in range(I.shape[1]):
                # Check if pixel (x,y) in on the edge of the non-filled corruption
                # 6-nbh apprach
                if y==0:
                    if m[y,x]==1 and m[y+1,x]==0:
                        step_mask[y,x]=True
                elif y==I.shape[0]-1:
                    if m[y,x]==1 and m[y-1,x]==0:
                        step_mask[y,x]=True
                else:
                    if m[y,x]==1 and (m[y-1,x]==0 or m[y+1,x]==0):
                        step_mask[y,x]=True
                if x==0:
                    if m[y,x]==1 and m[y,x+1]==0:
                        step_mask[y,x]=True
                elif x==I.shape[1]-1:
                    if m[y,x]==1 and m[y,x-1]==0:
                        step_mask[y,x]=True
                else:
                    if m[y,x]==1 and (m[y,x-1]==0 or m[y,x+1]==0):
                        step_mask[y,x]=True
        if np.sum(step_mask)>0:
            flag = True
            if advanced_average:
                I = fill_corruption_with_advanced_avg_step(I, blurred_image, m, step_mask ,ksize)
            else:
                I = fill_corruption_with_avg(I, mask=m, step_mask=step_mask, ksize=ksize)
            m[step_mask] = False #Remove the border that we just filled from the mask
    return I

def fill_corruption_with_advanced_avg_step(
    original_img, blurred_img, mask, step_mask=None, ksize=3
):
    """
    Method to fill a corrupted area of a 2D image with an approximation computed
    starting from a neighbor of every corrupted pixel

    Direct use of this method is to be avoided, as a point positioned in a
    fully corrupted neighborhood will stay corrupted

    img: input greyscale 2D image (0 is white, 1 is black)
    mask: corruption characteristic function (0 is image, 1 is corruption)
    step_mask: boundary of the corrupted area. If None, the whole mask is used
    ksize: kernel size for neighbors

    @returns: A filled 2D image
    """

    #If the boundary mask is not provided then the whole mask is used as boundary
    if step_mask is None:
        step_mask = mask

    I = original_img.copy()
    sm = step_mask.astype(np.bool_)
    kern = np.ones((ksize, ksize))

    #Compute the numerator and the denominator of the formula for the Advanced Averaging
    fh = ((original_img**(-1)) * (blurred_img**(-1)))*(1-mask)
    ff = (original_img**(-2)) * (1-mask)

    num = cv2.filter2D(np.array(fh), -1, kern)
    den = cv2.filter2D(np.array(ff), -1, kern)

    #For every point in the step mask (boundary), substitute the pixel with the average
    I[sm]=blurred_img[sm]*num[sm]/den[sm]

    return I

def compute_mask_border(mask):
    m = mask.copy()

    step_mask = np.zeros(m.shape, dtype=np.bool_)
    for y in range(mask.shape[0]):
        for x in range(mask.shape[1]):
            # Check if pixel (x,y) in on the edge of the non-filled corruption
            # 6-nbh apprach
            if y==0:
                if m[y,x]==1 and m[y+1,x]==0:
                    step_mask[y,x]=True
            elif y==mask.shape[0]-1:
                if m[y,x]==1 and m[y-1,x]==0:
                    step_mask[y,x]=True
            else:
                if m[y,x]==1 and (m[y-1,x]==0 or m[y+1,x]==0):
                    step_mask[y,x]=True
            if x==0:
                if m[y,x]==1 and m[y,x+1]==0:
                    step_mask[y,x]=True
            elif x==mask.shape[1]-1:
                if m[y,x]==1 and m[y,x-1]==0:
                    step_mask[y,x]=True
            else:
                if m[y,x]==1 and (m[y,x-1]==0 or m[y,x+1]==0):
                    step_mask[y,x]=True
    return step_mask


def imshow(img, r=None, mode = 2):
    """
    Handler method to show a greyscale numpy/opencv image using matplotlib

    TODO: extend the method to support color images

    img: a greyscale image
    r: range expressed as a tuple (min, max). If not provided (0,1) is used

    @returns: None
    """
    plt.figure(figsize=(8,8))
    if r is None:
        r= (0,1)
    if mode==1:
        plt.imshow(img, vmin=r[0], vmax=r[1], cmap="gray")
    elif mode==2:
        plt.imshow(1-img, vmin=r[0], vmax=r[1], cmap="gray")
    plt.axis('off')
    plt.show()


def ang_d(a, b, unit="radians", absolute=True):
    """
    Distance on S^1 the unit circle, expressed in angle
    Works with integers, floats and numpy arrays

    a:  first angle
    b:  second angle
    unit:   "radians" or "degrees", unit of measure of the angles
    absolute: if the required distance is regardless of orientation

    @returns: the distance on S^1 between the two angles
    """

    u = np.pi #radians

    if unit=="degrees":
        u = 180

    if unit=="P1":
        u = np.pi/2

    c = (a - b) % (2*u)
    if isinstance(a, (int, float)):  # scalar case
        if c > u:
            c -= 2*u
    else:  # numpy array case
        c[c>u] = c[c>u]-2*u
    if absolute:
        c = np.abs(c)
    return c