# ---
# 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] tags=[]
# # Dog Breed Classification
#
# This example is based on a very popular [Udacity project](https://github.com/udacity/dog-project), upgraded to use TensorFlow `2.3.0` with GPU accelaration. The goal is to classify images of dogs according to their breed.
#
# In this notebook, we take the first steps towards developing an algorithm that could be used as part of a mobile or web app. At the end of this project, our code will accept any user-supplied image as input. If a dog is detected in the image, it will provide an estimate of the dog's breed. In this real-world setting, we will piece together a series of state-of-the-art computer vision models to perform different tasks (Dog detection -> Breed classification).

# + [markdown] tags=[]
# ## Table of contents
#
# We break the notebook into separate steps. Feel free to use the links below to navigate the notebook.
#
# * [Step 0](#step0): Download Datasets and Install Dependencies
# * [Step 1](#step1): Import Datasets
# * [Step 2](#step2): Detect Dogs
# * [Step 3](#step3): Create a CNN to Classify Dog Breeds (from Scratch)
# * [Step 4](#step4): Create a CNN (VGG16) to Classify Dog Breeds (using Transfer Learning)
# * [Step 5](#step5): Create a CNN (ResNet-50) to Classify Dog Breeds (using Transfer Learning)
# * [Step 6](#step6): Write your Own Dog Classifier
# * [Step 7](#step7): Test Your Classifier

# + [markdown] colab_type="text" id="VvAyLIR-R0af" tags=[]
# <a id='step0'></a>
# ## Step 0: Download Datasets and Install Dependencies
#
# For this task we use TensorFlow `2.3.0`, as well as a few helper libraries like `Pillow`. Also, we need to download and extract the dataset we will use to train our classifier.

# + [markdown] colab_type="text" id="Fm1cOPKZSJQB" tags=[]
# ### Download the dataset
#
# The dataset is available [here](https://s3-us-west-1.amazonaws.com/udacity-aind/dog-project/dogImages.zip). The following cell downloads it as a zip file, extracts it and moves it to the corresponding folder. Finally, it removes the `zip` file.

# + colab={} colab_type="code" id="5MOpDJZCQGRb" tags=["skip"]
# !wget https://s3-us-west-1.amazonaws.com/udacity-aind/dog-project/dogImages.zip
# !unzip -qo dogImages.zip
# !rm dogImages.zip

# + [markdown] colab_type="text" id="q3iHERkNSNY-" tags=[]
# ### Install dependencies
#
# The task requires `Pillow` the friendly PIL fork by [<NAME> and Contributors](https://github.com/python-pillow/Pillow/graphs/contributors). PIL is the Python Imaging Library by <NAME> and Contributors. Below, we install the `Pillow` package using `pip`.

# + colab={} colab_type="code" id="WU7JZLVSRyAT" tags=["skip"]
# !pip3 install --user -r requirements/requirements-v2.txt

# + [markdown] colab_type="text" id="67p9fU72SQE9" tags=[]
# ### Import the necessary libraries
#
# We use Tensorflow `2.3.0` to build and train our dog breed classifier. We also need `Pillow` to load the images in memory, which we specifically instruct to load any truncated images also.

# + colab={} colab_type="code" id="7X11SD82SUhb" tags=["imports"]
import os
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

from tensorflow.keras.preprocessing.image import ImageDataGenerator

from glob import glob
from PIL import Image
from PIL import ImageFile


ImageFile.LOAD_TRUNCATED_IMAGES = True

# + colab={"base_uri": "https://localhost:8080/", "height": 34} colab_type="code" id="DKq8rHg5SWXq" outputId="7056f7a7-d327-4622-a59a-32b2108977f4" tags=["skip"]
print(f"Version of TensorFlow in use: {tf.__version__}")
print(f"Using GPU device: {tf.config.list_physical_devices('GPU')}")

# + [markdown] colab_type="text" id="ac4gM-_KSkV4" tags=[]
# <a id='step1'></a>
# ## Step 1: Import Datasets
#
# First, let us define the pipeline-parameters cell. We use it to define the hyperparametes we would like to tune later. These variables will be converted to KFP pipeline parameters, so we should make sure they are used as global variables throughout the notebook.

