NOTE: This Use Case is not purposed for resource constrained devices.
Gemstone Classification using Deep Learning¶
Credit: AITS Cainvas Community
Photo by Ivan Mesaros on Dribbble
A gemstone (gem, fine gem, jewel, precious stone, or semi-precious stone) is a piece of mineral crystal which, in cut and polished form, is used to make jewelry or other adornments. In this notebook, we will be classifying the type of gemstone based on the given image.¶
Import all the required libraries¶
In [1]:
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import wget
import os
import cv2
from random import randint
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
from sklearn.metrics import classification_report, log_loss, accuracy_score
from sklearn.model_selection import train_test_split
In [2]:
!pip install wget
Unzip the dataset to use in our notebook¶
In [3]:
!wget -N "https://cainvas-static.s3.amazonaws.com/media/user_data/cainvas-admin/gemstones.zip"
!unzip -qo gemstones.zip
In [4]:
directory = 'gemstones/train/'
In [5]:
#printing all the gemstone categories present in our dataset
Name=[]
for file in os.listdir(directory):
Name+=[file]
print(Name)
print(len(Name))
Map and display all the categories present in our dataset. There are total 68 different kinds of gemstones.¶
In [6]:
gems_map = dict(zip(Name, [t for t in range(len(Name))]))
print(gems_map)
r_gems_map=dict(zip([t for t in range(len(Name))],Name))
In [7]:
img_w, img_h = 100, 100
Create functions to read images and labels of gemstones from the training dataset.¶
In [8]:
#function which reads images and class names
def read_images():
Images, Labels = [], []
for root, dirs, files in os.walk('gemstones/train/'):
f = os.path.basename(root)
for file in files:
Labels.append(f)
try:
image = cv2.imread(root+'/'+file) # read the image (OpenCV)
image = cv2.resize(image,(int(img_w), int(img_h))) # resize the image (images are different sizes)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # converts an image from BGR color space to RGB
Images.append(image)
except Exception as e:
print(e)
Images = np.array(Images)
return (Images,Labels)
In [9]:
#function which converts string labels to numbers
def get_class_index(Labels):
for i, n in enumerate(Labels):
for j, k in enumerate(Name):
if n == k:
Labels[i] = j
Labels = np.array(Labels)
return Labels
Read images and labels from the training set.¶
In [10]:
Train_Imgs, Train_Lbls = read_images()
Train_Lbls = get_class_index(Train_Lbls)
print('Shape of train images: {}'.format(Train_Imgs.shape))
print('Shape of train labels: {}'.format(Train_Lbls.shape))
Printing some random images from the gemstone training set.¶
In [11]:
dim = 5
f,ax = plt.subplots(dim,dim)
f.subplots_adjust(0,0,2,2)
for i in range(0,dim):
for j in range(0,dim):
rnd_number = randint(0,len(Train_Imgs))
cl = Train_Lbls[rnd_number]
ax[i,j].imshow(Train_Imgs[rnd_number])
ax[i,j].set_title(Name[cl]+': ' + str(cl))
ax[i,j].axis('off')
Split the training dataset into train and validation sets.¶
In [12]:
from sklearn.model_selection import train_test_split
X_train, X_val, y_train, y_val = train_test_split(Train_Imgs, Train_Lbls, shuffle = True, test_size = 0.2, random_state = 42)
print('Shape of X_train: {}, y_train: {} '.format(X_train.shape, y_train.shape))
print('Shape of X_val: {}, y_val: {} '.format(X_val.shape, y_val.shape))
In [13]:
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense,Conv2D,MaxPooling2D,Dropout,Flatten,Activation,BatchNormalization, AveragePooling2D
from tensorflow.keras.models import model_from_json
from tensorflow.keras.models import load_model
Create a sequential model.¶
In [14]:
model = Sequential()
model.add(Conv2D(filters=32, kernel_size=(3,3),input_shape=(100,100,3), activation='relu', padding = 'same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(filters=64, kernel_size=(3,3), activation='relu', padding = 'same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(filters=64, kernel_size=(3,3), activation='relu', padding = 'same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(filters=64, kernel_size=(3,3), activation='relu', padding = 'same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(filters=64, kernel_size=(3,3), activation='relu', padding = 'same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Conv2D(filters=64, kernel_size=(3,3), activation='relu', padding = 'same'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(256))
model.add(Activation('relu'))
model.add(Dropout(0.2))
model.add(Dense(len(gems_map)))
model.add(Activation('softmax'))
model.summary()
In [15]:
model.compile(optimizer='adam',loss='sparse_categorical_crossentropy',metrics=['accuracy'])
Perform data augmentation on the images so that we can achieve more relevant data.¶
In [16]:
train_datagen = ImageDataGenerator(vertical_flip=True,
horizontal_flip=True,
rotation_range=40,
width_shift_range=0.2,
height_shift_range=0.2,
zoom_range=0.1,
validation_split=0.2)
val_datagen = ImageDataGenerator()
In [17]:
batch_size = 32
In [18]:
n = randint(0,len(X_train))
samples = np.expand_dims(X_train[n], 0)
it = train_datagen.flow(samples, batch_size=batch_size)
cols = 7
fig, ax = plt.subplots(nrows=1, ncols=cols, figsize=(15, 10))
ax[0].imshow(X_train[n], cmap='gray')
ax[0].set_title('Original', fontsize=10)
for i in range(1,cols):
batch = it.next() # generate batch of images
image = batch[0].astype('uint32') # convert to unsigned int for viewing
ax[i].set_title('augmented {}'.format(i), fontsize=10)
ax[i].imshow(image, cmap='gray')
In [19]:
train_gen = train_datagen.flow(X_train, y_train, batch_size=batch_size)
val_gen = val_datagen.flow(X_val, y_val, batch_size=batch_size)
In [20]:
EPOCHS = 80
iter_per_epoch = len(X_train) // batch_size
val_per_epoch = len(X_val) // batch_size
Train the sequential model.¶
In [21]:
m = model.fit(
train_gen,
steps_per_epoch= iter_per_epoch,
epochs=EPOCHS,
validation_data = val_gen,
validation_steps = val_per_epoch,
verbose = 1
)
Plot accuracy and loss graphs to evaluate the performance of our model.¶
In [22]:
plt.plot(m.history['loss'])
plt.plot(m.history['val_loss'])
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Model loss')
plt.show()
plt.plot(m.history['accuracy'])
plt.plot(m.history['val_accuracy'])
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Model accuracy')
plt.show()
Read the images and labels from the test set.¶
In [23]:
def read_imagest():
Images, Labels = [], []
for root, dirs, files in os.walk('gemstones/test/'):
f = os.path.basename(root)
for file in files:
Labels.append(f)
try:
image = cv2.imread(root+'/'+file) # read the image (OpenCV)
image = cv2.resize(image,(int(img_w), int(img_h))) # resize the image (images are different sizes)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # converts an image from BGR color space to RGB
Images.append(image)
except Exception as e:
print(e)
Images = np.array(Images)
return (Images,Labels)
In [24]:
Test_Imgs, Test_Lbls = read_imagest()
Test_Lbls = get_class_index(Test_Lbls)
Make predictions on random images from the test set.¶
In [25]:
f,ax = plt.subplots(5,5)
f.subplots_adjust(0,0,2,2)
for i in range(0,5,1):
for j in range(0,5,1):
rnd_number = randint(0,len(Test_Imgs))
pred_image = np.array([Test_Imgs[rnd_number]])
pred_class = model.predict_classes(pred_image)[0]
pred_prob = model.predict(pred_image).reshape(40)
act =Name[Test_Lbls[rnd_number]]
ax[i,j].imshow(Test_Imgs[rnd_number])
ax[i,j].imshow(pred_image[0])
if(Name[pred_class] != Name[Test_Lbls[rnd_number]]):
t = '{} [{}]'.format(Name[pred_class], Name[Test_Lbls[rnd_number]])
ax[i,j].set_title(t, fontdict={'color': 'darkred'})
else:
t = '[OK] {}'.format(Name[pred_class])
ax[i,j].set_title(t)
ax[i,j].axis('off')