#!/usr/bin/env python
# coding: utf-8

# <hr style="height:2px;">
# 
# # Demo: Training data generation for isotropic reconstruction of Zebrafish retina
# 
# This notebook demonstrates training data generation for an isotropic reconstruction task, where the anisotropic distortions along the undersampled Z axis are simulated from isotropic 2D slices.
# 
# Note that training data can be created from existing acquisitions.
# 
# We will use a single Retina stack for training data generation, whereas in your application you should aim to use stacks from different developmental timepoints to ensure a well trained model. 
# 
# More documentation is available at http://csbdeep.bioimagecomputing.com/doc/.

# In[1]:


from __future__ import print_function, unicode_literals, absolute_import, division
import numpy as np
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")

from tifffile import imread
from csbdeep.utils import download_and_extract_zip_file, plot_some, axes_dict
from csbdeep.io import save_training_data
from csbdeep.data import RawData, create_patches
from csbdeep.data.transform import anisotropic_distortions


# <hr style="height:2px;">
# 
# # Download example data
# 
# First we download some example data, consisting of a single 3D Zebrafish retina stack.

# In[2]:


download_and_extract_zip_file (
    url       = 'http://csbdeep.bioimagecomputing.com/example_data/retina.zip',
    targetdir = 'data',
)


# We plot XY and XZ slices of the training stack:

# In[3]:


x = imread('data/retina/cropped_farred_RFP_GFP_2109175_2color_sub_10.20.tif')
subsample = 10.2
print('image size         =', x.shape)
print('Z subsample factor =', subsample)

plt.figure(figsize=(16,15))
plot_some(np.moveaxis(x,1,-1)[[5,-5]],
          title_list=[['XY slice','XY slice']],
          pmin=2,pmax=99.8);

plt.figure(figsize=(16,15))
plot_some(np.moveaxis(np.moveaxis(x,1,-1)[:,[50,-50]],1,0),
          title_list=[['XZ slice','XZ slice']],
          pmin=2,pmax=99.8, aspect=subsample);


# <hr style="height:2px;">
# 
# # Generate training data for isotropic CARE
# 
# We first need to create a `RawData` object, which defines how to get pairs of images and the semantics of each axis (e.g. which one is considered a color channel, etc.).
# 
# In contrast to the standard CARE approach (e.g. [3D denoising](../denoising3D/1_datagen.ipynb)), we don't have pairs of low/high-SNR images here, just a single image.
# 
# Nevertheless, we can use `RawData.from_folder` and simply indicate the same folder as both source and target.  
# We also set `axes = 'ZCYX'` to indicate the semantic order of the image axes. 

# In[4]:


raw_data = RawData.from_folder (
    basepath    = 'data',
    source_dirs = ['retina'],
    target_dir  = 'retina',
    axes        = 'ZCYX',
)


# Furthermore, we must define how to modify XY slices to mimic the axial distortions of a real microscope as closely as possible. To that end, we define a `Transform` object that will take our `RawData` as input and return the modified image. Here, we use `anisotropic_distortions` to accomplish this.
# 
# The most important parameters are the subsampling factor along Z of the raw data and the *anisotropic part* $h_{aniso}(z,x)$ of the full PSF of the microscope $h_{full}(x,y,z)$. Specifically, $h_{aniso}(z,x)$ is the effective two-dimensional PSF with which lateral YX slices need to blurred such that they show the same image characteristics as an actual axial ZX slice. To find a correct $h_{aniso}$ for a given (e.g. measured) $h_full$ is in general a ill-posed deconvolution problem. In practice we find that using a simple gaussian approximation that uses the difference between the lateral and axial standard deviation ($\sigma_x$ and $\sigma_z$) of $h_{full}$ is often sufficcient (see function below).
# 
# More details can be found in our publication:
# Weigert et al. *Isotropic Reconstruction of 3D Fluorescence Microscopy Images Using Convolutional Neural Networks*. MICCAI 2017. https://doi.org/10.1007/978-3-319-66185-8_15
# 
# 

# In[5]:


def gaussian_anisotropic_psf(sigma_x, sigma_z):
    # create anisotropic psf based on lateral and axial standard deviation (in pixels) of the full PSF 
    _kx, _kz = int(4*sigma_x+1), int(4*sigma_z+1)
    _X, _Z = np.meshgrid(np.arange(-_kx,_kx+1), np.arange(-_kz,_kz+1), indexing='ij')
    return np.exp(-(_X**2/2/sigma_x**2+_Z**2/2/(sigma_z-sigma_x)**2))
    
psf = gaussian_anisotropic_psf(1, 3)
plt.imshow(psf) 


# In[6]:


anisotropic_transform = anisotropic_distortions (
    subsample     = 10.2,
    psf           = psf,
    poisson_noise = True,
    psf_axes      = 'YX',
)


# From the raw image stack and its synthetically distorted copy, we now generate corresponding patches. As a general rule, use a patch size that is a power of two along XYZT, or at least divisible by 8.  
# Typically, you should use more patches the more trainings stacks you have. By default, patches are sampled from non-background regions (i.e. that are above a relative threshold), see the documentation of `create_patches` for details.
# 
# Note that returned values `(X, Y, XY_axes)` by `create_patches` are not to be confused with the image axes X and Y.  
# By convention, the variable name `X` (or `x`) refers to an input variable for a machine learning model, whereas `Y` (or `y`) indicates an output variable.

# In[7]:


X, Y, XY_axes = create_patches (
    raw_data            = raw_data,
    patch_size          = (1,2,128,128),
    n_patches_per_image = 512,
    transforms          = [anisotropic_transform],
)


# In[8]:


assert X.shape == Y.shape
print("shape of X,Y =", X.shape)
print("axes  of X,Y =", XY_axes)


# Since the isotropic CARE model operates on 2D (+ channel) images, we need to remove the (singleton) Z axis before saving the training data.

# In[9]:


z = axes_dict(XY_axes)['Z']
X = np.take(X,0,axis=z)
Y = np.take(Y,0,axis=z)
XY_axes = XY_axes.replace('Z','')


# In[10]:


assert X.shape == Y.shape
print("shape of X,Y =", X.shape)
print("axes  of X,Y =", XY_axes)


# In[11]:


save_training_data('data/my_training_data.npz', X, Y, XY_axes)


# ## Show
# 
# This shows some of the generated patch pairs (odd rows: *source*, even rows: *target*)

# In[12]:


for i in range(2):
    plt.figure(figsize=(16,4))
    sl = slice(8*i, 8*(i+1))
    plot_some(np.moveaxis(X[sl],1,-1),np.moveaxis(Y[sl],1,-1),title_list=[np.arange(sl.start,sl.stop)])
    plt.show()
None;