# + colab={} colab_type="code" id="ifMA-P1ySYK2" tags=["pipeline-parameters"]
LR = 6e-4
BATCH_SIZE = 32
NUMBER_OF_NODES = 256
EPOCHS = 4
IMG_SIZE = 224


# + [markdown] colab_type="text" id="EsWsiKbhTN4E" tags=[]
# ### Processing the Dataset
#
# We use TensorFlow native generators to load and transform the data. Pay attention to the `train_datagen` which also includes several transformations to augment our dataset (e.g. width and height shift, brightness alterations and horizontal flip). These transformations are taking place in memory, leaving the original data untouched.

# + tags=["block:data_processing"]
def get_train_generator():
    data_datagen = ImageDataGenerator(
        rescale=1./255,
        width_shift_range=.2,
        height_shift_range=.2,
        brightness_range=[0.5,1.5],
        horizontal_flip=True
    )
    return data_datagen.flow_from_directory(
        "dogImages/train/",
        target_size=(IMG_SIZE, IMG_SIZE),
        batch_size=BATCH_SIZE,
        
    )

def get_valid_generator():
    data_datagen = ImageDataGenerator(rescale=1./255)
    return data_datagen.flow_from_directory(
        "dogImages/valid/",
        target_size=(IMG_SIZE, IMG_SIZE),
        batch_size=BATCH_SIZE
    )

def get_test_generator():
    data_datagen = ImageDataGenerator(rescale=1./255)
    return data_datagen.flow_from_directory(
        "dogImages/test/",
        target_size=(IMG_SIZE, IMG_SIZE),
        batch_size=BATCH_SIZE
    )


# + [markdown] colab_type="text" id="r021ns1ZUzi4" tags=[]
# <a id='step2'></a>
# ## Step 2: Detect Dogs
#
# In this section, we use a pre-trained [ResNet V2](https://link.springer.com/chapter/10.1007/978-3-319-46493-0_38) model to detect dogs in images. First, we download a pretrained ResNet-50 model on [ImageNet](http://www.image-net.org/), a very large, very popular dataset used for image classification and other computer vision tasks. ImageNet contains over 10 million URLs, each linking to an image containing an object from one of 1000 categories. Given an image, this pre-trained ResNet-50 model returns a prediction (derived from the available categories in ImageNet) for the object that is contained in the image.

# + colab={} colab_type="code" id="PgWSAhN-TFNf" tags=["block:dog_detector", "prev:data_processing", "limit:nvidia.com/gpu:1"]
dog_classifier = tf.keras.applications.ResNet50V2(
    weights="imagenet",
    input_shape=(IMG_SIZE, IMG_SIZE, 3)
)


# + [markdown] tags=[]
# The categories corresponding to dogs appear in an uninterrupted sequence referring to dictionary keys 151-268, inclusive, to include all categories from 'Chihuahua' to 'Mexican hairless'. Thus, in order to check if an image is predicted to contain a dog by the pre-trained ResNet model, we need only check if the function below returns a value between 151 and 268 (inclusive).
#
# We use these ideas to complete the `is_dog` function below, which returns True if a dog is detected in an image.

# + colab={} colab_type="code" id="lguOZx_PWE73" tags=[]
def is_dog(data):
    probs = dog_classifier.predict(data)
    preds = tf.argmax(probs, axis=1)
    return ((preds >= 151) & (preds <= 268))


# + [markdown] tags=[]
# To test the classifier we request a batch from our training data generator and feed it through the network. The accuracy, as expected, is really high.

# + colab={} colab_type="code" id="7AkKSHeaYIcg" tags=[]
train_generator = get_train_generator()
batch = train_generator.next()
predictions = is_dog(batch)

# + colab={"base_uri": "https://localhost:8080/", "height": 34} colab_type="code" id="wGeUt1jidjat" outputId="19f96954-ed31-4ab5-917a-53178f696e10" tags=[]
n_dog = np.sum(predictions)
dog_percentage = n_dog/BATCH_SIZE

