gemstone_classification_nn.py

# -*- coding: utf-8 -*-
"""Gemstone Classification NN

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1FFPDu0hDd4rO6Ky4wo8T927dsHhAtAL8

# Explore the Dataset

Root folders for the dataset. One full dataset and one reduced dataset
"""

gemstones_all_folder = str("/content/drive/MyDrive/Colab/gemstones")
gemstones_less_folder = str("/content/drive/MyDrive/Colab/gemstones_less")
gemstones_tests_folder = str("/content/drive/MyDrive/Colab/gemstones_tests")

gemstones_folder = gemstones_all_folder

"""Get a copy of the directory path to every gemstone image"""

import os

def get_gemstones():
    Images, Labels = [], []
    for root, folders, images in os.walk(gemstones_folder):
        for image in images:
            if image.endswith('.jpg'):
                Labels.append(os.path.join(root, image))
            try:
                img = cv2.imread(os.path.join(root, image))              # read the image (OpenCV)
                img = cv2.resize(img,(int(img_w*1.5), int(img_h*1.5)))       # resize the image (images are different sizes)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # converts an image from BGR color space to RGB
                Images.append(img)
            except Exception as e:
                print(e)
    Images = np.array(Images)
    
    return (Images, Labels)

(Images, Labels) = get_gemstones()

"""Plot some random gemstones to view"""

import random
from matplotlib import pyplot as plot
from matplotlib import image as plot_image

def plot_random_gemstones(rows=5, cols=6, img_size=2.5):
    random_gemstones = plot.figure(figsize=(img_size*cols, img_size*rows))
    images = []
    for i in range(1, rows*cols+1):
        images.append(plot_image.imread(random.choice(Labels)))
    for i in range(1, rows*cols+1):
        image = images[i-1]
        ax = random_gemstones.add_subplot(rows, cols, i)
        ax.axis('off')
        plot.imshow(image)
    plot.show()

"""Run this cell to see a new random subset of the gemstones"""

plot_random_gemstones()

"""###### WIP, not relevant to final submission"""

import cv2
from random import randint

img_w = 256
img_h = 256

"""read in images and labels

try and crop the images to fit the gemstone
"""

def edge_and_cut(img, path):
    try:
        edges = cv2.Canny(img, img_w, img_h)            
        
        if(np.count_nonzero(edges)>edges.size/10000):           
            pts = np.argwhere(edges>0)
            y1,x1 = pts.min(axis=0)
            y2,x2 = pts.max(axis=0)
            
            new_img = img[y1:y2, x1:x2]           # crop the region
            new_img = cv2.resize(new_img,(img_w, img_h))  # Convert back
        else:
            new_img = cv2.resize(img,(img_w, img_h))
    
    except Exception as e:
        print(e)
        new_img = cv2.resize(img,(img_w, img_h))
    
    cv2.imwrite(path, new_img)
    return new_img

i = 2

new_img = edge_and_cut(Images[i], Labels[i])
fig, ax = plot.subplots(nrows=1, ncols=2, figsize=(4, 3))
ax[0].axis('off')
ax[0].imshow(Images[i], cmap='gray')
ax[0].set_title('Original Image', fontsize=14)
ax[1].axis('off')
ax[1].imshow(new_img, cmap='gray')
ax[1].set_title('Cropped Images', fontsize=14)

def show_cropped(img):
    emb_img = img.copy()
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    blurred = cv2.GaussianBlur(gray, (3, 3), 0)
    edges = cv2.Canny(blurred, 15, 255)
    
    if(np.count_nonzero(edges)>edges.size/10000):
        pts = np.argwhere(edges>0)
        y1,x1 = pts.min(axis=0)
        y2,x2 = pts.max(axis=0)

        new_img = img[y1:y2, x1:x2]  

        edge_size = 2 #replace it with bigger size for larger images            

        emb_img[y1-edge_size:y1+edge_size, x1:x2] = [255, 0, 0]
        emb_img[y2-edge_size:y2+edge_size, x1:x2] = [255, 0, 0]
        emb_img[y1:y2, x1-edge_size:x1+edge_size] = [255, 0, 0]
        emb_img[y1:y2, x2-edge_size:x2+edge_size] = [255, 0, 0]

        new_img = cv2.resize(new_img,(img_w, img_h))  # Convert to primary size  
        
    else:
        new_img = cv2.resize(img,(img_w, img_h))
    
    fig, ax = plot.subplots(nrows=1, ncols=4, figsize=(10, 10))
    ax[0].axis('off')
    ax[0].imshow(img, cmap='gray')
    ax[0].set_title('Original Image', fontsize=14)
    ax[1].axis('off')
    ax[1].imshow(edges, cmap='gray')
    ax[1].set_title('Canny Edges', fontsize=14)
    ax[2].axis('off')
    ax[2].imshow(emb_img, cmap='gray')
    ax[2].set_title('Bounding Box', fontsize=14)       
    ax[3].axis('off')
    ax[3].imshow(new_img, cmap='gray')
    ax[3].set_title('Cropped', fontsize=14)

