# %% [markdown]
# # Histogram equalization
# [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/YoniChechik/AI_is_Math/blob/master/c_02a_basic_image_processing/histogram_equalization.ipynb)

# %%
# to run in google colab
import sys
if 'google.colab' in sys.modules:
    import subprocess
    subprocess.call('apt-get install subversion'.split())
    subprocess.call('svn export https://github.com/YoniChechik/AI_is_Math/trunk/c_02a_basic_image_processing/Unequalized_Hawkes_Bay_NZ.jpg'.split())

# %%
import numpy as np
import matplotlib.pyplot as plt
import cv2

figsize = (10,10)
# %% [markdown]
# First, read the image as grayscale

# %%
# read as grayscale
I = cv2.imread("Unequalized_Hawkes_Bay_NZ.jpg",0)

plt.figure(figsize=figsize)
plt.imshow(I, cmap='gray', vmin=0, vmax=255)
plt.title("Original image")
# %% [markdown]
# Let's start by calculating and showing the original histogram
# %%
bins_edges_min_max = [0,256]
num_bins=256
bin_count,bins_edges = np.histogram(I,num_bins,bins_edges_min_max)
bins_start = bins_edges[:-1]
# %%
def draw_hist(x_axis,input):
    fig,ax = plt.subplots(figsize=figsize)
    # why not using plt.hist? because we want to plot also some derivations of this hist, so this is easier
    plt.bar(x_axis, input, width=input.shape[0]/(x_axis[-1]-x_axis[0]+1))
    return fig,ax

draw_hist(bins_start,bin_count)
plt.title("Original histogram")

# %% [markdown]
# Normalize the histogram to gat a discrete PDF
# %%
pdf = bin_count/np.sum(bin_count)

draw_hist(bins_start,pdf)
plt.title("Original PDF")

# %% [markdown]
# Get the CDF by calculating the cumulative sum of the pdf data
# %%
cdf = np.cumsum(pdf)

plt.figure(figsize=figsize)
plt.plot(cdf)
plt.title("Original CDF")

# %%
fig,ax = draw_hist(bins_start,pdf)
ax.plot(cdf*np.max(pdf),'r')
plt.title("Original PDF+ const*CDF to show the connection between the two")

# %% [markdown]
# The final step is to un-normalize the CDF to become the equalization function
# %%
f_eq = np.round(cdf*255).astype(int)

f_eq
# %% [markdown]
# Use the equalization function to get the equalized image
# %%
I_eq = f_eq[I]

plt.figure(figsize=figsize)
plt.imshow(I_eq, cmap='gray', vmin=0, vmax=255)
plt.title("equalized image")
# %% [markdown]
# Plot the equalized histogram, PDF and CDF
# %%
bin_count,bins_edges = np.histogram(I_eq,num_bins,bins_edges_min_max)
bins_start = bins_edges[:-1]

draw_hist(bins_start,bin_count)
plt.title("equalized histogram")
# %%
pdf = bin_count/np.sum(bin_count)
cdf = np.cumsum(pdf)

fig,ax = draw_hist(bins_start,pdf)
ax.plot(cdf*np.max(pdf),'r')
plt.title("equalized PDF and const*CDF")
# %% [markdown]
# ## cv2 histogram equalization function
# %%
I_eq_cv2 = cv2.equalizeHist(I)

plt.figure(figsize=figsize)
plt.imshow(I_eq_cv2, cmap='gray', vmin=0, vmax=255)
plt.title("cv2.equalizeHist() result")

# %%
