uz_assignments/assignment3/uz_framework/image.py

import numpy as np
import cv2 as cv2
from matplotlib import pyplot as plt
from PIL import Image
from typing import Union
import numpy.typing as npt
import enum

class ImageType(enum.Enum):
    uint8 = 0
    float64 = 1

class DistanceMeasure(enum.Enum):
    euclidian_distance = 0
    chi_square_distance = 1
    intersection_distance = 2
    hellinger_distance = 3

def imread(path: str, type: ImageType) -> Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]]:
    """
    Reads an image in RGB order. Image type is transformed from uint8 to float, and
    range of values is reduced from [0, 255] to [0, 1].
    """
    I = Image.open(path).convert('RGB')  # PIL image.
    I = np.asarray(I)  # Converting to Numpy array.
    if type == ImageType.float64:
        I = I.astype(np.float64) / 255
        return I
    elif type == ImageType.uint8:
        return I

    raise Exception(f"Unrecognized image format! {type}")


def imread_gray(path: str, type: ImageType) -> Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]]:
    """
    Reads an image in gray. Image type is transformed from uint8 to float, and
    range of values is reduced from [0, 255] to [0, 1].
    """

    I = Image.open(path).convert('L')  # PIL image opening and converting to gray.
    I = np.asarray(I)  # Converting to Numpy array.

    if type == ImageType.float64:
        I = I.astype(np.float64) / 255
        return I
    elif type == ImageType.uint8:
        return I

    raise Exception("Unrecognized image format!")

def signal_show(*signals):
    """
    Plots all given 1D signals in the same plot.
    Signals can be Python lists or 1D numpy array.
    """
    for s in signals:
        if type(s) == np.ndarray:
            s = s.squeeze()
        plt.plot(s)
    plt.show()


def convolve(I: np.ndarray, *ks):
    """
    Convolves input image I with all given kernels.

    :param I: Image, should be of type float64 and scaled from 0 to 1.
    :param ks: 2D Kernels
    :return: Image convolved with all kernels.
    """
    for k in ks:
        k = np.flip(k)  # filter2D performs correlation, so flipping is necessary
        I = cv2.filter2D(I, cv2.CV_64F, k)
    return I

def transform_coloured_image_to_grayscale(image: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    """
    Accepts float64 picture with three colour channels and returns float64 grayscale image
    with one channel.
    """
    grayscale_image = np.zeros(image.shape[:2])

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            grayscale_image[i, j] = (image[i, j, 0] + image[i,j, 1] + image[i, j, 2]) / 3


    return grayscale_image

def invert_coloured_image_part(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]], startx: int, endx: int, starty: int, endy: int) -> Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]]:
    """
    Accepts image, starting position end end position for axes x & y. Returns whole image with inverted part.
    """
    inverted_image = image.copy() 

    if image.dtype.type == np.float64:     
        for i in range(startx, endx):
            for j in range(starty, endy):
                inverted_image[i, j, 0] = 1 - image[i, j, 0]
                inverted_image[i, j, 1] = 1 - image[i, j, 1]
                inverted_image[i, j, 2] = 1 - image[i, j, 2]
        return inverted_image
    elif image.dtype.type == np.uint8:
        for i in range(startx, endx):
            for j in range(starty, endy):
                inverted_image[i, j, 0] = 255 - image[i, j, 0]
                inverted_image[i, j, 1] = 255 - image[i, j, 1]
                inverted_image[i, j, 2] = 255 - image[i, j, 2]
        return inverted_image

    raise Exception("Unrecognized image format!")

def invert_coloured_image(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]]) -> Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]]:
    """
    Accepts image and inverts it
    """
    return invert_coloured_image_part(image, 0, image.shape[0], 0, image.shape[1])