for x in range(0,3):
    show_cropped(Images[randint(0,len(Images))])

"""# Data Preparation

Set up the ImageDataGenerators to normalize the rgb values, do some slight translation, and split the dataset into validation and training data
"""

import random
from keras.preprocessing.image import ImageDataGenerator

datagen_kwargs_augment = dict(
    rotation_range = 2,
    width_shift_range = 0.05,
    height_shift_range = 0.05,
    shear_range = 1,
    # fill_mode = "nearest",
    horizontal_flip = True,
    vertical_flip = True,
    rescale = 1/255,
    validation_split = 0.10,
)

datagen_kwargs_default = dict(rescale = 1/255)

target_resolution = (256, 256)
random_seed = random.randrange(0, 10000, 1)
# random_seed = 1369

training_datagen = ImageDataGenerator(**datagen_kwargs_augment)
training_generator = training_datagen.flow_from_directory(
    gemstones_folder,
    target_size = target_resolution,
    color_mode = "rgb",
    shuffle = random_seed,
    seed = 1369,
    subset = "training",
    batch_size = 12
)

validation_datagen = ImageDataGenerator(**datagen_kwargs_augment)
validation_generator = validation_datagen.flow_from_directory(
    gemstones_folder, 
    target_size = target_resolution,
    color_mode = "rgb",
    shuffle = True,
    seed = random_seed,
    subset = "validation",
    batch_size = 12
)

test_datagen = ImageDataGenerator(**datagen_kwargs_default)
test_generator = training_datagen.flow_from_directory(
    gemstones_folder, 
    target_size = target_resolution,
    color_mode = "rgb",
    shuffle = False,
    subset = "validation",
    batch_size = 12,
)

full_datagen = ImageDataGenerator(**datagen_kwargs_default)
full_generator = training_datagen.flow_from_directory(
    gemstones_folder, 
    target_size = target_resolution,
    color_mode = "rgb",
    shuffle = False,
    batch_size = 12,
)

"""# Model Architecture"""

from keras.models import Model
from keras.layers import *
from keras.regularizers import *

"""Colvolutional layers"""

image_input = Input(shape = (256, 256, 3))

# Block 1
layer = Conv2D(32, 2, activation='relu', kernel_initializer='he_uniform', kernel_regularizer=l2(0.0001), padding='same')(image_input)
layer = AveragePooling2D(2)(layer)
layer = Conv2D(32, 2, activation='relu', kernel_initializer='he_uniform', kernel_regularizer=l2(0.0001), padding='same')(layer)
layer = AveragePooling2D(2)(layer)
layer = Conv2D(32, 2, activation='relu', kernel_initializer='he_uniform', kernel_regularizer=l2(0.0001), padding='same')(layer)
layer = BatchNormalization()(layer)
layer = MaxPooling2D(2)(layer)
layer = Dropout(0.33)(layer)

# Block 2
layer = Conv2D(64, 2, activation='relu', kernel_initializer='he_uniform', kernel_regularizer=l2(0.0001), padding='same')(layer)
layer = AveragePooling2D(2)(layer)
layer = Conv2D(64, 2, activation='relu', kernel_initializer='he_uniform', kernel_regularizer=l2(0.0001), padding='same')(layer)
layer = BatchNormalization()(layer)
layer = MaxPooling2D(2)(layer)
layer = Dropout(0.33)(layer)

# Block 3
layer = Conv2D(128, 2, activation='relu', kernel_initializer='he_uniform', kernel_regularizer=l2(0.0001), padding='same')(layer)
layer = AveragePooling2D(2)(layer)
layer = Conv2D(128, 2, activation='relu', kernel_initializer='he_uniform', kernel_regularizer=l2(0.0001), padding='same')(layer)
layer = BatchNormalization()(layer)
layer = MaxPooling2D(2)(layer)
layer = Dropout(0.5)(layer)

