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

# + colab={"base_uri": "https://localhost:8080/"} id="O-82JAxpWHys" outputId="f916b215-6a48-428d-b3f2-d3515c8b1782"
from google.colab import drive

drive.mount('/proj')

# + colab={"base_uri": "https://localhost:8080/"} id="lwpbOzCBWqjJ" outputId="b9358d00-0d28-4ac1-82cb-0c9a3cbadb01"
# !pip install -q -U umap-learn[plot] hdbscan tensorflow-addons #opencv-python==4.5.1.48

# + id="6OnwBWQ-XK38"
#Imports

import tensorflow as tf
import tensorflow_addons as tfa
import tensorflow_datasets as tfds

import io
import numpy as np

###UMAP seems to take a while to import due to a pynndescent dependency (which uses numba).
import umap
import umap.plot
import hdbscan
import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import random
import seaborn as sns
import shutil

from pathlib import Path
from sklearn.neighbors import KNeighborsClassifier

# + id="Hw-u4KdTX9Ns"
CHECKPOINT_DIR="/proj/MyDrive/158780_Project_Dolphin_Computer_Vision/models/triplet_loss_dolphin/training/"
SEED = 100
VAL_PROP = 0.15
DIM=128
CHANNELS=3
data_dir = "/content/final_pigmentation_catalogue_2016"
test_dir = "/content/test"

# + id="Tuf_PA-DEPKv"
shutil.rmtree(data_dir, ignore_errors=True)

# + colab={"base_uri": "https://localhost:8080/"} id="h-vLGLnhYE1s" outputId="064f349d-ede4-4d01-ee78-1d65ed1322fe"
# !unzip /proj/MyDrive/158780_Project_Dolphin_Computer_Vision/images/final_pigmentation_catalogue_2016.zip -d /content

# + [markdown] id="6pXlbfgD9Qry"
# We will build the stratified testing set with 10% of the dataset per class. As far as I'm aware, keras can't be used for this.

# + id="rTX8t16N6Z3L"
shutil.rmtree(test_dir, ignore_errors=True)
for dir in Path(data_dir).glob('*'):
  class_test_dir = f"/content/test/{dir.name}"
  Path(class_test_dir).mkdir(parents = True, exist_ok = True)
  images = list(dir.glob('*.png'))
  test_prop = round(len(images) * 0.1)
  for i in range(test_prop):
    test_file = random.choice(images)
    while os.path.exists(f"{class_test_dir}/{test_file.name}"): test_file = random.choice(images)
    shutil.move(str(test_file), class_test_dir)

# + colab={"base_uri": "https://localhost:8080/"} id="GwMjLmzFDNYz" outputId="af1027b1-d2d3-4aa6-9ce2-a2335ff23b7f"
# !ls /content/test/0004

# + [markdown] id="YW3T1LXbcqux"
# We'll use transfer learning with a ResNet50 backbone with an L2 normaliser and online image augmentation. We load ResNet50, inialised with imagenet weights. Critically, the final dense layer has no activiation function - we don't want to perform any non-linear transformations prior to nornalisation.
#
# We use a 2048 unit dense layer, using Leaky ReLu as its activation function. This is followed by batch normalisation for regularisation. Next, a 1024 unit dense layer with ,batch normalisation, a final 256 unit dense layer with no activation function, and L2 normalisation.

# + id="8FwIC8qmaVil"
batch_size=128
from tensorflow.keras.applications.resnet import ResNet50
base_model = ResNet50(include_top=False, pooling='avg', weights='imagenet')
base_model.trainable = False

model = tf.keras.Sequential([
    tf.keras.layers.experimental.preprocessing.RandomFlip("horizontal", input_shape=(DIM,DIM,CHANNELS)),
    tf.keras.layers.experimental.preprocessing.RandomRotation(0.15),
    tf.keras.layers.experimental.preprocessing.RandomZoom(0.15),
    tf.keras.layers.experimental.preprocessing.RandomTranslation((-.15, .15), (-.15, .15), fill_mode='reflect',interpolation='bilinear'),
    tf.keras.layers.experimental.preprocessing.Rescaling(1./255, input_shape=(DIM,DIM,CHANNELS)),
    base_model,
    tf.keras.layers.Dense(2048, activation=tf.keras.layers.LeakyReLU(alpha=0.3)),  
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dense(1024, activation=tf.keras.layers.LeakyReLU(alpha=0.3)),  
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dense(256, activation=None),  
    tf.keras.layers.Lambda(lambda x: tf.math.l2_normalize(x, axis=1))  
])

# + colab={"base_uri": "https://localhost:8080/"} id="0EPBfrWcodNq" outputId="f008d611-be04-4490-d9b3-674d71eade32"
model.build(input_shape=(DIM,DIM,CHANNELS))
model.summary()


# + [markdown] id="0Gh3dlFPcnvl"
# We load the dataset, visualise the classes, and apply both pre-fetching and caching for performance.

