SBF_filter.py

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
import cv2
import sys
import math
import statistics
import random
import scipy.ndimage as ndimage
from IPython.display import Image
from nitorch.data import load_nifti


# In[2]:


def sp_noise(image, prob):
    '''
    Add salt and pepper noise to image
    prob: Probability of the noise
    '''
    output = np.zeros(image.shape,int)
    thres = 1 - prob 
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            rdn = random.random()
            if rdn < prob:
                output[i][j] = 0
            elif rdn > thres:
                output[i][j] = 255
            else:
                output[i][j] = image[i][j]
    return output


# In[3]:


def impulse_noise(image, prob): # add random valued impulse noise in grayscale 

    output = np.zeros(image.shape,int)

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):  
            rdn = random.random()
            if rdn < prob:
                output[i][j] = random.randint(0,255)
            else:
                output[i][j] = image[i][j]
    return output


# In[4]:


def gauss_noise(image, var, prob):
    mean = 0
    sigma = var**0.5
    output = np.zeros(image.shape, int)

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):  
            rdn = random.random()
            if rdn < prob:
                #print(image[i][j] + np.random.normal(mean, sigma))
                output[i][j] = (image[i][j] + np.random.normal(mean, sigma)).astype(int)
            else:
                output[i][j] = image[i][j]
                
    output = np.where(output>255,255,output)
    output = np.where(output<0,0,output)
            
    return output


# In[5]:


def median_calc(i,j,image,N):
    Q3 = image[(i-1):(i+N+1), (j-N-1):(j)]
    Q2 = image[(i-N-1):(i), (j-N-1):(j)]
    Q4 = image[(i-1):(i+N+1), (j-1):(j+N+1)]
    Q1 = image[(i-N-1):(i), (j-1):(j+N+1)]
    Q0 = image[(i-N-1):(i+N+1), (j-N-1):(j+N+1)]

    m0 = np.median(Q0)
    m1 = np.median(Q1)
    m2 = np.median(Q2)
    m3 = np.median(Q3)
    m4 = np.median(Q4)   
    
    return m0, m1, m2, m3, m4


# In[6]:


def find_dav(i, j, image, N):
    vert_dav = (image[i+N-1][j-1]+image[i][j-1]+image[i-N][j-1]+image[i-N-1][j-1])/4
    hor_dav = (image[i-1][j-N-1]+image[i-1][j-N]+image[i-1][j]+image[i-1][j+N-1])/4
    diag_dav = (image[i][j-1]+image[i-1][j]+image[i-1][j-N]+image[i-N][j-1])/4

    vert_diff = np.abs((image[i-1][j-1]-vert_dav)**2)
    hor_diff = np.abs((image[i-1][j-1]-hor_dav)**2)
    diag_diff = np.abs((image[i-1][j-1]-diag_dav)**2)

    if np.median(vert_diff) == min(np.median(vert_diff), np.median(hor_diff), np.median(diag_diff)):
        dav = vert_dav
    elif np.median(hor_diff) == min(np.median(vert_diff), np.median(hor_diff), np.median(diag_diff)):
        dav = hor_dav
    elif np.median(diag_diff) == min(np.median(vert_diff), np.median(hor_diff), np.median(diag_diff)):
        dav = diag_dav  

    return dav


# In[7]:


def reference_median(m2, m3, m4, dav, SQRM, SQMDC, pi_val):
    
    if (SQMDC<=pi_val): #weak edge
        SQRM = (m2+m3)/2
    elif (SQMDC>=pi_val) and ((m3 <= np.median(dav) <= m4) or (m4 <= np.median(dav) <= m3)): #no edge
        SQRM = m3
    else:
        SQRM = m2  
        
    return SQRM


# In[8]:


def noise_type(i,j,image,RM,Tk1, Tk2):    
    
    if abs(np.median(image[i-1][j-1])-RM)>=Tk1:
        s1 = 1
        s2 = 1
    elif abs(np.median(image[i-1][j-1])-RM)>=Tk2:
        s1 = 1
        s2 = 0
    else:
        s1 = 0
        s2 = 0
    return s1, s2


# In[9]:


def SBF(N, i, j, image, sigma_s, sigma_r, S1, S2, SQRM):
    
    u = 0
    
    if S1 == 0:
        u = image[i-1][j-1]
        
    else:
        
        nom = 0
        denom = 0
        I = 0
        W_G = 0
        W_SR = 0
        
        if S2 == 0:
            I = np.median(image[i-1][j-1])
        elif S2 == 1:
            I = SQRM

        for s in np.arange(-N, (N+1)):
            for t in np.arange(-N, (N+1)):

                W_G = np.exp(-np.sqrt((i+s)**2+(j+t)**2)/(2*sigma_s**2))
                W_SR = np.exp(-np.abs(I - np.median(image[i+s-1][j+t-1]))/(2*sigma_r**2))

                nom += (W_G*W_SR*image[i+s-1][j+t-1])
                denom += (W_G*W_SR)                   

        u = nom/denom
    
    return u 


# In[10]:


def SwitchingBilateralFilter(image, N=2, pi_val=50, Tk1=20, Tk2=5, sigma_r=25):
    
    r, c = image.shape
    exit_ft = image

    for i in range(N+1, r-N+1):
        for j in range(N+1, c-N+1):

            m0, m1, m2, m3, m4 = median_calc(i,j,image,N)

            dav = 0
            dav = find_dav(i, j, image, N)   

            SQRM = 0
            SQMDC = m3-m2

            SQRM = reference_median(m2, m3, m4, dav, SQRM, SQMDC, pi_val)

            S1, S2 =  noise_type(i,j,image,SQRM,Tk1, Tk2)               

            if (SQRM!=0):
                sigma_s = 3
            else:
                sigma_s = 1 

            u = SBF(N, i, j, image, sigma_s, sigma_r, S1, S2, SQRM)

            exit_ft[i-1][j-1] = u

    return exit_ft


# In[11]:


def SBFilter(image, N, pi_val, Tk1, Tk2, sigma_r):
    image = cv2.normalize(image, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
    r, c = image.shape
    exit_ft = image
    exit_ft = SwitchingBilateralFilter(image, N, pi_val, Tk1, Tk2, sigma_r)
    return exit_ft


# In[ ]: