Batch Input Operation Question

[thudepwcp13]

Hi, I try your code recently and it works well and fast. Thanks!
However, I really need to do some simple deconvolution process for thousands of times. This means I need some like batch input and output operation in Tensorflow and GPU rather than simple for loop. So could you give me some advice to add this function in your code?
Thanks again! @armish @hammer @eczech @russellb @eric-czech

[thudepwcp13]

I found you have implemented the run_batch function in restoration.py but it's also by simple loop. Why can not change the input dim to [batch_size, [None] * self.n_dims] ?
def _run_batch(self, acquisition_batch, **kwargs):
return [self._run(acq, **kwargs) for acq in acquisition_batch]

    datah = self._wrap_input(tf.placeholder(self.dtype, shape=[None] * self.n_dims, name='data'))
    kernh = self._wrap_input(tf.placeholder(self.dtype, shape=[None] * self.n_dims, name='kernel'))

[eric-czech]

Hey @thudepwcp13 , I like that idea for sure and would love to support a batch dimension. I just pushed a minor change to get things moving on this. Here's an example, assuming you pull the latest:

%matplotlib inline
from flowdec.nb import utils as fd_nb_utils
from flowdec import restoration as fd_restoration
from flowdec import data as fd_data

# NOTE: pad_mode = 'none' is crucial here -- batching doesn't work with padding yet
algo = fd_restoration.RichardsonLucyDeconvolver(4, pad_mode='none').initialize()

# Load a single 3D volume and repeat its image/PSF 3 times as if it was a batch of different images
acq = fd_data.bars_25pct()
acq = fd_data.Acquisition(data=np.stack([acq.data]*3), kernel=[acq.kernel]*3)

# Deconvolve the 4D array
res = algo.run(acq, niter=25)

print(res.data.shape)
# > (3, 32, 64, 64)

# Plot the first deconvolved volume
# fd_nb_utils.plot_rotations(res.data[0])

This is still very rough but it should be possible to deconvolve multiple images at once this way. Remaining things to take care of are:

1. Make padding work with a batch dimension
2. Make it possible to determine whether or not a batch dimension is present (or assume it's always there) for determination of the appropriate TF spectral functions
3. Check that this works for batches of 1 or 2D signals
4. Add tests to ensure that the batched images are deconvolved identically to those in isolation
5. Rethink the "Acquistion" wrapper since it's awkward with batched contents

I haven't done any validation on results of batch inputs so take that with a grain of salt, but I thought I'd share that much in case you want to give it a shot -- it might be a while before I get a chance to add that feature properly. Let me know if you have any success with that though or if you have other suggestions.

[naureeng]

I'm trying to set up batch deconvolution as well. I have a series of 2 color channel composite stacks that I would like to deconvolve with 2 different PSFs. Here is what I have so far based on previous answers. My question is how can I use the 2 different PSFs and automate saving the output ?

``# deconvolution by flowdec

from flowdec import data as fd_data
from flowdec import restoration as fd_restoration
from skimage import io
import numpy as np
import glob
import skimage

image_list = []
for filename in glob.glob('/nfs/winstor/isogai/stella/processed/180712Merfish_3D/*.tif'): #tif files
    img = np.stack([ skimage.io.imread(f) for f in filename ])
    psf_561 = io.imread('/flowdec/datasets/PSF561_38z.tif')
    psf_647 = io.imread('/flowdec/datasets/PSF647_38z.tif')
    assert img.ndim == 3
    assert psf_647.ndim == 3
    algo = fd_restoration.RichardsonLucyDeconvolver(3,pad_mode='none').initialize()
    res = np.stack([ 
    algo.run(fd_data.Acquisition(img[..., i], psf_647[..., i])).data 
    for i in range(img.shape[-1])
])

    

[eric-czech]

Hi @naureeng , I'm not quite sure how your images are stored in the tif files but it looks like the skimage.io.imread(f) for f in filename bit isn't what you want since that should try to load an image for each character in the filename. This is a bit of guesswork that would still depend on whether or not the channels are in the same file or separate files as well as what order they will load in, but it should be helpful:

import tqdm
wavelengths = [561, 647]
path = '/nfs/winstor/isogai/stella/processed/180712Merfish_3D/*.tif'
psfs = [io.imread('/flowdec/datasets/PSF%s_38z.tif' % w) for w in wavelengths]
algo = fd_restoration.RichardsonLucyDeconvolver(3, pad_mode='none').initialize()
images = {}

# Using tqdm to report progress
for filename in tqdm.tqdm(glob.glob(path)):
    # Assume each image would load as [z, y, x, channels] (which will
    # depend on how the files were written) and that there are two
    # channels (first at 561nm, second at 657nm)
    img = skimage.io.imread(filename)
    assert img.ndim == 4 
    assert img.shape[-1] == 2 # Make sure there are two channels

    res = [ 
      algo.run(fd_data.Acquisition(img[..., i], psfs[i])).data 
      for i in range(img.shape[-1])
    ]
    
    # Create new axis at end so that result is also [z, y, x, channels], 
    # not [channels, z, y, x]
    res = np.stack(res, -1)
    
    # Use imsave to store res somewhere 
    # skimage.io.imsave('/some/path/image_deconvolved.tif', res)

If nothing else, I would definitely suggest that you move the fd_restoration.RichardsonLucyDeconvolver(3,pad_mode='none').initialize() part outside of the loop. That will speed things up a good as will not loading the psfs for every image.

If that's slow, then we could talk about batching in the sense intended by the OP (i.e. as a performance optimization rather than control flow).

[naureeng]

Thank you so much for your help Eric! The current output is a 2 x 2 square of 2046 x 2046 pixels. However, 1 of those 4 squares should be 2046 x 2046. How would I be able to save each one as a separate image rather than "image_deconvolved.tif" ? My output is 2 x 2, representing the 2 channels of both 2 color images
``

# -*- coding: utf-8 -*-
from flowdec import data as fd_data
from flowdec import restoration as fd_restoration
from skimage import io
import tqdm
import skimage
import os
import glob
import numpy as np

wavelengths = [561, 647]
path = '/nfs/nhome/live/naureeng/flowdec/datasets/test/*.tif'
psfs = [io.imread('PSF%s_30z.tif' % w) for w in wavelengths]
algo = fd_restoration.RichardsonLucyDeconvolver(3, pad_mode='none').initialize()
images = {}

# Using tqdm to report progress
for filename in tqdm.tqdm(glob.glob(path)):
    # Assume each image would load as [z, x, y, channels] and that there are two
    # channels (first at 561 nm, second at 657 nm)
    img = skimage.io.imread(filename)
    [z, channels, x, y] = img.shape
    img_final = np.reshape(img, (z, x, y, channels))
    assert img_final.ndim == 4
    assert img_final.shape[-1] == 2 # Make sure there are 2 channels
    
    res = [
            algo.run(fd_data.Acquisition(img_final[..., i], psfs[i]), niter=25).data
            for i in range(img_final.shape[-1])
            ]
    
    # Create new axis at the end so that result is also [z,y,x, channels],
    # not [channels, z,y,x]
    res = np.stack(res, -1)
  
    # Use imsave to store res
    skimage.io.imsave('/nfs/nhome/live/naureeng/flowdec/datasets/test/image_deconvolved.tif', res)

[eric-czech]

Np @naureeng, are you sure the image is loading as [z, channels, x, y] (where x = columns and y = rows)? The first step to making it work will be figuring that out. But assuming that this is truly the shape of the image, then you don't want to use reshape like that (it is probably reordering the data in some non-sensical way). Instead you want:

# Move channels axis (1) to last axis (-1) to give result shape [z, x, y, channels]
img_final = np.moveaxis(img, 1, -1) 
# Alternatively this will accomplish the same thing in a more general way: 
# img_final = np.transpose(img, (0, 2, 3, 1))

To save each channel separately:

for i in range(res.shape[-1]): # Loop over channels
    res_path = '/nfs/nhome/live/naureeng/flowdec/datasets/test/image_deconvolved_ch{}.tif'.format(i+1)
    skimage.io.imsave(res_path , res[..., i])

[naureeng]

Thank you so much! I modified it to prevent over-writing of the 2 image channels:

# Python3 code to convert tuple  
# into string 
def convertTuple(tup): 
    str =  ''.join(tup) 
    return str
 filename_deconv = os.path.splitext(filename)
    final_name = filename_deconv[0] + "_D{}.tif",
    final_name_str = convertTuple(final_name)
  
    for i in range(res.shape[-1]): # Loop over channels
        res_path = final_name_str.format(i+1)
        skimage.io.imsave(res_path , res[..., i])

[eric-czech]

Sounds like you're running that outside of the first loop. This is what you want:

for filename in tqdm.tqdm(glob.glob(path)):
    # stuff
    for i in range(res.shape[-1]): 
        # save files

not:

for filename in tqdm.tqdm(glob.glob(path)):
    # stuff
for i in range(res.shape[-1]): 
    # save files