# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# This is the sequel of tutorial_training_dnn.ipynb.
#
# Here, we will discuss about loading a trained network, forwarding data to it and how to extract features maps.

## Set up the sys.path in order to be able to import our modules
import os
import sys
from matplotlib import pyplot as plt
module_path = os.path.abspath(os.path.join('../python'))
if module_path not in sys.path:
    sys.path.append(module_path)

# First, you can load a model using load_model from Keras to load an hdf5 file containing a network (saved using save_model or a ModelCheckpoint also defined in Keras).
#
# Otherwise, using the DeepNet class, you can call the constructor with the file.

# +
# Input the hdf5 file, here we load a network trained on Tikhonov
net_file = 'Deconv_net_tuto_copy.hdf5'
### Method 1
#from keras.models import load_model
#dnn = load_model(net_file) #dnn is instance of the model class

### Method 2
from DeepDeconv.deepnetFCS.DeepNet import DeepNet
dnn = DeepNet(model_file=net_file) #dnn is instance of DeepNet class
# -

dnn_alexis=DeepNet(model_file="../nets/DeconvNet2D_FCS_sc4_layer4x5x6x7_relu_growthRate12_reshfl_SNR20to100.hdf5")

dnn.model.summary()

# We also load a traditionnal CNN network

from keras.models import load_model
kk=load_model('/data/DeepDeconv/model/archive/unet_vs_cnn/cnn.hdf5')

# We then print the structure of the network:

kk.summary()

# And display the structure of the network:

# +
from keras.utils import plot_model
plot_model(kk, to_file='{0}.png'.format("CNN_10_layers"),show_shapes=True)
from IPython.display import SVG
from keras.utils.vis_utils import model_to_dot

pydot_obj=model_to_dot(kk,show_shapes=True,show_layer_names=False)
SVG(pydot_obj.create(prog='dot', format='svg'))
# -

# To forward data through the network use the predict method available for both classes. For the method 2:

# +
from DeepDeconv.utils.batch_utils import get_batch_from_fits
import numpy as np

# Input the file containing the galaxies and psfs for testing
testset_file = '/data/DeepDeconv/data/vsc_euclidpsfs/reshuffle/image-shfl-0-multihdu.fits'
noiseless_img_hdu = 0
psf_hdu = 1
targets_hdu = 2

# Create the set of test with 10 observations at SNR 50
test_data, target_data = get_batch_from_fits(testset_file, idx_list=np.arange(10), SNR=50,
                        noiseless_img_hdu=noiseless_img_hdu, targets_hdu=targets_hdu, psf_hdu=psf_hdu,
                        image_dim=96, image_per_row=100,
                        deconv_mode='TIKHONOV')
noisy_data, target_data = get_batch_from_fits(testset_file, idx_list=np.arange(10), SNR=50,
                        noiseless_img_hdu=noiseless_img_hdu, targets_hdu=targets_hdu, psf_hdu=psf_hdu,
                        image_dim=96, image_per_row=100,
                        deconv_mode=None)

dnn_reconstruction = dnn.predict(test_data, verbose=1)
print(dnn_reconstruction.shape)
# -

# Using DNN from Alexis:

dnn_reconstruction_alexis = dnn_alexis.predict(test_data, verbose=1)
print(dnn_reconstruction.shape)

plt.figure(figsize=(12,8))
plt.subplot(221),plt.imshow(target_data[0,:,:,0])
plt.subplot(222),plt.imshow(noisy_data[0,:,:,0])
plt.subplot(223),plt.imshow(dnn_reconstruction_alexis[0,:,:,0])
plt.subplot(224),plt.imshow(dnn_reconstruction[0,:,:,0])

# The DeepNet class also provides a method to extract intermediate features maps:
#
#     get_layer_output(test_data, layer_idx)
#     
# layer_idx (int): idx of the layer whose features maps will be extracted.
#
# Use dnn.model.summary() to retrieve the layers list.

dnn.model.summary()

# +
# Suppose we want the last concatenation layer
idx = -4

# We print the name to check if you have the correct layer
print(dnn.model.layers[idx].name)

# We extract the features maps for the same data set as above.
# The layer has 48 neurons so we get an output of shape (10,96,96,48)
feat_maps = dnn.get_layer_output(test_data, layer_idx=idx)
print(feat_maps.shape)

# +
plt.figure(figsize=(16,16))
fig,axes=plt.subplots(4,4)
for kx in range(4):
    for ky in range(4):
        axes[kx, ky].imshow(feat_maps[kx,:,:,ky])


# -



StarcoderdataJupyterScriptsDedupFiltered