print('{:.0%} of the files have a detected dog'.format(dog_percentage))

# + [markdown] colab_type="text" id="8WJ6MCYrf4sR" tags=[]
# <a id='step3'></a>
# ## Step 3: Create a CNN to Classify Dog Breeds (from Scratch)
#
# Now that we have a function for detecting dogs in images, we need a way to predict the dog breed from images. In this step, you will create a CNN that classifies dog breeds. We should be careful with adding too many trainable layers though. More parameters means longer training and you may also fall in the trap of overfitting. Thankfully, `tf.keras` provides a handy estimate of the time that each epoch is likely to take; you can extrapolate this estimate to figure out how long it will take for your algorithm to train.
#
# We mention that the task of assigning breed to dogs from images is considered exceptionally challenging. To see why, consider that *even a human* would have great difficulty in distinguishing between a Brittany and a Welsh Springer Spaniel.  
#
# Brittany | Welsh Springer Spaniel
# - | - 
# <img src="images/Brittany_02625.jpg" width="100"> | <img src="images/Welsh_springer_spaniel_08203.jpg" width="200">
#
# It is not difficult to find other dog breed pairs with minimal inter-class variation (for instance, Curly-Coated Retrievers and American Water Spaniels).  
#
# Curly-Coated Retriever | American Water Spaniel
# - | -
# <img src="images/Curly-coated_retriever_03896.jpg" width="200"> | <img src="images/American_water_spaniel_00648.jpg" width="200">
#
# Likewise, recall that labradors come in yellow, chocolate, and black. Your vision-based algorithm will have to conquer this high intra-class variation to determine how to classify all of these different shades as the same breed.  
#
# Yellow Labrador | Chocolate Labrador | Black Labrador  
# - | - | -
# <img src="images/Labrador_retriever_06457.jpg" width="150"> | <img src="images/Labrador_retriever_06455.jpg" width="240"> | <img src="images/Labrador_retriever_06449.jpg" width="220">
#
# We also mention that random chance presents an exceptionally low bar: setting aside the fact that the classes are slightly imabalanced, a random guess will provide a correct answer roughly 1 in 133 times, which corresponds to an accuracy of less than 1%. 

# + [markdown] tags=[]
# ### Model architecture
#
# Create a CNN to classify dog breed.  At the end of your code cell block, summarize the layers of your model by executing `model.summary()`.

# + colab={} colab_type="code" id="cX4dZkEjfe6u" tags=["block:custom_classifier", "prev:data_processing", "limit:nvidia.com/gpu:1"]
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(16, 3, activation="relu", input_shape=(IMG_SIZE, IMG_SIZE, 3)),
    tf.keras.layers.MaxPool2D(),
    tf.keras.layers.Conv2D(32, 3, activation="relu"),
    tf.keras.layers.MaxPool2D(),
    tf.keras.layers.Conv2D(64, 3, activation="relu"),
    tf.keras.layers.MaxPool2D(),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(NUMBER_OF_NODES, activation="relu"),
    tf.keras.layers.Dense(133, activation="softmax")
])

# + colab={"base_uri": "https://localhost:8080/", "height": 459} colab_type="code" id="iMmibLNahFRp" outputId="31a84afd-576e-4262-efea-978fd98d0255" tags=[]
model.summary()

# + [markdown] tags=[]
# The next step is to compile the model. For this, we need to pass an optimizer and a loss function. We can also pass a list of metrics we want. In this example, we pass the _accuracy_ metric.

# + colab={} colab_type="code" id="9uRX3_UyhN90" tags=[]
model.compile(
    optimizer=tf.optimizers.Adam(learning_rate=LR),
    loss=tf.losses.categorical_crossentropy,
    metrics=["accuracy"]
)

# + [markdown] tags=[]
# Finally, we can train the model using the `fit` method. This runs on batches yielded by the data generator and prints out the _loss_ and _accuracy_ both for train and validation sets.

