############################################################################
#
#   Name:           filter_stacking.py
#   Python:         3.3.5 on win32
#   Autor:          Adrian Haas
#   
############################################################################
#  focus stacking with filter method
############################################################################
#
# Algorithm:
#
#   1. for each image create a filtered image:
#
#       a.  gaussian blur
#       b.  laplacian
#       c.  dilate
#
#   2. Map each pixel from image with the highest filtered image value.
#
############################################################################

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]

#create maximum map

x=img_list[0].shape[0]
y=img_list[0].shape[1]
points=np.zeros(x*y).reshape(x,y)


#filter each image

i=1

for row_img in img_list:

    gray = cv2.cvtColor(row_img,cv2.COLOR_BGR2GRAY)

    #gaussian
    dst=cv2.blur(gray,(3,3))

    #laplacian
    laplacian = abs(cv2.Laplacian(dst,cv2.CV_64F,ksize=3))

    #dilation
    kernel = np.ones((5,5),np.uint8)
    laplacian = cv2.dilate(laplacian,kernel,iterations = 3)

    #fill in maximum map

    if i==1:
        points[:,:]=i
        laplacian_temp=laplacian.copy()

    else:

        mask=laplacian_temp<laplacian
        laplacian_temp[mask]=laplacian[mask]
        points[mask]=i

    i+=1

#map all pictures to first pic

im_out=img_list[0].copy()
i=1
while i<len(img_list):
    mask=points==i+1
    im_out[mask]=img_list[i][mask]
    i+=1

#save output
cv2.imwrite('filter_stack.jpg',im_out)