# + colab={"base_uri": "https://localhost:8080/"} id="tfUkMTBUwRJe" outputId="e409288c-3943-46ee-d872-d8d3a83ace51"
def build_data(batch_size):
  train_dataset = tf.keras.preprocessing.image_dataset_from_directory(
      data_dir,
      labels="inferred",
      label_mode="int",
      class_names=None,
      color_mode="rgb",
      batch_size=batch_size,
      image_size=(DIM, DIM),
      shuffle=True,
      seed=SEED,
      validation_split=VAL_PROP,
      subset="training",
      interpolation="bilinear",
      follow_links=False,
  )

  val_dataset = tf.keras.preprocessing.image_dataset_from_directory(
      data_dir,
      labels="inferred",
      label_mode="int",
      class_names=None,
      color_mode="rgb",
      batch_size=batch_size,
      image_size=(DIM, DIM),
      shuffle=True,
      seed=SEED,
      validation_split=VAL_PROP,
      subset="validation",
      follow_links=False,
  )

  full_dataset = tf.keras.preprocessing.image_dataset_from_directory(
      data_dir,
      labels="inferred",
      label_mode="int",
      color_mode="rgb",
      # batch_size=batch_size,
      image_size=(DIM, DIM),
      shuffle=True,
      seed=SEED,
  )

  test_dataset = tf.keras.preprocessing.image_dataset_from_directory(
      test_dir,
      labels="inferred",
      label_mode="int",
      color_mode="rgb",
      # batch_size=batch_size,
      image_size=(DIM, DIM),
      shuffle=True,
      seed=SEED,
  )
  return train_dataset, val_dataset, full_dataset, test_dataset

def prefetch_data(batch_size):
  train_dataset, val_dataset, full_dataset, test_dataset = build_data(batch_size)
  train_dataset = train_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
  val_dataset = val_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)
  full_dataset = full_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
  test_dataset = test_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
  return train_dataset, val_dataset, full_dataset, test_dataset

train_dataset, val_dataset, full_dataset, test_dataset = build_data(68)

# + colab={"base_uri": "https://localhost:8080/", "height": 399} id="VODR1haNwW7p" outputId="bbbbb036-636a-4c24-f8e2-11cac33bf9f5"
import matplotlib.pyplot as plt

class_names = train_dataset.class_names
plt.figure(figsize=(10, 10))
for images, labels in train_dataset.take(1):
  for i in range(6):
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(images[i].numpy().astype("uint8"))
    plt.title(class_names[labels[i]])
    plt.axis("off")

# + colab={"base_uri": "https://localhost:8080/"} id="h56CGG0Twnf_" outputId="6e3bf4d4-27c2-458c-baaa-b66b2ceba709"
train_dataset, val_dataset, full_dataset, test_dataset = prefetch_data(68)

# + [markdown] id="PFZpa-UOd3YI"
# We first perform online learning with triplet hard loss and a soft margin of 2.

# + id="rKE2IBbsKIG1"
model.compile(
  optimizer=tf.keras.optimizers.Adamax(0.001),
  loss=tfa.losses.TripletHardLoss(2, soft = True)
)


# + id="66jr6kj2RmFh"
def scheduler(epoch, lr):
  if epoch < 10:
    return lr
  else:
    return lr * tf.math.exp(-0.1)


# + colab={"base_uri": "https://localhost:8080/"} id="plHC1V4IgZyP" outputId="c4fd9d60-9240-4c49-bfc8-481f3bf83983"
history = model.fit(
  train_dataset,
  validation_data = val_dataset,
  epochs = 50,
  callbacks = [
    tf.keras.callbacks.LearningRateScheduler(scheduler),            
    tf.keras.callbacks.EarlyStopping(monitor='val_loss', mode='min', restore_best_weights = True, patience = 30),
    tf.keras.callbacks.ModelCheckpoint(f"{CHECKPOINT_DIR}/training_semi_resnet_aug", monitor= 'val_loss', save_best_only = True, mode = "min", save_weights_only = True),
    tf.keras.callbacks.ReduceLROnPlateau(monitor="val_loss", factor=0.1, patience=10, mode="min", min_delta=0.0001, min_lr = 0)
  ]
)


# + [markdown] id="wptw0En_2hyo"
# We visualise the training results.

# + colab={"base_uri": "https://localhost:8080/", "height": 281} id="GZCkagHdw7Bz" outputId="514465e9-5e5a-40e4-fee7-3b00400f02c8"
def plot_loss(history):
  epochs = len(history.history['loss'])
  loss = history.history['loss']
  val_loss = history.history['val_loss']
  epochs_range = range(epochs)
  plt.plot(epochs_range, loss, label='Training Loss')
  plt.plot(epochs_range, val_loss, label='Validation Loss')
  plt.legend(loc='upper right')
  plt.title('Training and Validation Loss')
  plt.show()

plot_loss(history)


# + [markdown] id="rmVG4Z5b2oNu"
# Not bad, but we consistently reach a loss plateau.