# + colab={"base_uri": "https://localhost:8080/", "height": 102} colab_type="code" id="Vj67Ok9dhbOC" outputId="2e64d6cf-0c21-4fac-c51f-d88661d85923" tags=[]
train_generator = get_train_generator()
valid_generator = get_valid_generator()

tb_callback = tf.keras.callbacks.TensorBoard(log_dir="custom_classifier_logs")

model.fit(train_generator, epochs=2,
    validation_data=valid_generator,
    callbacks=[tb_callback]
)

# + [markdown] tags=[]
# ### Evaluation
#
# To evaluate the final model we feed it with the test dataset and call the `evaluate` method.

# + tags=["block:custome_classifier_eval", "prev:custom_classifier", "limit:nvidia.com/gpu:1"]
test_generator = get_test_generator()

test_loss_custom, test_accuracy_custom = model.evaluate(test_generator)

print(f"The accuracy in the test set is {test_accuracy_custom:.3f}.")

# + [markdown] colab_type="text" id="Q-kl1HXWo0mP" tags=[]
# <a id='step4'></a>
# ## Step 4: Create a CNN (VGG16) to Classify Dog Breeds (using Transfer Learning)
#
# To reduce training time without sacrificing accuracy, we train a CNN using Transfer Learning. Transfer Learning is the fine-tuning of a network that was pre-trained on some big dataset with new classification layers. The idea behind is that we want to keep all the good features learned in the lower levels of the network (because there's a high probability the new images will also have those features) and just learn a new classifier on top of those. This tends to work well, especially with small datasets that don't allow for a full training of the network from scratch (it's also much faster than a full training).
#
# One way of doing Transfer Learning is by loading a pretrained model up to a point, usually chopping off the final dense part of the model and adding a fully connected layer with the output that we want (e.g. an 133-node classifier). Then, we freeze the first part of the model (i.e. the body) and train only the final layer we added.

# + colab={} colab_type="code" id="j6fsA2R8nVBa" tags=["block:vgg16_classifier", "prev:data_processing", "limit:nvidia.com/gpu:1"]
vgg_body = tf.keras.applications.VGG16(
    weights="imagenet",
    include_top=False,
    input_shape=(IMG_SIZE, IMG_SIZE, 3)
)

# + colab={} colab_type="code" id="szJl_qQ9sZuz" tags=[]
vgg_body.trainable = False

# + colab={} colab_type="code" id="YdbFhBtIp-Jx" tags=[]
inputs = tf.keras.layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
# We make sure that the vgg_body is running in inference mode here,
# by passing `training=False`. This is important for fine-tuning, as you will
# learn in a few paragraphs.
x = vgg_body(inputs, training=False)
# Convert features of shape `vgg_body.output_shape[1:]` to vectors
x = tf.keras.layers.GlobalAveragePooling2D()(x)
# A Dense classifier (categorical classification)
outputs = tf.keras.layers.Dense(133, activation="softmax")(x)

vgg_model = tf.keras.Model(inputs, outputs)

# + colab={"base_uri": "https://localhost:8080/", "height": 289} colab_type="code" id="9I2yPjo_qMTL" outputId="25b81b5f-e3d3-411c-d011-a20e18d38a11" tags=[]
vgg_model.summary()

# + colab={} colab_type="code" id="kTMJ5joCqpTE" tags=[]
vgg_model.compile(
    optimizer=tf.optimizers.Adam(learning_rate=LR),
    loss=tf.losses.categorical_crossentropy,
    metrics=["accuracy"]
)

# + colab={"base_uri": "https://localhost:8080/", "height": 459} colab_type="code" id="8IoRWrQnq5c2" outputId="3f69f065-815e-4c9c-b665-ecc2a10337b6" tags=[]
train_generator = get_train_generator()
valid_generator = get_valid_generator()

vgg_model.fit(train_generator, epochs=2,
    validation_data=valid_generator
)

# + [markdown] tags=[]
# ### Evaluation
#
# To evaluate the model on the test set we call the same `evaluate` method.

# + tags=["block:vgg16_classifier_eval", "prev:vgg16_classifier"]
test_generator = get_test_generator()

