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

# <hr style="height:2px;">
# 
# # Demo: Apply trained CARE model for denoising of *Tribolium castaneum*
# 
# This notebook demonstrates applying a CARE model for a 3D denoising task, assuming that training was already completed via [2_training.ipynb](2_training.ipynb).  
# The trained model is assumed to be located in the folder `models` with the name `my_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 Path, download_and_extract_zip_file, plot_some
from csbdeep.io import save_tiff_imagej_compatible
from csbdeep.models import CARE


# <hr style="height:2px;">
# 
# # Download example data
# 
# The example data (also for testing) should have been downloaded in [1_datagen.ipynb](1_datagen.ipynb).  
# Just in case, we will download it here again if it's not already present.

# In[2]:


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


# <hr style="height:2px;">
# 
# # Raw low-SNR image and associated high-SNR ground truth
# 
# Plot the test stack pair and define its image axes, which will be needed later for CARE prediction.

# In[3]:


y = imread('data/tribolium/test/GT/nGFP_0.1_0.2_0.5_20_14_late.tif')
x = imread('data/tribolium/test/low/nGFP_0.1_0.2_0.5_20_14_late.tif')

axes = 'ZYX'
print('image size =', x.shape)
print('image axes =', axes)

plt.figure(figsize=(16,10))
plot_some(np.stack([x,y]),
          title_list=[['low (maximum projection)','GT (maximum projection)']], 
          pmin=2,pmax=99.8);


# <hr style="height:2px;">
# 
# # CARE model
# 
# Load trained model (located in base directory `models` with name `my_model`) from disk.  
# The configuration was saved during training and is automatically loaded when `CARE` is initialized with `config=None`.

# In[4]:


model = CARE(config=None, name='my_model', basedir='models')


# ## Apply CARE network to raw image
# 
# Predict the restored image (image will be successively split into smaller tiles if there are memory issues).

# In[5]:


get_ipython().run_cell_magic('time', '', 'restored = model.predict(x, axes)\n')


# Alternatively, one can directly set `n_tiles` to avoid the time overhead from multiple retries in case of memory issues.
# 
# **Note**: *Out of memory* problems during `model.predict` can also indicate that the GPU is used by another process. In particular, shut down the training notebook before running the prediction (you may need to restart this notebook).

# In[6]:


get_ipython().run_cell_magic('time', '', 'restored = model.predict(x, axes, n_tiles=(1,4,4))\n')


# ## Save restored image
# 
# Save the restored image stack as a ImageJ-compatible TIFF image, i.e. the image can be opened in ImageJ/Fiji with correct axes semantics.

# In[7]:


Path('results').mkdir(exist_ok=True)
save_tiff_imagej_compatible('results/%s_nGFP_0.1_0.2_0.5_20_14_late.tif' % model.name, restored, axes)


# <hr style="height:2px;">
# 
# # Raw low/high-SNR image and denoised image via CARE network
# 
# Plot the test stack pair and the predicted restored stack (middle).

# In[8]:


plt.figure(figsize=(16,10))
plot_some(np.stack([x,restored,y]),
          title_list=[['low (maximum projection)','CARE (maximum projection)','GT (maximum projection)']], 
          pmin=2,pmax=99.8);