from google.colab import drive
drive.mount('/content/drive')
!nvidia-smi
RUN_TRAINING = True
LOAD_PREV_MODEL = False
RUN_TESTING = True
if RUN_TRAINING and LOAD_PREV_MODEL:
raise Exception('Redundant training is being run. Change RUN_TRAINING or LOAD_PREV_MODEL value.')
if not RUN_TRAINING and not LOAD_PREV_MODEL:
raise Exception('No weights will be loaded for the testing phase. Change RUN_TRAINING or LOAD_PREVIOUS_MODEL.')
# Variables and hyperparameters for model
BATCH_SIZE = 32
TOTAL_LOSS_VAE_MULTIPLIER = 100
NUM_EPOCHS = 5
LAMBDA = 1e-3
LEARNING_RATE = 0.0001
# Necessary tensorflow addons for triplet loss
!pip install -q -U tensorflow-addons
# For saving models
!pip install -q pyyaml h5py
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.pyplot import imshow
import io
import tensorflow_addons as tfa
import os
vae = keras.models.load_model('/content/drive/MyDrive/saved_model')
vae.build((None,256,256,3))
class ConvexCombination(layers.Layer):
'''
Custom layer for convex combination.
Input: 2 x None x 256 x 256 x 3 tensor, where [0] is batch of original vehicle images and [1] is a batch
of vehicle images with salient features highlighted by the VAE.
Output: None x 256 x 256 x 3 tensor (a batch of convex-combined vehicle images)
'''
def __init__(self):
super(ConvexCombination, self).__init__()
init_weight_tensor = tf.convert_to_tensor(0.01)
self.w = tf.Variable(initial_value=init_weight_tensor, trainable=True)
def call(self, inputs):
clipped_w = tf.clip_by_value(self.w, 0, 1)
orig_img_part = clipped_w * inputs[0]
recon_img_part = (1 - clipped_w) * inputs[1]
return orig_img_part + recon_img_part
# Import the pre-trained ResNet50 model
resnet50 = tf.keras.applications.ResNet50(include_top=False, input_shape=(256,256,3), weights=None)
def load_aic21():
'''
Loads AIC21 onto local disk.
'''
if not (tf.io.gfile.exists('/content/AIC21')):
!mkdir /content/AIC21
print("Zipping AIC21...")
!zip -FF /content/drive/MyDrive/AIC21/AIC21_Track2_ReID.zip --out /content/AIC21/AIC21_Track2_ReID_full.zip
%cd /content/AIC21
print("Unzipping AIC21...")
!unzip -FF -q AIC21_Track2_ReID_full.zip
print("Unzipping AIC21 done!")
return "/content/AIC21/AIC21_Track2_ReID"
from IPython.display import clear_output
train_dir = load_aic21()
clear_output()
already_run = False
if not already_run:
%cd /content/AIC21/AIC21_Track2_ReID/image_train/
# Copy the script which creates the AIC21 image folders into the correct location
!cp /content/drive/MyDrive/AIC21/format_aic21_ids_linux.sh /content/AIC21/AIC21_Track2_ReID/image_train
!chmod 777 /content/AIC21/AIC21_Track2_ReID/image_train/format_aic21_ids_linux.sh
!/content/AIC21/AIC21_Track2_ReID/image_train/./format_aic21_ids_linux.sh
already_run = True
class ReId(keras.Model):
def __init__(self, vae, num_ids, **kwargs):
super(ReId, self).__init__(**kwargs)
self.vae = vae
self.convex_combination = ConvexCombination()
self.resnet_50 = resnet50
self.glob_avg_pool = keras.layers.GlobalAveragePooling2D()
self.bnneck = keras.layers.BatchNormalization()
self.final_fc_layer = keras.layers.Dense(num_ids)
self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
self.vae_loss_tracker = keras.metrics.Mean(
name="vae_loss"
)
self.classification_loss_tracker = keras.metrics.Mean(name="classification_loss")
self.triplet_loss_tracker = keras.metrics.Mean(name="triplet_loss")
@property
def metrics(self):
return [
self.total_loss_tracker,
self.vae_loss_tracker,
# Classification loss is cross-entropy loss here
self.classification_loss_tracker,
self.triplet_loss_tracker,
]
def train_step(self, data):
x, y = data
y = tf.cast(y, tf.int32)
with tf.GradientTape() as tape:
# Calculate VAE loss
z_mean, z_log_var, z = self.vae.encoder(x)
reconstruction = self.vae.decoder(z)
reconstruction_loss = tf.reduce_mean(
keras.losses.mean_squared_error(x, reconstruction)
)
kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
# Calculate VAE loss
vae_loss = reconstruction_loss + (LAMBDA * kl_loss)
embedding_output = self.call(x)
bnneck_output = self.bnneck(embedding_output)
training_output = self.final_fc_layer(bnneck_output)
# Calculate triplet loss
triplet_loss = tfa.losses.triplet_semihard_loss(y_true=y, y_pred=embedding_output)
# Calculate cross-entropy loss
ce_loss = keras.losses.sparse_categorical_crossentropy(y, training_output)
# Sum up the losses
total_loss = triplet_loss + ce_loss + (TOTAL_LOSS_VAE_MULTIPLIER * vae_loss)
grads = tape.gradient(total_loss, self.trainable_weights)
self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
self.total_loss_tracker.update_state(total_loss)
self.classification_loss_tracker.update_state(ce_loss)
self.triplet_loss_tracker.update_state(triplet_loss)
self.vae_loss_tracker.update_state(vae_loss)
return {
"loss": self.total_loss_tracker.result(),
"classification loss:": self.classification_loss_tracker.result(),
"triplet loss:": self.triplet_loss_tracker.result(),
"vae loss:": self.vae_loss_tracker.result(),
}
def call(self, data):
recon_img = self.vae(data)
convex_combo_input = [data, (data - recon_img)]
convex_combo_output = self.convex_combination(convex_combo_input)
final_output = self.glob_avg_pool(self.resnet_50(convex_combo_output))
return final_output
def normalize_vehicle_img(input_img):
rescaled_input_img = input_img / 255.0
output_img = rescaled_input_img - (tf.reduce_mean(rescaled_input_img))
output_img = output_img / (tf.math.reduce_std(rescaled_input_img))
output_img = (output_img * 0.5) + 0.5
return tf.clip_by_value(output_img, 0, 1)
import tensorflow_datasets as tfds
import os
def read_data(training_directory):
'''
Returns a three-tuple. The first item is a list of lists of paths to images,
the second is a list of lists of labels, and the third is a dictionary mapping
the class directory names to the labels in the second list.
Parameters:
training_directory: Path to directory containing directories of images where
each directory only contains images of one ID.
'''
image_list_final = []
label_list_final = []
label_map_dict = {}
count_label = 0
# Iterate through each class directory within training directory
for class_name in os.listdir(training_directory):
image_list = []
label_list = []
# Get path of class directory
class_path = os.path.join(training_directory, class_name)
# Ensure path is a directory and not a file
if not (os.path.isdir(class_path)):
continue
# Entering name of directory as key in dictionary returns the index of the
# class
label_map_dict[class_name] = count_label
# Iterate through each image in the class directory
for image_name in os.listdir(class_path):
# Get path of the image
image_path = os.path.join(class_path, image_name)
# Add index of the class to the label list
label_list.append(count_label)
# Add path to image to the image list
image_list.append(image_path)
count_label += 1
# Add the new lists to the final lists
image_list_final.append(image_list)
label_list_final.append(label_list)
return image_list_final, label_list_final, label_map_dict
def _parse_function(filename, label):
'''
Function used to parse files into dataset.
'''
image_string = tf.io.read_file(filename, "file_reader")
image_decoded = tf.image.decode_jpeg(image_string, channels=3)
image = tf.cast(image_decoded, tf.float32)
image = tf.image.resize(image, (256,256))
image = normalize_vehicle_img(image)
return image, label
def gen_list_of_id_datasets(training_directory):
'''
Returns a list of datasets where each dataset contains the images from a
single ID.
Parameters:
training_directory: Path to directory containing directories of images where
each directory only contains images of one ID.
'''
image_list_final, label_list_final, label_map_dict = read_data(training_directory)
dataset_list = []
for i,j in zip(image_list_final, label_list_final):
dataset = tf.data.Dataset.from_tensor_slices((tf.constant(i), tf.constant(j)))
dataset = dataset.shuffle(len(i))
#dataset = dataset.repeat(epochs)
dataset = dataset.map(_parse_function)
dataset_list.append(dataset)
img_count = sum( [ len(inner_lst) for inner_lst in image_list_final])
print("Created datasets for", str(len(dataset_list)), "IDs, containing a total of", str(img_count), "images.")
return dataset_list
import random
def gen_list_of_equalcut_datasets(training_directory, N, split=1.0):
'''
Returns a list of datasets where each dataset contains the images from a
single ID.
Parameters:
training_directory: Path to directory containing directories of images where
each directory only contains images of one ID.
N: Factor to divide BATCH_SIZE by in order to determine how many instances of each ID will exist
within a batch.
split=1.0: Proportion of the dataset to take images from for cases where a smaller sample size is needed.
'''
# Calculate how many instances of each ID should be within each batch.
K = int(BATCH_SIZE / N)
# Get list of image file names, label indeces and a dictionary linking indices to class names
image_list_final, label_list_final, label_map_dict = read_data(training_directory)
# Create dictionary to map labels to paths
id_split_dict = {}
for i,j in zip(image_list_final, label_list_final):
for a,b in zip(i,j):
if b in id_split_dict:
id_split_dict[b].append(a)
else:
id_split_dict[b] = [a]
# Shuffle dataset
kys = list(id_split_dict.keys())
random.shuffle(kys)
# Value to keep track of current index in each of the lists
i = 0
# Lists of paths and labels which will be in the final dataset
final_dataset_paths = []
final_dataset_labels = []
while True:
# Identify the lists which still have images that haven't been added to the dataset yet
remaining_list_ids = []
for id in kys:
if len(id_split_dict[id]) > i:
remaining_list_ids.append(id)
# Ending condition for loop: when there aren't any (or there's just one) lists left with image paths in them
if len(remaining_list_ids) <= 1:
break
# Concatenate K image paths from each list into the new list of paths
for id in remaining_list_ids:
final_dataset_paths.extend((id_split_dict[id])[i:i+K])
ext_length = len(id_split_dict[id][i:i+K])
final_dataset_labels.extend([id]*ext_length)
# Move forward K places
i += K
# Create dataset from the list of paths
dataset = tf.data.Dataset.from_tensor_slices((tf.constant(tuple(final_dataset_paths)), tf.constant(tuple(final_dataset_labels))))
#dataset = dataset.repeat(epochs)
# Convert each of the paths into actual images
dataset = dataset.map(_parse_function).batch(BATCH_SIZE)
print("Created dataset of length", len(dataset))
return dataset
dl = gen_list_of_equalcut_datasets('/content/AIC21/AIC21_Track2_ReID/image_train', 8)
reid = ReId(vae, 440)
!mkdir -p /content/tr
checkpoint_path = "/content/training_1/cp.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)
# Create a callback that saves the model's weights
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
save_weights_only=True,
verbose=1)
if RUN_TRAINING:
reid.compile(optimizer=keras.optimizers.Adam(learning_rate=LEARNING_RATE))
reid.fit(dl, epochs=NUM_EPOCHS, batch_size=BATCH_SIZE, callbacks=[cp_callback])
if RUN_TRAINING:
# Save the entire model as a SavedModel.
!mkdir -p /content/saved_model
reid.build((None,256,256,3))
reid.save('/content/drive/MyDrive/full_pipeline_vehicle_reid_prev_model/reid_latest')
reid.summary()
from tensorflow.keras.preprocessing.image import ImageDataGenerator
datagen = ImageDataGenerator(rescale=1.0/255.0,samplewise_center=True, samplewise_std_normalization=True)
test_it = datagen.flow_from_directory(train_dir, classes=['image_test'], batch_size=BATCH_SIZE,
class_mode=None)
if LOAD_PREV_MODEL:
reid. = keras.models.load_model('/content/drive/MyDrive/full_pipeline_vehicle_reid_prev_model/reid_09_06_21_12_51/reid_09_06_21_12_51.data-00000-of-00001')
def show_imgs(arr):
'''
Shows multiple images in an array of images.
Parameters:
arr: Array of images
'''
for img in arr:
plt.figure()
plt.imshow(img)
img1 = test_it[0][7]
img2 = test_it[0][11]
show_imgs([img1,img2])
from datetime import datetime
before = datetime.now()
pred1 = reid.predict(np.expand_dims(img1, axis=0))
after = datetime.now()
print("Running pred1 took",(after-before))
from datetime import datetime
before = datetime.now()
pred2 = reid.predict(np.expand_dims(img2, axis=0))
after = datetime.now()
print("Running pred2 took",(after-before))
# The image to query the gallery set with
gallery_test_img = test_it[0][8]
show_imgs([gallery_test_img])
# Pass query image through model
query_embedding = reid.predict(np.expand_dims(gallery_test_img, axis=0))
preds = reid.predict(test_it, workers=1, use_multiprocessing=False, verbose=1)
preds.shape
# Calculate the Euclidean distances for each image from the query image
euc_distances = []
i = 0
for pred in preds:
dist = query_embedding - pred
euc_distance = np.linalg.norm(dist)
euc_distances.append((euc_distance, i))
i += 1
def euc_sort(e):
return e[0]
euc_distances.sort(key=euc_sort)
euc_distances[0:10]
tst_img = test_it[375][31]
show_imgs([tst_img])
import math
nan_ratios = []
print("Starting...")
for weight in reid.weights:
num_nans = 0
num_weights = 0
for single_weight in weight.numpy().flatten():
num_weights += 1
if math.isnan(single_weight):
num_nans += 1
nan_ratios.append(num_nans / num_weights)
import matplotlib.pyplot as plt
plt.plot(nan_ratios)
plt.ylabel('some numbers')
plt.show()