stack.py

import numpy as np
import scipy as sp
from ipdb import set_trace as st
import skimage as ski
import utils
from matplotlib import pyplot as plt
import matplotlib as mpl
from common import *
from munch import Munch as M
from scipy import sparse
from scipy.interpolate import Rbf
import os
from scipy.fft import fftn, ifftn, set_workers
from tqdm.contrib.itertools import product
# from itertools import product
from tqdm import tqdm
from tqdm.contrib import tenumerate
import tifffile
import imageio as ii
from scipy import misc
import random

class Stacker:
    def __init__(
        self,
        work_dir,
        params = M(
            noise_mul = 32.0, # fudge factor (also try 4, 16, 64)
            patch_size = 32,
            cache_path = '.cache',
        ),
        custom_name = 'out',
        method = 'fuse',
    ):

        self.custom_name = custom_name
        self.method = method

        self.work_dir = work_dir
        for k in params:
            setattr(self, k, params[k])

    def fuse_clip(self, patches, iters=2, k=2):
        valid = np.ones_like(patches)

        for _ in range(iters):
            #compute the masked mean
            num = (valid*patches).sum(axis=0)
            denom = valid.sum(axis=0)+1E-9
            mu = num/denom

            diff = patches-mu[None]
            stds = (diff**2).sum(axis=0)**0.5/np.sqrt(patches.shape[0])
            valid = np.abs(diff) < stds[None]*k

            #print('valid mean is', valid.mean())

        return mu

    def fuse(self, patches):
        avg = patches.mean(axis = (0))
        dist = np.abs(patches-avg[None]).max(axis = (1,2))

        ffts = fftn(patches, axes = (1,2), norm='ortho')
        reffft = ffts[dist.argmin()] #reference patch

        num_frames = patches.shape[0]
        noise = (patches.std(axis=0)**2).mean()
        # fudge factors from hasinoff paper
        noise *= num_frames**2 / 8 * self.noise_mul
        
        fft_err = np.abs(reffft[None] - ffts)**2
        err_scores = fft_err / (fft_err + noise)
        # print('score', err_scores.min(), err_scores.mean(), err_scores.max())
        # if err_scores.max() > 0.5 and dist.max() < 2E-3:
        #     st()
        
        correction = err_scores * (reffft[None] - ffts)
        fused_fft = (ffts + correction).mean(axis = 0)

        fused_patch = ifftn(fused_fft, axes = (0,1), norm='ortho').real #imag components 0
        # if dist.max() < 2E-3:
        #     st()
        
        return fused_patch
            
    def __call__(self, paths):
        N = len(paths)
        H,W,C = np.load(paths[0]).shape

        PS = self.patch_size
        HPS = PS//2 #img size must be a multiple of this

        padH = np.ceil(H/HPS).astype(np.int)*HPS-H
        padW = np.ceil(W/HPS).astype(np.int)*HPS-W
        paddedH, paddedW = H+padH+PS, W+padW+PS
        
        imgs = np.memmap(self.cache_path, dtype=np.float32, mode='w+', shape=(N,C,paddedH,paddedW))
        for i, path in tenumerate(paths):
            imgs[i] = np.pad(
                np.load(path, mmap_mode='r').transpose((2,0,1)),
                ((0,0), (HPS, padH+HPS), (HPS, padW+HPS)),
                mode='edge'
            )
        
        out = np.zeros((C, paddedH, paddedW))

        xs, ys = np.meshgrid(np.arange(PS), np.arange(PS))
        foo = lambda z: 0.5 - 0.5*np.cos( 2 * np.pi * (z.astype(np.float32)+0.5)/PS)
        window_fn = foo(xs) * foo(ys)

        ####
        # from pathos.pools import ProcessPool
        # pool = ProcessPool(nodes=20)
        
        # patch_slices = [
        #     np.index_exp[
        #         i*HPS:(i+2)*HPS,
        #         j*HPS:(j+2)*HPS,
        #     ]
        #     for (i,j) in product(range(paddedH//HPS-1), range(paddedW//HPS-1))
        # ]
            
        # patchws = [
        #     imgs[np.index_exp[:,c]+patch_slice] * window_fn[None]
        #     for patch_slice in tqdm(patch_slices)
        #     for c in range(C)
        # ]

        # results = pool.map(self.fuse, patchws)

        # for i, patch_slice in tenumerate(patch_slices):
        #     for c in range(C):
        #         out[c][patch_slice] += results[i*3+c]
        
        for (i, j) in product(range(paddedH//HPS-1), range(paddedW//HPS-1)):
            patch_slice = np.index_exp[
                i*HPS:(i+2)*HPS,
                j*HPS:(j+2)*HPS,
            ]

            patchws = [imgs[np.index_exp[:,c]+patch_slice] * window_fn[None] for c in range(C)]

            if True:
                from pathos.pools import ProcessPool 
                pool = ProcessPool(nodes=3)
                results = pool.map(self.fuse, patchws)
            else:
                results = [getattr(self, self.method)(patchw) for patchw in patchws]

            for c in range(C):
                out[c][patch_slice] += results[c]
            
        out = out[:,HPS:H+HPS,HPS:W+HPS]
        out = np.clip(out, 0.0, 1.0).transpose(1,2,0)
        np.save(f'{self.work_dir}/stacked/out', out)
        tifffile.imwrite(f'{self.work_dir}/stacked/out.tiff', out)
        ii.imsave(f'{self.work_dir}/stacked/{self.custom_name}.jpg', np.round(out*255).astype(np.uint8), quality=100)
        os.remove(self.cache_path)

if __name__ == '__main__':
    #let's test it out...
    
    face_ = misc.face() #an RGB image
    N = 20
    sigma = 100
    sigma2 = 500
    
    Nblur = 4
    Nalign = 4

    #1. add misalignments
    #2. add blurring

    np.random.seed(0)
    
    paths = []
    print('generating fake data')
    for i in tqdm(range(N)):
        face = (
            face_.astype(np.float32) +
            np.random.randn(*face_.shape) * sigma +
            utils.blur(np.random.randn(*face_.shape), 7) * sigma2
        )

        if i < Nblur:
            face = utils.blur(face, [1,2,3,4][i])
        elif Nblur <= i < Nalign+Nblur:
            face = np.roll(face, [5,7,9,3][i-Nblur], axis = 0)
            face = np.roll(face, [2,-4,5,-6][i-Nblur], axis = 1)
        else:
            face = utils.blur(face, 0.5)
            L = 2
            face = np.roll(face, random.randint(-L,L), axis=0)
            face = np.roll(face, random.randint(-L,L), axis=1)

        #let's not clip actually...
        #face = np.clip(face, 0.0, 255.0)/255.0
        face = face/255.0
        
        path = f'debug/testimgs/{i:02d}.npy'
        np.save(path, face)
        ii.imsave(path.replace('npy', 'jpg'), np.round(np.clip(face,0,1)*255.0).astype(np.uint8), quality=98)
        paths.append(path)
    
    bar = Stacker('debug', method='fuse', custom_name='hasinoff32')
    bar.patch_size=64
    bar(paths)

    bar.custom_name='hasinoff8'
    bar.noise_mul=1.0
    bar(paths)
    
    bar.custom_name='hasinoff128'
    bar.noise_mul=10000.0
    bar(paths)
    
    
    foo = Stacker('debug', method='fuse_clip', custom_name='clipavg')
    foo.patch_size=256
    foo(paths)