Ex Data, Scientia

Home Contact

Building an AI-based image-classifier application

import numpy as np
import pandas as pd
import os
import shutil
from tensorflow.keras import layers, models
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# n_digits = 10
# weights_dir='MNIST'
# weights_name = 'classifier.h5'
# sample_dir = 'MNIST/Sample'

n_digits = int(n_digits)

l1 = layers.Input(shape = [27,27,3])
l2 = layers.Conv2D(30, [3,3], activation = 'relu')(l1)
l2a = layers.MaxPool2D([2,2])(l2)
l3 = layers.Conv2D(30, [3,3], activation = 'relu')(l2a)
l4 = layers.Flatten()(l3)
l5 = layers.Dense(300, activation = 'relu')(l4)
l6 = layers.Dense(50, activation = 'relu')(l5)
l7 = layers.Dense(n_digits, activation = 'softmax')(l6)

classifier = models.Model([l1], [l7])

classifier.load_weights(weights_dir + '/' + weights_name)

os.mkdir(sample_dir + '/dummy_class')
for i in os.listdir(sample_dir):
    if i != 'dummy_class':
        os.rename(sample_dir + '/' + i, sample_dir + '/dummy_class/' + i)
    
global datagen
datagen = ImageDataGenerator(rescale=1./255)

global sample_generator
sample_generator = datagen.flow_from_directory(
        directory = sample_dir,
        target_size = [27,27],
        batch_size = 1,
        class_mode = 'categorical',
        shuffle = False)

global n_sample_imgs
n_sample_imgs = len(sample_generator.filenames)

global preds

preds = classifier.predict_generator(
    sample_generator
    )

global preds_argmax

preds_argmax = np.argmax(preds, axis = 1)

os.mkdir(sample_dir + '/classified')
for i in range(len(np.unique(preds_argmax))):
    os.mkdir(sample_dir + '/classified/' + str(np.unique(preds_argmax)[i]))

global image_names
image_names = sample_generator.filenames

for i in range(len(image_names)):
    shutil.copy2(sample_generator.filepaths[i], 
                 sample_dir + '/classified/' + str(preds_argmax[i]) + '/' + (image_names[i]).split('/')[1])