# os imports
import os
# numopy and progress bar imports
import numpy as np
from tqdm import tqdm
import pickle
# tensorflow imports
import tensorflow
# CNN image encoders
from tensorflow.keras.applications.densenet import DenseNet201 as dn201
from tensorflow.keras.applications.densenet import DenseNet121 as dn121
from tensorflow.keras.applications.densenet import DenseNet169 as dn169
from tensorflow.keras.applications.efficientnet import EfficientNetB0
from tensorflow.keras.applications.efficientnet import EfficientNetB5 as enb5
from tensorflow.keras.applications.efficientnet import EfficientNetB7 as enb7
from tensorflow.keras.applications.resnet_v2 import ResNet50V2 as rn50v2
from tensorflow.keras.applications.inception_resnet_v2 import InceptionResNetV2 as incres
# preprocessing functions
from tensorflow.keras.applications.densenet import preprocess_input as dense_preprocess
from tensorflow.keras.applications.efficientnet import preprocess_input as efficient_preprocess
from tensorflow.keras.applications.resnet_v2 import preprocess_input as resnet_preprocess
from tensorflow.keras.applications.inception_resnet_v2 import preprocess_input as inceptionresnet_preprocess
# layers imports
from tensorflow.keras.layers import Flatten, Dropout, Dense, Conv2D, MaxPooling2D, BatchNormalization
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing import image as img
# cotnet import
# from keras_cv_attention_models import cotnet
def load_encoded_vecs(filename:str) -> dict:
""" Loads the image embeddings for each image id, we extracted offline during my research
Args:
filename (str): the whole path of npy file
Returns:
dict: encoded_vectors from filename
"""
with open(filename, 'rb') as f:
print("Image Encoded Vectors loaded from directory path:", filename)
return pickle.load(f)
def save_encoded_vecs(image_vecs:np.array, output_path:str, filename:str) -> None:
""" Function which helps us to save the encoded images into a pickle file
Args:
image_vecs (np.array): the encoded images vectors that we extracted using the encode_images function
output_path (str): the output path where we want to save our image embeddings
filename (str): a name we want to use for our npy file (ex. densenet201_image_vecs). It's not necessary to write '.pkl' at the end!
"""
path = output_path + filename + '.pkl'
with open(path, 'wb') as f:
pickle.dump(image_vecs, f, pickle.HIGHEST_PROTOCOL)
print("Image Encoded Vectors stored in:", path)
class ImageEncoder:
def __init__(self, encoder:str, images_dir_path:str, weights:str='imagenet'):
""" This class helps us to extract image embeddings with different Keras CNNs.
Args:
encoder (str): encoder name you want to use (ex. densenet201 for DenseNet201)
images_dir_path (str): The directory to store our extracted vectors
weights (str, optional): the pretrained weights you want to use for your model. It's common to use imagenet as default pretrained weights.. Defaults to 'imagenet'.
"""
self.encoder_weights = weights
self.image_dir_path = images_dir_path
# we extracted the last average pooling layer for each encoder
if encoder == 'densenet201':
self.image_shape = 224
self.preprocess = 'densenet'
model = dn201(include_top=True, weights=self.encoder_weights,
input_shape=(self.image_shape, self.image_shape, 3))
self.image_encoder = Model(
inputs=model.input, outputs=model.get_layer('avg_pool').output)
elif encoder == 'densenet121':
self.image_shape = 224
self.preprocess = 'densenet'
model = dn121(include_top=True, weights=self.encoder_weights,
input_shape=(self.image_shape, self.image_shape, 3))
self.image_encoder = Model(
inputs=model.input, outputs=model.get_layer('avg_pool').output)
elif encoder == 'densenet169':
self.image_shape = 224
self.preprocess = 'densenet'
model = dn169(include_top=True, weights=self.encoder_weights,
input_shape=(self.image_shape, self.image_shape, 3))
self.image_encoder = Model(
inputs=model.input, outputs=model.get_layer('avg_pool').output)
elif encoder == 'efficientnet5':
self.image_shape = 456
self.preprocess = 'efficientnet'
model = enb5(include_top=True, weights=self.encoder_weights,
input_shape=(self.image_shape, self.image_shape, 3))
self.image_encoder = Model(
inputs=model.input, outputs=model.get_layer('avg_pool').output)
elif encoder == 'efficientnet0':
self.image_shape = 224
self.preprocess = 'efficientnet'
model = EfficientNetB0(include_top=True, weights=self.encoder_weights,
input_shape=(self.image_shape, self.image_shape, 3))
self.image_encoder = Model(
inputs=model.input, outputs=model.get_layer('avg_pool').output)
elif encoder == 'resnet50v2':
self.image_shape = 224
self.preprocess = 'resnet'
model = rn50v2(include_top=True, weights=self.encoder_weights,
input_shape=(self.image_shape, self.image_shape, 3))
self.image_encoder = Model(
inputs=model.input, outputs=model.get_layer('avg_pool').output)
elif encoder == 'inceptionresnet':
self.image_shape = 299
self.preprocess = 'inceptionresnet'
model = incres(include_top=True, weights=self.encoder_weights,
input_shape=(self.image_shape, self.image_shape, 3))
self.image_encoder = Model(
inputs=model.input, outputs=model.get_layer('avg_pool').output)
elif encoder == 'cotnet':
self.image_shape = 224
self.preprocess = 'cotnet'
model = cotnet.CotNet50(pretrained="imagenet", num_classes=0)
self.image_encoder = Model(
inputs=model.input, outputs=model.output)
else:
print("You have to initialize a valid version of image encoder\n"
"Choices are: [densenet201, densenet121, densenet169, efficientnet0, efficientnet5, resnet50v2, inceptionresnet, cotnet]")
print("Exiting...")
return
def get_preprocessor(self) -> str:
""" Gets the pre-processing function
Returns:
str: The pre-processing name we initialized
"""
return self.preprocess
def get_image_shape(self) -> int:
""" Gets the input shape
Returns:
int: The input shape for the employed encoder
"""
return self.image_shape
def get_image_encoder(self) -> Model:
""" Gets the image encoder we built
Returns:
Model: The CNN encoder
"""
return self.image_encoder
def get_images_dirpath(self) -> str:
""" Gets the image directory path to store our vectors
Returns:
str: The image directory path
"""
return self.image_dir_path
def encode(self, _image:str, verbose:int=0) -> np.array:
""" Loads an image and it passes it in CNN encoder to extract its image embeddings.
Args:
_image (str): The image id, for which its image we want to encode
verbose (int, optional): If we want to display the extraction. Defaults to 0.
Returns:
np.array: The encoded version of the given image
"""
# case CoTNet
if self.get_preprocessor() == 'cotnet':
image = img.load_img(self.image_dir_path + _image + '.jpg')
image_array = img.img_to_array(image)
imm = tensorflow.keras.applications.imagenet_utils.preprocess_input(image_array, mode='torch')
image_encoded = self.image_encoder(
tensorflow.expand_dims(tensorflow.image.resize(imm, self.image_encoder.input_shape[1:3]), 0)).numpy()
else:
# case othe encoders
# load the image and convert it to np.array
image = img.load_img(self.image_dir_path + _image + '.jpg',
target_size=(self.image_shape, self.image_shape))
image_array = img.img_to_array(image)
image_array = np.expand_dims(image_array, axis=0)
# pre-process array in order to fit with the employed encoder
if self.get_preprocessor() == 'densenet':
preprocessed_image_array = dense_preprocess(image_array)
elif self.get_preprocessor() == 'efficientnet':
preprocessed_image_array = efficient_preprocess(image_array)
elif self.get_preprocessor() == 'resnet':
preprocessed_image_array = resnet_preprocess(image_array)
elif self.get_preprocessor() == 'inceptionresnet':
preprocessed_image_array = inceptionresnet_preprocess(
image_array)
# extract image embeddings
image_encoded = self.image_encoder.predict(preprocessed_image_array, verbose=verbose)
return image_encoded
def encode_images(self, images:list) -> np.array:
""" Loads an image list with image ids, and extract their image embeddings
Args:
images (list): Image IDs list
Returns:
np.array: All image vectors
"""
image_vecs = {_image: self.encode(_image) for _image in
tqdm(images, desc="Encoding images", position=0, leave=True)}
return image_vecs