test_loss_vgg, test_accuracy_vgg = vgg_model.evaluate(test_generator)

print(f"The accuracy in the test set is {test_accuracy_vgg:.3f}.")

# + [markdown] tags=[]
# <a id='step5'></a>
# ## Step 5: Create a CNN (ResNet-50) to Classify Dog Breeds (using Transfer Learning)
#
# In this section, we will use the same procedure but with a pretrained ResNet-50 model.

# + tags=["block:resnet50_classifier", "prev:data_processing", "limit:nvidia.com/gpu:1"]
resnet_body = tf.keras.applications.ResNet50V2(
    weights="imagenet",
    include_top=False,
    input_shape=(IMG_SIZE, IMG_SIZE, 3)
)

# + tags=[]
resnet_body.trainable = False

# + tags=[]
inputs = tf.keras.layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
# We make sure that the vgg_body is running in inference mode here,
# by passing `training=False`. This is important for fine-tuning, as you will
# learn in a few paragraphs.
x = resnet_body(inputs, training=False)
# Convert features of shape `vgg_body.output_shape[1:]` to vectors
x = tf.keras.layers.Flatten()(x)
# A Dense classifier (categorical classification)
outputs = tf.keras.layers.Dense(133, activation="softmax")(x)

resnet_model = tf.keras.Model(inputs, outputs)

# + tags=[]
resnet_model.compile(
    optimizer=tf.optimizers.Adam(learning_rate=LR),
    loss=tf.losses.categorical_crossentropy,
    metrics=["accuracy"]
)

# + tags=[]
train_generator = get_train_generator()
valid_generator = get_valid_generator()

resnet_model.fit(train_generator, epochs=EPOCHS,
    validation_data=valid_generator
)

# + [markdown] tags=[]
# ### Evaluation

# + tags=["block:resnet50_classifier_eval", "prev:resnet50_classifier", "limit:nvidia.com/gpu:1"]
test_generator = get_test_generator()

test_loss_resnet, test_accuracy_resnet = resnet_model.evaluate(test_generator)

print(f"The accuracy in the test set is {test_accuracy_resnet:.3f}.")

# + [markdown] tags=[]
# <a id='step6'></a>
# ## Step 6: Write your Own Dog Classifier
#
# To create our own classifier we need a class to predict if there is a dog in the image and if that's true, return the breed. For the first part we use the `dog_classifier` method and then predict the breed using the `predict_breed` method.

# + tags=["skip"]
idx_to_class = {value: key for key, value in train_generator.class_indices.items()}

def predict_breed(images):
    probs = resnet_model.predict(images)
    pred = tf.argmax(probs, axis=1)
    label = idx_to_class[pred.numpy()[0]]
    return label.split(".")[-1]


# + tags=["skip"]
def predict_dog(image):
    image = image[None,...]
    if is_dog(image):
        pred =  predict_breed(image)
        print(f"This photo looks like a(n) {pred}.")
        return

    print("No dog detected")

image = train_generator.next()[0][0]
plt.imshow(image)
plt.show()
pred = predict_dog(image)

# + [markdown] tags=[]
# <a id='step7'></a>
# ## Step 7: Test Your Classifier
#
# In the last section, we take your new algorithm for a spin; if you have a dog, does it predict your dog's breed accurately? If you have a cat, does it mistakenly think that your cat is a dog?

# + tags=["skip"]
for img_path in sorted(glob("check_images/*")):
    print(img_path)
    img = Image.open(img_path)
    img = img.resize((224, 224))
    
    plt.imshow(img)
    plt.show()
    
    img = np.array(img)
    predict_dog(img)

# + [markdown] tags=[]
# ### Pipeline metrics
#
# This is the pipeline-metrics cell. Use it to define the pipeline metrics that KFP will produce for every pipeline run. Kale will associate each one of these metrics to the steps that produced them. Also, you will have to choose one these metrics as the Katib search objective metric.

# + tags=["pipeline-metrics"]
print(test_accuracy_resnet)

StarcoderdataJupyterScriptsDedupFiltered