############################################################################
#
#   Name:           pyramid_stacking.py
#   Python:         3.3.5 on win32
#   Autor:          Adrian Haas
#   
############################################################################
#  focus stacking with pyramid method
############################################################################
#
# Algorithm:
#
#   1. for each image create gauss and laplace pyramid
#
#   2. stack highest laplace value at each level to a new image
#
#   3. recreate with the new images pyramid the output image
#
#   The basic pyramid algorithm is under
#     https://opencv-python-tutroals.readthedocs.io/en/latest/
#        ->image processing in opencv -> image pyramids 
#
#
############################################################################

import cv2
import numpy as np

############################################################################

#read images

img_fn=["IMG_4428.JPG","IMG_4429.JPG","IMG_4430.JPG","IMG_4431.JPG",
        "IMG_4432.JPG","IMG_4433.JPG","IMG_4434.JPG","IMG_4435.JPG",
        "IMG_4436.JPG"]


img_list = [cv2.imread(fn) for fn in img_fn]

#pyriamid levels

deep=6

#function for generating pyramids

def fgauss(G,deep):

    gP = [G]

    for i in range(deep+1):

        G = cv2.pyrDown(G)
        gP.append(G)

    return gP

def flaplace(gP,deep):

    lP = [gP[deep]]

    for i in range(deep,0,-1):

        size = (gP[i-1].shape[1], gP[i-1].shape[0])
        GE = cv2.pyrUp(gP[i], dstsize = size)
        #laplacian negative (use int16 istead uint8)
        L=gP[i-1].astype(np.int16)-GE.astype(np.int16)
        lP.append(L)

    return lP


#create gauss pyramid image list
gp_list=[fgauss(img,deep) for img in img_list]

#create laplace pyramid image list
gl_list=[flaplace(gp,deep) for gp in gp_list]


#stack highest laplace value at each level
LS = []
for i in range(0,deep+1):
    ls=gl_list[0][i]

    for gl in gl_list:

        mask=abs(ls)<abs(gl[i])
        ls[mask]=gl[i][mask]

    LS.append(ls)

# now reconstruct

ls_ = LS[0]

for i in range(1,deep+1):

    size = (LS[i].shape[1], LS[i].shape[0])
    ls_ = cv2.pyrUp(ls_, dstsize = size)
    ls_=ls_+LS[i]

    #clip 0..255
    maskmax=ls_>255
    maskmin=ls_<0
    ls_[maskmax]=255
    ls_[maskmin]=0

#save output
cv2.imwrite("pyramid_stack.jpg",ls_.astype(np.uint8))