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

# 训练集图像
img_paths = glob.glob('./../data/vesselseg/DRIVE/training/images/*.tif')

img_paths.sort()
img_paths

# 读取其中一张图像
rgb_img = Image.open(img_paths[0])

rgb_img.size

plt.imshow(rgb_img)

rgb_img = np.asarray(rgb_img)
print(rgb_img.shape)

# Blue 通道
plt.imshow(rgb_img[:,:,0], cmap="gray")

# Green 通道
plt.imshow(rgb_img[:,:,1], cmap="gray")

# Red 通道
plt.imshow(rgb_img[:,:,2], cmap="gray")

rgb_img = Image.open(img_paths[0])
print(rgb_img.size)
# Convert the dimension of imgs to [N,H,W,C]
rgb_img = np.expand_dims(rgb_img,0)
print(rgb_img.shape)
# Convert the dimension of imgs to [N,C,H,W]
rgb_img = np.transpose(rgb_img,(0,3,1,2))
print(rgb_img.shape)

#convert RGB image in black and white
def rgb2gray(rgb):
    assert (len(rgb.shape)==4)  #4D arrays
    assert (rgb.shape[1]==3)
    # 给 Green 通道对比度较高，给更大的权重
    bn_imgs = rgb[:,0,:,:]*0.299 + rgb[:,1,:,:]*0.587 + rgb[:,2,:,:]*0.114
    bn_imgs = np.reshape(bn_imgs,(rgb.shape[0],1,rgb.shape[2],rgb.shape[3]))
    return bn_imgs

gray_img = rgb2gray(rgb_img)

# 明明已经是单通道了，为什么还是彩色的？
plt.imshow(gray_img[0,0,:,:])

plt.imshow(gray_img[0,0,:,:], cmap="gray")

def dataset_normalized(imgs):
    assert (len(imgs.shape)==4)  #4D arrays
    assert (imgs.shape[1]==1)  #check the channel is 1
    imgs_normalized = np.empty(imgs.shape)
    imgs_std = np.std(imgs)
    imgs_mean = np.mean(imgs)
    imgs_normalized = (imgs-imgs_mean)/imgs_std
    for i in range(imgs.shape[0]):
        imgs_normalized[i] = ((imgs_normalized[i] - np.min(imgs_normalized[i])) / (np.max(imgs_normalized[i])-np.min(imgs_normalized[i])))*255
    return imgs_normalized

img_norm = dataset_normalized(gray_img)

img_norm.shape

plt.imshow(img_norm[0,0,:,:], cmap="gray")

# CLAHE (Contrast Limited Adaptive Histogram Equalization)
#adaptive histogram equalization is used. In this, image is divided into small blocks called "tiles" (tileSize is 8x8 by default in OpenCV). Then each of these blocks are histogram equalized as usual. So in a small area, histogram would confine to a small region (unless there is noise). If noise is there, it will be amplified. To avoid this, contrast limiting is applied. If any histogram bin is above the specified contrast limit (by default 40 in OpenCV), those pixels are clipped and distributed uniformly to other bins before applying histogram equalization. After equalization, to remove artifacts in tile borders, bilinear interpolation is applied
def clahe_equalized(imgs):
    assert (len(imgs.shape)==4)  #4D arrays
    assert (imgs.shape[1]==1)  #check the channel is 1
    #create a CLAHE object (Arguments are optional).
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    imgs_equalized = np.empty(imgs.shape)
    for i in range(imgs.shape[0]):
        imgs_equalized[i,0] = clahe.apply(np.array(imgs[i,0], dtype = np.uint8))
    return imgs_equalized

img_clahe = clahe_equalized(img_norm)

img_clahe.shape

plt.imshow(img_clahe[0,0,:,:], cmap="gray")

# laplacian kernle
K = [[0., 1., 0.],[1., -4., 1.], [0., 1., 0.]]

# Image sharpening by laplacian filter

def laplacian_sharpening(img, K_size=3):
    
    H, W = img.shape[2], img.shape[3]
    # zero padding
    pad = K_size // 2
    out = np.zeros((H + pad * 2, W + pad * 2), dtype=np.float)
    out[pad: pad + H, pad: pad + W] = img.copy().astype(np.float)
    tmp = out.copy()
    
    # laplacian kernle
    K = [[0., 1., 0.],[1., -4., 1.], [0., 1., 0.]]
    
    # filtering and adding image -> Sharpening image
    for y in range(H):
        for x in range(W):
            # core code
            out[pad + y, pad + x] = (-1) * np.sum(K * (tmp[y: y + K_size, x: x + K_size])) + tmp[pad + y, pad + x]

    out = np.clip(out, 0, 255)
    out = out[pad: pad + H, pad: pad + W].astype(np.uint8)

    return out

img_sharp = laplacian_sharpening(img_clahe)

img_sharp.shape

plt.imshow(img_sharp[:,:], cmap='gray')

def adjust_gamma(imgs, gamma=1.0):
    assert (len(imgs.shape)==4)  #4D arrays
    assert (imgs.shape[1]==1)  #check the channel is 1
    # build a lookup table mapping the pixel values [0, 255] to
    # their adjusted gamma values
    invGamma = 1.0 / gamma
    table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
    # apply gamma correction using the lookup table
    new_imgs = np.empty(imgs.shape)
    for i in range(imgs.shape[0]):
        new_imgs[i,0] = cv2.LUT(np.array(imgs[i,0], dtype = np.uint8), table)
    return new_imgs

img_gamma = adjust_gamma(img_clahe) 

plt.imshow(img_gamma[0,0,:,:], cmap='gray')

  
# Reading the image named 'input.jpg'
input_image = cv2.imread(img_paths[0])
input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2GRAY)

# Getting the kernel to be used in Top-Hat
filterSize =(3, 3)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, 
                                   filterSize)

# Applying the Top-Hat operation
tophat_img = cv2.morphologyEx(img_gamma[0,0,:,:], 
                              cv2.MORPH_TOPHAT,
                              kernel)

plt.imshow(img_gamma[0,0,:,:] - tophat_img, cmap='gray')