Untitled

from pickletools import uint8
import numpy as np
import cv2
import math

def show(image):
    cv2.imshow('image',image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def get_gs(window_size, sigma_s):
    '''return gs, size: (window_size, window_size)'''
    p = int(window_size/2)
    gs = np.zeros((window_size, window_size))
    for i in range(window_size):
        for j in range(i, window_size):
            result = np.exp( -((i-p)**2 + (j-p)**2) * 0.5 / (sigma_s**2) )
            gs[i][j] = result
            gs[j][i] = result
    return gs

def get_gr_value(p, q, gr_dict):
    p_red, p_green, p_blue = p[0], p[1], p[2]
    q_red, q_green, q_blue = q[0], q[1], q[2]
    gr = gr_dict[abs(p_red - q_red)] * gr_dict[abs(p_green -  q_green)] * gr_dict[abs(p_blue-q_blue)]
    return gr

class Joint_bilateral_filter(object):
    def __init__(self, sigma_s, sigma_r):
        self.sigma_r = sigma_r
        self.sigma_s = sigma_s
        self.wndw_size = 6*sigma_s+1
        self.pad_w = 3*sigma_s
    
    def joint_bilateral_filter(self, img, guidance):
        BORDER_TYPE = cv2.BORDER_REFLECT
        padded_img = cv2.copyMakeBorder(img, self.pad_w, self.pad_w, self.pad_w, self.pad_w, BORDER_TYPE).astype(np.int32)
        padded_guidance = cv2.copyMakeBorder(guidance, self.pad_w, self.pad_w, self.pad_w, self.pad_w, BORDER_TYPE).astype(np.int32)
        window_size = self.wndw_size
        
        
        ### TODO ###
        # save Gs kernel dict
        print("window size: ", window_size)
        print("padded image size: ", padded_img.shape)
        print("padded guidance size: ", padded_guidance.shape)
        gs = get_gs(window_size, self.sigma_s)
        print("gs shape:", gs.shape)
        print("=============")
        print("gr")
        gr_dict = [np.exp(-(x * 0.5 / ((self.sigma_r**2)* (255**2)))) for x in range(256)]
        
        
        output = np.zeros(img.shape)
        
        i, j = 0, 0
        for _ in range(padded_img.shape[0]):
            print(i, j)
            gr_value = get_gr_value(padded_guidance[i][j], padded_guidance[int(window_size/2)][int(window_size/2)], gr_dict)
            denominator = gs[i][j] * gr_value
            numerator = denominator * padded_guidance[i][j]
            output[i][j] = numerator/denominator
            j += 1
            if j == window_size:
                i += 1
                j = 0
                
        print("=============")
        print(output)
        return np.clip(output, 0, 255).astype(np.uint8)