"""Flattened & Dense layers"""

layer = Flatten()(layer)
layer = Dense(128, activation='softmax', kernel_regularizer=l2(0.0001))(layer)
layer = BatchNormalization()(layer)
layer = Dropout(0.5)(layer)
output = Dense(full_generator.num_classes, activation='softmax')(layer)

model = Model(image_input, output)

"""Model Summary"""

model.summary()

"""# Model Training

Training Parameters
"""

import numpy as np
from keras.optimizers import Adam
from keras.metrics import TopKCategoricalAccuracy

adam = Adam(learning_rate=0.0002)
model.compile(loss='categorical_crossentropy', optimizer=adam, metrics=["mse", "accuracy", TopKCategoricalAccuracy(3, "top-3 accuracy")])

def get_callbacks():
    return
    [
        tfdocs.modeling.EpochDots(),
        tf.keras.callbacks.EarlyStopping(monitor='val_mse', mode=auto, patience=10, restore_best_weights=True),
    ]

train_steps = np.ceil(training_generator.samples / training_generator.batch_size)
val_steps = np.ceil(validation_generator.samples / validation_generator.batch_size)

"""Training the model"""

# history = model.fit(
#     training_generator,
#     epochs = 250,
#     steps_per_epoch = train_steps,
#     validation_data = validation_generator,
#     validation_steps = val_steps,
#     callbacks = get_callbacks(),
#     verbose=1)

"""Save Model"""

# model.save("/content/drive/MyDrive/Colab/gemstone_model_250")

"""# Evaluation

Load a Previous Model
"""

from keras.models import load_model

model = load_model("/content/drive/MyDrive/Colab/gemstone_model_250")

"""Evaluate the Test Dataset"""

import numpy as np

test_steps = np.ceil(test_generator.samples / test_generator.batch_size)

evaluation = model.evaluate(test_generator, steps=test_steps)
# print(evaluation)

predictions = model.predict(test_generator, test_steps)
predictions = np.argmax(predictions, axis=1)
print(test_generator.classes)
print(predictions)

from sklearn.metrics import confusion_matrix, classification_report

matrix = confusion_matrix(test_generator.classes, predictions)

import seaborn as sn
import pandas as pd

df_cm = pd.DataFrame(matrix)
plot.figure(figsize=(25,17))
sn.set(font_scale=1) # for label size
sn.heatmap(df_cm, annot=True, annot_kws={"size": 5}) # font size

plot.show()

report = classification_report(test_generator.classes, predictions, target_names=test_generator.class_indices)
print(report)

"""Evaluate the Full Dataset"""

full_steps = np.ceil(full_generator.samples / full_generator.batch_size)

evaluation = model.evaluate(full_generator, steps=full_steps)
# print(evaluation)

predictions = model.predict(full_generator, full_steps)
predictions = np.argmax(predictions, axis=1)
print(full_generator.classes)
print(predictions)

from sklearn.metrics import confusion_matrix, classification_report

matrix = confusion_matrix(full_generator.classes, predictions)

import seaborn as sn
import pandas as pd

df_cm = pd.DataFrame(matrix)
plot.figure(figsize=(25,17))
sn.set(font_scale=1) # for label size
sn.heatmap(df_cm, annot=True, annot_kws={"size": 5}) # font size

plot.show()

report = classification_report(full_generator.classes, predictions, target_names=full_generator.class_indices)
print(report)

"""Diagnotic Summary"""

def summarize_diagnostics(history):

    summary, (ax1, ax2, ax3) = plot.subplots(3, figsize=(12,18))
    
    ax1.set_title('Top 3 Classification Accuracy')
    ax1.set(ylim=(0.0, 1.05))
    ax1.plot(history.history['top-3 accuracy'], color='blue', label='train')
    ax1.plot(history.history['val_top-3 accuracy'], color='orange', label='test')

    ax2.set_title('Classification Accuracy')
    ax2.set(ylim=(0.0, 1.05))
    ax2.plot(history.history['accuracy'], color='blue', label='train')
    ax2.plot(history.history['val_accuracy'], color='orange', label='test')

    ax3.set_title("Mean Squared Error")
    ax3.plot(history.history['mse'], color='blue', label='train')
    ax3.plot(history.history['val_mse'], color='orange', label='test')

    plot.show()
    plot.close()

"""Show Graphs"""

summarize_diagnostics(history)