# ---
# 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
# ---

# + [markdown] id="mMlSeyeEJKrc"
# ## Example code to reproduce the results of *Developing a Victorious Strategy to the Second Strong Gravitational lensing Data Challenge* by Bom et al.
#
# This notebook is used with a subsample of the ISGLC dataset, using the HJY+VIS(repeated) configuration. The network will be retrained on a given number of images, using the same number for validation and the rest for testing.

# + id="F3r1lKG4JKrf"
import os
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import tensorflow as tf
from tensorflow import keras
import cv2
import sklearn.metrics as sk_metrics
import numpy as np
import matplotlib.pyplot as plt
import h5py
import pandas as pd
import scipy
from PIL import Image
from tensorflow_addons import optimizers

# + [markdown] id="JdlgGoMTJKrh"
# Include the RADAM optimizer to tf namespace so the model can be loaded

# + id="4FOfQ9CUJKrh"
tf.keras.optimizers.RectifiedAdam = optimizers.RectifiedAdam

# + [markdown] id="OV0v5vwtJKri"
# Set the parameter:
#
# - n_img: number of images available for training (same number for validation). Evaluation will be done on the rest
# - batch_size: batch size to feed model.fit. Must be smaller than n_img
# - n_epochs: number of epochs for training

# + id="Le9q_PjTJKri"
n_img = 150
batch_size = 50
n_epochs = 10


# + id="c7i9rvLXJKrj"
def plot_all_bands(img, ax=None):
    img_norm = np.zeros_like(img)
    for i in range(img.shape[-1]):
        norm = ((img[:,:,i:i+1] - img[:,:,i:i+1].min()) / np.ptp(img[:,:,i:i+1]))#*255
        #norm = cv2.normalize(img[:,:,i:i+1], np.zeros_like(img[:,:,i:i+1]), 0, 1.0, cv2.NORM_MINMAX)
        img_norm[:,:,i:i+1] = norm
    if ax == None:
        f, ax = plt.subplots()
    ax.imshow(img_norm, interpolation='nearest', cmap='gray')
    return None

def resize_imgs(imgs, newsize):
    #newsize is a square
    result = np.zeros((imgs.shape[0], newsize, newsize, imgs.shape[-1]))
    for i, img in enumerate(imgs):
        for j in range(imgs.shape[-1]):
            resized = cv2.resize(img[:,:,j], (newsize, newsize), interpolation = cv2.INTER_NEAREST)
            result[i,:,:,j] = resized
    return result
    
    
def combine_bands(imgs, alone=0):
    n_bands = imgs.shape[-1]
    assert alone < n_bands, 'chosen band must be between 0 and {}'.format(n_bands - 1)
    solitary = imgs[:,:,:,alone:alone+1]
    rest_ids = [iid for iid in range(n_bands) if iid != alone]
    #print(rest_ids)
    rest = np.zeros((imgs.shape[0], imgs.shape[1], imgs.shape[2], n_bands-1))
    for j, band in enumerate(rest_ids):
        rest[:,:,:,j] = imgs[:,:,:,band]
        

    return rest, solitary

def fbeta(rec, prec, beta2=0.001):
    return (1+beta2) * (rec*prec) / ((beta2*prec) + rec)

def normalize(arr):
    result = np.zeros_like(arr)
    if arr.shape[-1] > 1:
        for i in range(arr.shape[-1]):
            maxi = np.max(arr[:,:,:,i])
            mini = np.min(arr[:,:,:,i])
            new_arr = (arr[:,:,:,i] - mini) / np.ptp(arr[:,:,:,i])
            result[:,:,:,i] = new_arr
    else:
        maxi = np.max(arr[:,:,:,0])
        mini = np.min(arr[:,:,:,0])
        new_arr = (arr[:,:,:,:1] - mini) / np.ptp(arr[:,:,:,:1])
        result = new_arr
    return result

def create_input_data(data, alone=0, norm=True):
    hjy, vis = combine_bands(data, alone)
    hjy_res = resize_imgs(hjy, 66)
    vis_res = resize_imgs(vis, 200)
    if norm:
        hjy_res = normalize(hjy_res)
        vis_res = normalize(vis_res)
    pad = np.zeros_like(vis_res)
    vis_res = np.concatenate([vis_res, pad, pad], axis=-1)
    return hjy_res, vis_res


# + id="FVQIx_rNJKrk"
# %config Completer.use_jedi=False

# + [markdown] id="Khyo75Y_JKrl"
# Set the directories where the data and model are, load them. One-hot encode the labels since the I SGLC provided them as 0 or 1.

# + id="V-xSRxdnJKrl"
data_dir = '/tf/bernardo/EscolaCBPF/lensdata'
model_dir = '/tf/bernardo/Challenge2/data'

# + id="EjOO2cbrJKrm"
imgs = np.load(os.path.join(data_dir, 'imgs_challenge1_git.npy'))
is_lens = np.load(os.path.join(data_dir, 'labels_challenge1_git.npy'))
is_lens = keras.utils.to_categorical(is_lens)

