Image blending via Laplacian pyramids

imblend.py

import cv2 as cv
import numpy as np

def build_gauss_pyr(img, n_layers):
    pyr = [img]
    for _ in range(n_layers):
        pyr.append(cv.pyrDown(pyr[-1]))
    return pyr

def build_lap_pyr(gauss_pyr):
    pyr = []
    reversed_gp = gauss_pyr[::-1]
    for small, original in zip(reversed_gp, reversed_gp[1:]):
        unsampled = cv.pyrUp(small)
        lap = cv.subtract(original, unsampled) 
        pyr.append(lap)
    return pyr

def imcombine(img1, img2, mask, clip_to_uint8=True):
    n_rows, n_cols, dts = img1.shape
    mask = mask.reshape(n_rows, n_cols, -1)
    result = img1 * mask + img2 * (1 - mask)
    if clip_to_uint8:
        result = result.clip(0, 255).astype(np.uint8)
    return result

def imblend(img1, img2, mask, n_layers=6):
    n_rows, n_cols, _ = img1.shape
    
    gauss_pyr1 = build_gauss_pyr(img1.astype(float), n_layers)
    lap_pyr1 = build_lap_pyr(gauss_pyr1)
    gauss_pyr2 = build_gauss_pyr(img2.astype(float), n_layers)
    lap_pyr2 = build_lap_pyr(gauss_pyr2)
    mask_pyr = build_gauss_pyr(mask, n_layers)
    
    combined_lap_pyr = []
    for lap1, lap2, submask in zip(lap_pyr1, lap_pyr2, mask_pyr[-2::-1]):
        combined = imcombine(lap1, lap2, submask, False)
        combined_lap_pyr.append(combined)

    smallest_mask = mask_pyr[-1]
    res = imcombine(gauss_pyr1[-1], gauss_pyr2[-1], smallest_mask, False)
    for combined_lap in combined_lap_pyr:
        res = cv.add(cv.pyrUp(res), combined_lap)

    return np.clip(res, 0, 255).astype(np.uint8)