# + id="bDtc4cIZxpHi"
def fit_model(epochs, batch_size):
  train_dataset, val_dataset, full_dataset, test_dataset = prefetch_data(batch_size)
  model.load_weights(f"{CHECKPOINT_DIR}/training_semi_resnet_aug")
  history = model.fit(
    train_dataset,
    validation_data = val_dataset,
    epochs = epochs,
    callbacks = [
      tf.keras.callbacks.LearningRateScheduler(scheduler),            
      tf.keras.callbacks.EarlyStopping(monitor='val_loss', mode='min', restore_best_weights = True, patience = 30),
      tf.keras.callbacks.ModelCheckpoint(f"{CHECKPOINT_DIR}/training_semi_resnet_aug", monitor= 'val_loss', save_best_only = True, mode = "min", save_weights_only = True),
      tf.keras.callbacks.ReduceLROnPlateau(monitor="val_loss", factor=0.1, patience=10, mode="min", min_delta=0.0001, min_lr = 0)    
    ]
  )
  return history



# + colab={"base_uri": "https://localhost:8080/", "height": 281} id="s9kXhyiGyk_4" outputId="76ed1cc0-a194-4da3-b312-de2c5458d79e"
history = fit_model(50, 64)
plot_loss(history)


# + [markdown] id="tQT4PD8AKnjG"
# It seems we've hit a loss plateau. We'll try decreasing the learning rate.
#

# + id="I1wLzTSmWMEQ"
def compile_model(lr, loss):
  model.compile(
    optimizer=tf.keras.optimizers.Adamax(lr),
    loss=loss
  )


# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="4_6gS0v_NxJp" outputId="92ce2cdc-d67f-4c56-89dc-dd391937db7a"
compile_model(1e-07, tfa.losses.TripletHardLoss(2, soft = False))
history = fit_model(50, 128)
plot_loss(history)

# + [markdown] id="mlVWjG61y3qm"
# Again, we've reached a plateau.

# + [markdown] id="MET8tyKC1lkQ"
# We'll switch back to semi-hard loss

# + [markdown] id="9sCPjwwZxqyb"
#  So, we will generate the embeddings, visualise them, and use a KNN classifier.

# + colab={"base_uri": "https://localhost:8080/"} id="wAUoSwKyzfXn" outputId="0bbf8f14-1238-4149-f1b8-a96a0d68f79c"
train_embeddings = model.predict(full_dataset)
test_embeddings = model.predict(test_dataset)
train_embeddings.shape

# + [markdown] id="wv40JpwtNbpD"
# We retrieve the class labels.

# + colab={"base_uri": "https://localhost:8080/"} id="a7E8iN1eMHgV" outputId="594ebc23-ebc1-4220-aaf3-606f8a503dee"
y_train = np.concatenate([y for x, y in full_dataset], axis=0)
y_test = np.concatenate([y for x, y in test_dataset], axis=0)
y_train[0:100]

# + [markdown] id="wr-vr0t9Ty2c"
# Next, we plot the results.

# + colab={"base_uri": "https://localhost:8080/", "height": 295} id="Vymd_lLdN56D" outputId="4a4c672f-30ef-41c6-e817-28796b3a2d8c"
mapper=umap.UMAP( n_neighbors = 100, min_dist= 0.1).fit(train_embeddings, y_train)
cluster = hdbscan.HDBSCAN(
    algorithm ='best', 
    approx_min_span_tree = True,
    gen_min_span_tree = False, 
    leaf_size = 40, 
    metric='euclidean', 
    min_cluster_size = 25,
    min_samples = 25, 
    p = None
).fit(mapper.embedding_)

plt.title("Clustered UMAP visualisation.")
scatter = plt.scatter(
    mapper.embedding_[:, 0], 
    mapper.embedding_[:, 1],
    c = cluster.labels_,
    s = 0.1, 
)

plt.legend(
  handles = scatter.legend_elements()[0], 
  labels = pd.Series(cluster.labels_.tolist()).astype("string").sort_values().unique(),
  loc="right",
  bbox_to_anchor=(1.2, 0.5)
)
plt.xlabel("D1")
plt.ylabel("D2")
plt.show();

# + [markdown] id="kfbw22vJUOVF"
# Not so good. The clusters are not well formed.

# + colab={"base_uri": "https://localhost:8080/"} id="FZTTN8DgRoq9" outputId="c4e38923-6f86-41a0-ba6e-5f4190eee6cd"
knn=KNeighborsClassifier(100)
knn.fit(mapper.embedding_, y_train)

# + colab={"base_uri": "https://localhost:8080/"} id="XxGpCevpTfsi" outputId="8b06d47d-384c-4327-be32-94044ebf811a"
knn.score(mapper.transform(test_embeddings), y_test)

# + [markdown] id="3z0FykwMUTaJ"
# The KNN classifier has not done a good job.

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="asqbnozyUlzr" outputId="3fadf98c-c5c8-4062-81d2-4a8e63e9b1ac"
# %%shell
jupyter nbconvert --to html /content/4_3_triplet_loss_dolphin_resnet_aug.ipynb

StarcoderdataJupyterScriptsDedupFiltered