# + [markdown] id="Yrf_rmZiJKrm"
# Prepare the data. Since the images are in four bands, we choose one of them to be the equivalent of VIS, while the rest are the equivalent to HJY. We found the best results when $r$ was alone, so we choose its index in `combine_bands` and `create_input_data`, as the images are $ugri$. `combine_bands` only combine the bands this way for plotting, while `create_input_data` also resizes them to the format of the II SGLC (200x200 and 66x66) and normalizes them.

# + id="ltl5DNdhJKrm"
hjy_plot, vis_plot = combine_bands(imgs, alone=2)
hjy, vis = create_input_data(imgs, alone=2)

# + [markdown] id="06ptT0qGJKrn"
# Select example lens and no lens for visualization

# + id="EjJCOgMqJKrn" outputId="c6186ef6-c779-4e20-f2d6-f89a747e1a0d"
lens = np.where(is_lens == 1)[0]
nolens = np.where(is_lens == 0)[0]

id_lens = np.random.choice(lens)
id_nolens = np.random.choice(nolens)

f, ax = plt.subplots(2,2,figsize=[10,10], gridspec_kw={'wspace':0.01, 'hspace':0.02})

ax[0,0].imshow(vis_plot[id_lens], interpolation='nearest', cmap='gray')
ax[0,0].set_axis_off()
ax[0,0].set_title('R', fontsize=20)
ax[0,0].annotate('LENS', (10,10), xytext=(-25,50), annotation_clip=False, fontsize=20, weight='bold')

plot_all_bands(hjy_plot[id_lens], ax=ax[0,1])
ax[0,1].set_axis_off()
ax[0,1].set_title('UGI', fontsize=20)

ax[1,0].imshow(vis_plot[id_nolens], interpolation='nearest', cmap='gray')
ax[1,0].set_axis_off()
ax[1,0].annotate('NO LENS', (10,10), xytext=(-40,50), annotation_clip=False, fontsize=20, weight='bold')

plot_all_bands(hjy_plot[id_nolens], ax=ax[1,1])
ax[1,1].set_axis_off()

plt.show()

# + [markdown] id="GJpfijcyJKrn"
# Load the model. Select random `n_img` from the dataset, split in train/validation/test

# + id="EeuI20bOJKro"
model = keras.models.load_model(os.path.join(model_dir, 'efn2_hjy_vis_newpre_repeated_aug_new_fold5'))

# + id="mWGtP5VlJKro"
ids_train = np.random.choice(hjy.shape[0], n_img*2, replace=False)
hjy_train = hjy[ids_train[:n_img]]
vis_train = vis[ids_train[:n_img]]
y_train = is_lens[ids_train[:n_img]]

hjy_val = hjy[ids_train[n_img:]]
vis_val = vis[ids_train[n_img:]]
y_val = is_lens[ids_train[n_img:]]

ids_test = np.array([i for i in range(hjy.shape[0]) if i not in ids_train])
hjy_test = hjy[ids_test]
vis_test = vis[ids_test]
y_test = is_lens[ids_test]

# + id="BCjNvh7gJKro" outputId="67f3fa73-dedb-48b1-bfc0-3dd5ae3341bd"
save_name = '/tf/bernardo/Challenge2/data/test_model'
checkpoint = keras.callbacks.ModelCheckpoint(save_name, monitor='val_loss', save_best_only=True, verbose=1)
history = model.fit([vis_train, hjy_train], y_train, batch_size=batch_size, epochs=n_epochs, 
                    validation_data=([vis_val, hjy_val], y_val), callbacks=[checkpoint])

# + [markdown] id="1eQ57T-uJKrp"
# Load the best model, with the lowest validation loss

# + id="F1n2lqMeJKrp" outputId="8137dcdd-d47e-4307-a8fc-2fbba469e535"
best_model = keras.models.load_model(save_name)

# + [markdown] id="R6b5fJoFJKrp"
# Predict on the test set and get the metrics

# + id="1Fm1_5dTJKrp"
preds = best_model.predict([vis_test, hjy_test])

# + id="BYKge4bnJKrp"
fpr, tpr, _ = sk_metrics.roc_curve(y_test[:,1], preds[:,1])
auroc = sk_metrics.roc_auc_score(y_test[:,1], preds[:,1])

precision, recall, _ = sk_metrics.precision_recall_curve(y_test[:,1], preds[:,1])
prauc = sk_metrics.average_precision_score(y_test[:,1], preds[:,1])

# + id="L9ZFK054JKrq" outputId="050f25f6-2674-4fc9-d201-ae390fbba018"
f, ax = plt.subplots(1,2, figsize=[14,6])

ax[0].plot(fpr, tpr, 'k-', label = 'AUC={:2.3f}'.format(auroc))
ax[0].set_xlabel('False Positive Rate')
ax[0].set_ylabel('True Positive Rate')
ax[0].legend(loc='lower right')

ax[1].plot(recall, precision, 'k-', label = 'AUC={:2.3f}'.format(prauc))
ax[1].set_xlabel('Recall')
ax[1].set_ylabel('Precision')
ax[1].legend(loc='lower left')

plt.tight_layout()
plt.show()

# + id="3WTf8k4OJKrq"


StarcoderdataJupyterScriptsDedupFiltered