def calculate_best_treshold_using_otsu_method(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]]) -> int:
    """
    Accepts image and returns best treshold using otsu method
    """

    if image.dtype.type == np.float64: 
        im = image.copy()
        im  = im * (255.0/im.max())
    elif image.dtype.type == np.uint8:
        im = image.copy()
    else:
        raise Exception("Unrecognized image format!")
        
    treshold_range = np.arange(np.max(im) + 1) 
    criterias = []

    for treshold in treshold_range:
        # create the thresholded image
        thresholded_im = np.zeros(im.shape)

        thresholded_im[im >= treshold] = 1

        # compute weights
        nb_pixels = im.size
        nb_pixels1 = np.count_nonzero(thresholded_im)
        weight1 = nb_pixels1 / nb_pixels
        weight0 = 1 - weight1

        # if one the classes is empty, eg all pixels are below or above the threshold, that threshold will not be considered
        # in the search for the best threshold
        if weight1 == 0 or weight0 == 0:
            continue 

        # find all pixels belonging to each class
        val_pixels1 = im[thresholded_im == 1]
        val_pixels0 = im[thresholded_im == 0]

        # compute variance of these classes
        var0 = np.var(val_pixels0) if len(val_pixels0) > 0 else 0
        var1 = np.var(val_pixels1) if len(val_pixels1) > 0 else 0

        criterias.append( weight0 * var0 + weight1 * var1)

    best_threshold = treshold_range[np.argmin(criterias)]
    return best_threshold


def get_image_bins_for_loop(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]], number_of_bins: int) -> npt.NDArray[np.float64]:
    """
    Accepts image in the float64 format or uint8, returns normailzed
    image bins, histogram
    """
    if image.dtype.type == np.float64 or image.dtype.type == np.uint8: 
        bin_restrictions = np.linspace(np.min(image), np.max(image), num=number_of_bins)
    else:
        raise Exception("Unrecognized image format!")

    bins = np.zeros(number_of_bins).astype(np.float64)

    for pixel in image.reshape(-1):
        # https://stackoverflow.com/a/16244044
        bins[np.argmax(bin_restrictions > pixel)] += 1

    return bins / np.sum(bins)


# Much faster implementation than for loop
def get_image_bins(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]], number_of_bins: int) -> npt.NDArray[np.float64]:
    """
    Accepts image in the float64 format or uint8, returns normailzed
    image bins, histogram
    """
    if image.dtype.type == np.float64 or image.dtype.type == np.uint8: 
        bins = np.linspace(np.min(image), np.max(image), num=number_of_bins)
    else:
        raise Exception("Unrecognized image format!")

    # Put pixels into classes
    # ex. binsize = 10 then 0.4 would map into 4
    binarray = np.digitize(image.reshape(-1), bins).astype(np.uint8)

    # Now count those values 
    binarray = np.unique(binarray, return_counts=True)

    counts = binarray[1].astype(np.float64) # Get the counts out of tuple

    # Check if there is any empty bin
    empty_bins = []
    bins = binarray[0]
    for i in range(1, number_of_bins + 1):
        if i not in bins:
            empty_bins.append(i)
    
    # Add empty bins with zeros
    if empty_bins != []:
        for i in empty_bins:
            counts = np.insert(counts, i - 1, 0)

    return counts

def get_image_bins_ND(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]], number_of_bins: int) -> npt.NDArray[np.float64]:
    """
    Accepts image in the float64 format or uint8 and number of bins
    Returns normailzed image histogram bins
    """ 
    bs = [] 
    hist = np.zeros((number_of_bins, number_of_bins, number_of_bins))
    bins = np.linspace(np.min(image), np.max(image), num=number_of_bins)

    for i in range(image.shape[2]):
        v = image[:, :, i].reshape(-1)
        bs.append(np.digitize(v, bins).astype(np.uint32))

    for i in range(len(bs[0])):
        hist[bs[2][i] -1, bs[1][i] -1, bs[0][i] - 1] += 1

    return hist / np.sum(hist)

def compare_two_histograms(h1: npt.NDArray[np.float64], h2: npt.NDArray[np.float64], method: DistanceMeasure) -> float:
    """
    Accepts two histograms and method of comparison
    Returns distance between them
    """

    if method == DistanceMeasure.euclidian_distance:
        d = np.sqrt(np.sum(np.square(h1 - h2)))
    elif method == DistanceMeasure.chi_square_distance:
        d = 0.5 * np.sum(np.square(h1 - h2) / (h1 + h2 + np.finfo(float).eps))
    elif method == DistanceMeasure.intersection_distance:
        d = 1 - np.sum(np.minimum(h1, h2))
    elif method == DistanceMeasure.hellinger_distance:
        d = np.sqrt(0.5 * np.sum(np.square(np.sqrt(h1) - np.sqrt(h2))))
    else:
        raise Exception('Unsuported method chosen!')

    return d.astype(float)
    

def apply_mask_on_image(image: Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]], mask: npt.NDArray[np.uint8]) -> Union[npt.NDArray[np.float64],  npt.NDArray[np.uint8]]:
    """
    Accepts image and applys mask to image
    """
    image = image.copy()

    mask = np.expand_dims(mask, axis=2) 
    image = mask * image

    return image

def simple_convolution(signal: npt.NDArray[np.float64], kernel: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    """
        Accepts: signal & kernel

        Returns: convolved signal with a kernel
    """
    N = int(np.ceil(len(kernel) / 2 - 1))
    n_conv =  signal.size - kernel.size  + 1
    print(N, signal.size -N, n_conv)

    convolved_signal = np.zeros(len(signal))
    rev_kernel = kernel[::-1].copy()

    for i in range(n_conv):
        convolved_signal[i] = np.dot(signal[i: i+kernel.size], rev_kernel) # Well if you would add i+N then you wuold shift this

    return convolved_signal

def simple_convolution_improved(signal: npt.NDArray[np.float64], kernel: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    """
        Accepts: signal & kernel

        Returns: convolved signal with a kernel
        Improved method replicates edges of an signal
    """
    signal_len = signal.size
    kernel_len = kernel.size
    signal = signal.copy()
    
    # Calculate which values to fill in and range
    EDGE_FRONT = signal[0]
    EDGE_BACK = signal[-1]
    EXTEND_RANGE = int(np.floor(kernel.size / 2))
    
    # Append end insert edges
    front = [ EDGE_FRONT for _ in range(EXTEND_RANGE)]
    back = [ EDGE_BACK for _ in range(EXTEND_RANGE)]
    signal = np.insert(signal, 1, front, axis=0)
    signal = np.append(signal, back, axis=0)
    
    convolved_signal = np.zeros(signal_len)
    rev_kernel = kernel[::-1].copy()
    n_conv =  signal_len - kernel_len  + 1

    for i in range(n_conv):
        convolved_signal[i] = np.dot(signal[i: i+kernel.size], rev_kernel) 

    return convolved_signal

def get_gaussian_kernel(sigma: float) -> npt.NDArray[np.float64]:
    """
    Accepts sigma
    Returns gaussian kernel
    """
    # https://github.com/mikepound/convolve/blob/master/run.gaussian.py
    kernel_size = int(2 * np.ceil(3*sigma) + 1)
    k_min_max = np.ceil(3*sigma)
    k_interval = np.arange(-k_min_max, k_min_max + 1.)

    result = (1. / (np.sqrt(2. * np.pi )* sigma)) * np.exp(- (np.square(k_interval)) / (2.*np.square(sigma)))

    assert(kernel_size == len(result))

    return result / np.sum(result)

def sharpen_image(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sharpen_factor=1.0) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Accepts: image & sharpen factor 

    Returns: sharpened image
    https://blog.demofox.org/2022/02/26/image-sharpening-convolution-kernels/ <-- sharpening kernel, but also on slides
    good explanation
    """
    sharpened_image = image.copy()

    KERNEL = np.array([[-1, -1, -1], 
                       [-1, 17, -1],
                       [-1, -1,-1]]) * 1./9. * sharpen_factor

    if image.dtype.type == np.float64: 
        sharpened_image = cv2.filter2D(sharpened_image, cv2.CV_64F, KERNEL)
    elif image.dtype.type == np.uint8:
        sharpened_image = cv2.filter2D(sharpened_image, cv2.CV_8U, KERNEL)

    return sharpened_image

def gaussfilter2D(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
        Accepts: image, sigma  
        Applies gaussian noise on image
        returns: filtered_image
    """
    filtered_image = image.copy()
    kernel = np.array([get_gaussian_kernel(sigma)])
    filtered_image = cv2.filter2D(filtered_image, cv2.CV_64F, kernel)
    filtered_image = cv2.filter2D(filtered_image, cv2.CV_64F, kernel.T)
    
    return filtered_image

def gaussdx(sigma: float) -> npt.NDArray[np.float64]:
    """
    Accepts sigma
    Returns gaussian kernel
    """
    kernel_size = int(2 * np.ceil(3*sigma) + 1)
    k_min_max = np.ceil(3*sigma)
    k_interval = np.arange(-k_min_max, k_min_max + 1.)

    result = - (1. / (np.sqrt(2. * np.pi )* np.power(sigma, 3))) * k_interval* np.exp(- (np.square(k_interval)) / (2.*np.square(sigma)))

    assert(kernel_size == len(result))

    return result / np.sum(np.abs(result))


def simple_median(signal: npt.NDArray[np.float64], width: int) -> npt.NDArray[np.float64]:
    """
    Accepts: signal & width
    returns signal improved using median filter
    """
    if width % 2 == 0:
        raise Exception('No u won\'t do that')

    signal = signal.copy()
    for i in range(len(signal) - int(np.ceil(width/2))):
        middle_element = int(i + np.floor(width/2))
        signal[middle_element] = np.median(signal[i:i+width])
    return signal

def apply_median_method_2D(image:Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], width: int) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Accepts: image & filter width
    returns: image with median filter applied
    """
    if width % 2 == 0:
        raise Exception('No u won\'t do that')

    image = image.copy()
    W_HALF = int(np.floor(width/2))
    padded_image = np.pad(image, W_HALF, mode='edge')
    IMAGE_HEIGHT = image.shape[0] # y
    IMAGE_WIDTH = image.shape[1] # x

    for x in range(W_HALF, IMAGE_WIDTH): # I think we can start from 0, cuz we padded an image
        for y in range(W_HALF, IMAGE_HEIGHT):
            median_filter = np.zeros(0)
            STARTX = x - W_HALF
            STARTY = y - W_HALF
            for m in range(width):
                median_filter = np.append(median_filter, padded_image[STARTY + m][STARTX: STARTX + width], axis=0)
            if image.dtype.type == np.uint8:
                image[y][x] = int(np.mean(median_filter))
            else:
                image[y][x] = np.mean(median_filter)
    return image

def filter_laplace(image:Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Accepts: image & sigma
    returns: image with laplace filter applied
    """
    # Prepare unit impulse and gauss kernel
    unit_impulse = np.zeros((1, 2 * int(np.ceil(3*sigma)) + 1))
    unit_impulse[0][int(np.ceil(unit_impulse.size /2)) - 1]= 1
    gauss_kernel = np.array([get_gaussian_kernel(sigma)])
    assert(len(gauss_kernel[0]) == len(unit_impulse[0]))

    laplacian_filter = unit_impulse - gauss_kernel[0]

    # Now apply laplacian filter
    applied_by_x = cv2.filter2D(image, -1, laplacian_filter)
    applied_by_y = cv2.filter2D(applied_by_x, -1, laplacian_filter.T)
    
    return applied_by_y

def gauss_noise(I, magnitude=.1) -> npt.NDArray[np.float64]:
    """
    Accepts: image & magnitude
    Returns: image with gaussian noise applied
    """
    # input: image, magnitude of noise
    # output: modified image
    I = I.copy()

    return I + np.random.normal(size=I.shape) * magnitude


def sp_noise(I, percent=.1) -> npt.NDArray[np.float64]:
    """
    Accepts: image & percent
    Returns: image with salt and pepper noise applied
    """
    # input: image, percent of corrupted pixels
    # output: modified image

    res = I.copy()
    res[np.random.rand(I.shape[0], I.shape[1]) < percent / 2] = 1
    res[np.random.rand(I.shape[0], I.shape[1]) < percent / 2] = 0

    return res

def sp_noise1D(signal, percent=.1) -> npt.NDArray[np.float64]:
    """
    Accepts: signal & percent
    Returns: signal with salt and pepper noise applied
    """
    signal = signal.copy()
    signal[np.random.rand(signal.shape[0]) < percent / 2] = 2
    signal[np.random.rand(signal.shape[0]) < percent / 2] = 1
    signal[np.random.rand(signal.shape[0]) < percent / 2] = 4
    signal[np.random.rand(signal.shape[0]) < percent / 2] = 0.4
    return signal

def sum_two_grayscale_images(image_a: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], image_b :Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Accepts: image_a, image_b
    Returns: image_a + image_b
    """
    # Merge image_a and image_b
    return (image_a + image_b)/ 2

def generate_dirac_impulse(size: int) -> npt.NDArray[np.float64]:
    """
    Accepts: size
    Returns: dirac impulse of size
    """

    dirac_impulse = np.zeros((size, size))
    dirac_impulse[int(size/2), int(size/2)] = 1

    return dirac_impulse

def derive_image_by_x(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Accepts: image
    Returns: image derived by x
    """
    image = image.copy()
    gaussd = np.array([gaussdx(sigma)])
    gauss = np.array([get_gaussian_kernel(sigma)])
    gaussd = np.flip(gaussd, axis=1)
    
    applied_by_y = cv2.filter2D(image, cv2.CV_64F, gauss.T)
    applied_by_x = cv2.filter2D(applied_by_y, cv2.CV_64F, gaussd)

    return applied_by_x

def derive_image_by_y(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Accepts: image
    Returns: image derived by y
    """
    image = image.copy()
    gaussd = np.array([gaussdx(sigma)])
    gauss = np.array([get_gaussian_kernel(sigma)])
    gaussd = np.flip(gaussd, axis=1)

    applied_by_x = cv2.filter2D(image, cv2.CV_64F, gauss)
    applied_by_y = cv2.filter2D(applied_by_x, cv2.CV_64F, gaussd.T)

    return applied_by_y

def derive_image_first_order(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float) -> tuple[Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]]:
    """
    Accepts: image
    returns: image derived by x, image derived by y
    """
    return derive_image_by_x(image, sigma), derive_image_by_y(image, sigma)

def derive_image_second_order(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float) -> tuple[tuple[Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]], tuple[Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]]]:
    """
    Accepts: image
    Returns: Ixx, Ixy, Iyx, Iyy 
    """
    derived_by_x = derive_image_by_x(image, sigma)
    derived_by_y = derive_image_by_y(image, sigma)

    return derive_image_first_order(derived_by_x, sigma), derive_image_first_order(derived_by_y, sigma)

def gradient_magnitude(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float) -> tuple[Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]]:
    """
    Accepts: image
    Returns: gradient magnitude of image and derivative angles
    """
    Ix, Iy = derive_image_first_order(image, sigma)
    return np.sqrt(Ix**2 + Iy**2), np.arctan2(Iy, Ix)


def find_edges_primitive(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float, theta: float) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Aceppts: image, sigma & theta
    Returns: image with edges
    """
    derivative_magnitude, _ = gradient_magnitude(image, sigma)

    binary_mask = np.zeros_like(derivative_magnitude)
    binary_mask[(derivative_magnitude >= theta)] = 1

    return binary_mask


def find_edges_nms(image: Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]], sigma: float, theta: float) -> Union[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Aceppts: image, sigma & theta
    Returns: image with edges
    """
    step_size = np.pi/8


    def get_gradient_orientation(angle: float) -> tuple[tuple[int, int], tuple[int, int]]:
        """
        Accepts: angle
        Returns: indexes of gradient orientation (x, y), (x, y)
        Basically walks around the unit circle and returns the indexes of the closest angle
        """
        angle = angle % np.pi

        for i in range(0, 8):
            if angle >= i * step_size and angle <= (i+1) * step_size:
                if i == 0 or i == 7:
                    return (0, 1), (0, -1)
                elif i == 1 or i == 2:
                    return (-1, 1), (1, -1)
                elif i == 3 or i == 4:
                    return (-1, 0), (1, 0)
                elif i == 5 or i == 6:
                    return (-1, -1), (1, 1)
        raise ValueError(f"Angle {angle} is not in range")

    derivative_magnitude, derivative_angles = gradient_magnitude(image, sigma)

    reduced_magnitude = np.zeros_like(derivative_magnitude)
    nms_mask = np.zeros_like(derivative_magnitude)
    reduced_magnitude[(derivative_magnitude >= theta)] = 1


    for y in range(reduced_magnitude.shape[0]):
        for x in range(reduced_magnitude.shape[1]):
            gp1, gp2 = get_gradient_orientation(derivative_angles[y, x])

            # Out of bounds checks
            if x + gp1[0] < 0 or x + gp2[0] < 0:
                continue
            elif y + gp1[1] < 0 or y + gp2[1] < 0:
                continue
            elif x + gp1[0] >= reduced_magnitude.shape[1] or x + gp2[0] >= reduced_magnitude.shape[1]:
                continue
            elif y + gp1[1] >= reduced_magnitude.shape[0] or y + gp2[1] >= reduced_magnitude.shape[0]:
                continue
            elif reduced_magnitude[y + gp1[1], x + gp1[0]] == 1 and reduced_magnitude[y + gp2[1], x + gp2[0]] == 1:
                nms_mask[y, x] = 1 
 
    return nms_mask