from google.colab import drive
drive.mount('/content/drive')
import os
import re
import sys
import math
import random
import logging
import requests
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow_probability as tfp
from tqdm.notebook import *
from sklearn import metrics
from sklearn.model_selection import train_test_split
import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.keras.layers import *
from tensorflow.keras import backend as K
from tensorflow.keras.optimizers import *
from tensorflow.keras import mixed_precision
try:
import albumentations as A
import tensorflow_addons as tfa
from vit_keras import vit, layers
import efficientnet.tfkeras as efn
from classification_models.tfkeras import Classifiers
except:
!pip install -qq gcsfs
!pip install -qq vit-keras
!pip install -qq efficientnet
!pip install -qq albumentations
!pip install -qq image-classifiers
!pip install -qq tensorflow-addons
!gdown --id 16MJuZ3wovKcc1B7V6eaUZKwkH1Dipykg
import albumentations as A
import tensorflow_addons as tfa
from vit_keras import vit, layers
import efficientnet.tfkeras as efn
from classification_models.tfkeras import Classifiers
logging.disable(logging.WARNING)
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', tpu.master())
#tf.config.set_soft_device_placement(True)
tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)
print("REPLICAS: ", strategy.num_replicas_in_sync)
AUTO = tf.data.experimental.AUTOTUNE
config = {
"seed": 1312,
"lr": 1e-4,
"valid_size": 0.05,
"epochs": 50,
"batch_size": 8 * strategy.num_replicas_in_sync,
"image_size": [800, 800, 3],
"embedding_dimensions": 1024,
"n_classes": 81313,
"image_paths": [
'gs://kds-02a5b8675d55c9a79251760390f626ffd3a0807438e67d2c7edea3cb',
'gs://kds-7aa79b5bc6f9af00ef7fd0c00f645a0abe32ebc8426ce4dd299077e6',
'gs://kds-971e58e5965ae894e73daa42bc53c93538d7afd16f8cc31a7e0ec68b',
'gs://kds-9c55456bf87ba673337d07f52d56d36e16e1ead2da1bb15e16610dd7',
'gs://kds-5949695dad43c3d30cd209773b1365e354a3613d9b63a25051eb4305',
'gs://kds-c4a1c215158b3e3002adac53b7de364a742bae7c6557212557d0378e',
'gs://kds-969690135ac88129bd11436ab06669e2aa1a23d66083087e9a692255',
'gs://kds-1793eae3b59c9d40461a1b04d82452040a42b8b0f063554dcd024ccf',
'gs://kds-73589684dcf1e8ebd2d37605cdb19ebfa650cccb17f36221e0ee48ac',
'gs://kds-d0e89541a75c0bf3642ef1f91a8f5ed5ff25630b417d504fba0d77ab',
'gs://kds-ebc1e3faa8dbe5cc846799207a330dd245c63befa4c93241ce526d1a',
'gs://kds-3b830115dd341cae8ac39c8d611fe4bc6fcad798e1e84ec336de6b36',
'gs://kds-8cb07756a69b20d2e4f294f26a928e7d285cf6d6db6640d9873c7ee1',
'gs://kds-e9f50c957467dcf68c27f8b865d87720e2cd4f1ff058733e474920d0',
'gs://kds-20d0c45f1756ef8664edcf2aa2db83ca4a0efbfb8ea533cd0bf85633',
'gs://kds-6ea4f0da52f0c996a7fbe5835459d1301a824236e793dadea8776b7d',
'gs://kds-05fdfc56d42ee43c29806c0dc9a06edb5d6f9a1cd82871b67538986a',
'gs://kds-6fb2e328eab4580255aa16b7f9bc7074babb1f68488144750b2e1c9f',
'gs://kds-02f60c478a0a861b6b80dba0ebb07e4d8547ad036ea01065a560d520',
'gs://kds-75583a573fde7550697fd8f591a2db35399acd2e772ca31eaecce602'
],
"encoded_csv_path": 'train_encoded.csv',
"save_path": "/content/drive/My Drive/2021_GLR/",
"margin": "AdaCos",
"backbone": "efficientnetv2-b3",
"last_epoch": 34,
"total_save": 5
}
config["save_path"] += config["backbone"]
if (config["last_epoch"] != 0):
config["weight_path"] = os.path.join(config["save_path"], "model_{}_{}_{}.h5".format(config["backbone"], config["margin"], (config["last_epoch"])%config["total_save"]))
else:
config["weight_path"] = None
augmentation = A.Compose([
A.HorizontalFlip(p = 0.5),
A.RandomBrightnessContrast(brightness_limit=0.1, p=0.2),
A.JpegCompression(quality_lower=95, quality_upper=100, p=0.25),
A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=15, p=0.25),
A.Cutout(num_holes=2, max_h_size=4, max_w_size=4, p=0.1),
])
def seed_everything(seed):
random.seed(seed)
np.random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
tf.random.set_seed(seed)
seed_everything(config["seed"])
def transform(image, label):
image = augmentation(image=image)["image"]
return image, label
def decode_image(image_data, image_size = config['image_size']):
image = tf.image.decode_jpeg(image_data, channels = 3)
image = tf.cast(image, tf.float32) / 255.0
#image = tf.image.resize(image, (config["image_size"][0], config["image_size"][1]))
image = tf.image.resize_with_pad(image, target_height = config["image_size"][0], target_width = config["image_size"][1])
image = tf.reshape(image, image_size)
return image
def count_data_items(filenames):
records = [int(re.compile(r"_([0-9]*)\.").search(filename).group(1)) for filename in filenames]
df = pd.read_csv(config["encoded_csv_path"])
n = df[df['group'].isin(records)].shape[0]
return n
def read_tfrecord(example):
TFREC_FORMAT = {
"image": tf.io.FixedLenFeature([], tf.string),
"target": tf.io.FixedLenFeature([], tf.int64)
}
example = tf.io.parse_single_example(example, TFREC_FORMAT)
image = decode_image(example['image'], config['image_size'])
target = tf.cast(example["target"], tf.int32)
if (config['margin'] != "ArcMargin"):
target = tf.one_hot(target, tf.constant(config["n_classes"], name = "C"), on_value = 1.0, off_value = 0.0, axis =-1)
return image, target
def load_dataset(filenames, ordered = False):
ignore_order = tf.data.Options()
if not ordered:
ignore_order.experimental_deterministic = False
dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads = AUTO)
dataset = dataset.with_options(ignore_order)
dataset = dataset.map(read_tfrecord, num_parallel_calls = AUTO)
return dataset
def arcface_format(image, target):
return {'input': image, 'label': target}, target
def get_training_dataset(filenames, ordered = False, do_aug = False):
dataset = load_dataset(filenames, ordered = ordered)
if (do_aug):
dataset = dataset.map(transform, num_parallel_calls = AUTO)
dataset = dataset.map(arcface_format, num_parallel_calls = AUTO)
dataset = dataset.repeat()
dataset = dataset.shuffle(config["seed"])
dataset = dataset.batch(config["batch_size"])
dataset = dataset.prefetch(AUTO)
return dataset
def get_validation_dataset(filenames, ordered = True, prediction = False):
dataset = load_dataset(filenames, ordered = ordered)
dataset = dataset.map(arcface_format, num_parallel_calls = AUTO)
if prediction:
dataset = dataset.batch(config["batch_size"])
else:
dataset = dataset.batch(config["batch_size"])
dataset = dataset.prefetch(AUTO)
return dataset
class CosineSimilarity(tf.keras.layers.Layer):
def __init__(self, num_classes, **kwargs):
super().__init__(**kwargs)
self.num_classes = num_classes
def build(self, input_shape):
input_dim = input_shape[-1]
self.W = self.add_weight(shape=(input_dim, self.num_classes),
initializer='random_normal',
trainable=True)
def call(self, inputs):
x = tf.nn.l2_normalize(inputs, axis=-1) # (batch_size, ndim)
w = tf.nn.l2_normalize(self.W, axis=0) # (ndim, nclass)
cos = tf.matmul(x, w) # (batch_size, nclass)
return cos
def get_config(self):
config = super().get_config().copy()
config.update({
'num_classes': self.num_classes
})
return config
class ArcFace(tf.keras.layers.Layer):
def __init__(self, num_classes, margin=0.5, scale=64, **kwargs):
super().__init__(**kwargs)
self.num_classes = num_classes
self.margin = margin
self.scale = scale
self.cos_similarity = CosineSimilarity(num_classes)
def call(self, inputs, training):
# If not training (prediction), labels are ignored
feature, labels = inputs
cos = self.cos_similarity(feature)
if training:
theta = tf.acos(tf.clip_by_value(cos, -1, 1))
cos_add = tf.cos(theta + self.margin)
mask = tf.cast(labels, dtype=cos_add.dtype)
logits = mask*cos_add + (1-mask)*cos
logits *= self.scale
return logits
else:
return cos
def get_config(self):
config = super().get_config().copy()
config.update({
'num_classes': self.num_classes,
'margin': self.margin,
'scale': self.scale
})
return config
class GeMPooling(tf.keras.layers.Layer):
def __init__(self, p=1.0, eps=1e-7):
super().__init__()
self.p = p
self.eps = 1e-7
def get_config(self):
config = super().get_config().copy()
config.update({
'p': self.p,
'eps': self.eps
})
return config
def call(self, inputs: tf.Tensor, **kwargs):
inputs = tf.clip_by_value(inputs, clip_value_min=self.eps, clip_value_max=tf.reduce_max(inputs))
inputs = tf.pow(inputs, self.p)
inputs = tf.reduce_mean(inputs, axis=[1, 2])
inputs = tf.pow(inputs, 1./self.p)
return inputs
class CosineSimilarity(tf.keras.layers.Layer):
"""
Cosine similarity with classwise weights
"""
def __init__(self, num_classes, **kwargs):
super().__init__(**kwargs)
self.num_classes = num_classes
def build(self, input_shape):
input_dim = input_shape[-1]
self.W = self.add_weight(shape=(input_dim, self.num_classes),
initializer='random_normal',
trainable=True)
def call(self, inputs):
x = tf.nn.l2_normalize(inputs, axis=-1) # (batch_size, ndim)
w = tf.nn.l2_normalize(self.W, axis=0) # (ndim, nclass)
cos = tf.matmul(x, w) # (batch_size, nclass)
return cos
def get_config(self):
config = super().get_config().copy()
config.update({
'num_classes': self.num_classes
})
return config
class ArcMargin(tf.keras.layers.Layer):
def __init__(self, n_classes, s=30, m=0.50, easy_margin=False, ls_eps=0.0, **kwargs):
super(ArcMargin, self).__init__(**kwargs)
self.n_classes = n_classes
self.s = s
self.m = m
self.ls_eps = ls_eps
self.easy_margin = easy_margin
self.cos_m = tf.math.cos(m)
self.sin_m = tf.math.sin(m)
self.th = tf.math.cos(math.pi - m)
self.mm = tf.math.sin(math.pi - m) * m
def get_config(self):
config = super().get_config().copy()
config.update({
'n_classes': self.n_classes,
's': self.s,
'm': self.m,
'ls_eps': self.ls_eps,
'easy_margin': self.easy_margin,
})
return config
def build(self, input_shape):
super(ArcMargin, self).build(input_shape[0])
self.W = self.add_weight(
name='W',
shape=(int(input_shape[0][-1]), self.n_classes),
initializer='glorot_uniform',
dtype='float32',
trainable=True,
regularizer=None
)
def call(self, inputs):
X, y = inputs
y = tf.cast(y, dtype=tf.int32)
cosine = tf.matmul(
tf.math.l2_normalize(X, axis=1),
tf.math.l2_normalize(self.W, axis=0)
)
sine = tf.math.sqrt(1.0 - tf.math.pow(cosine, 2))
phi = cosine * self.cos_m - sine * self.sin_m
if self.easy_margin:
phi = tf.where(cosine > 0, phi, cosine)
else:
phi = tf.where(cosine > self.th, phi, cosine - self.mm)
one_hot = tf.cast(
tf.one_hot(y, depth=self.n_classes),
dtype=cosine.dtype
)
if self.ls_eps > 0:
one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.n_classes
output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
output *= self.s
return output
class ArcFace(tf.keras.layers.Layer):
"""
Implementation of https://arxiv.org/pdf/1801.07698.pdf
"""
def __init__(self, num_classes, margin=0.5, scale=64, **kwargs):
super().__init__(**kwargs)
self.num_classes = num_classes
self.margin = margin
self.scale = scale
self.cos_similarity = CosineSimilarity(num_classes)
def call(self, inputs, training):
# If not training (prediction), labels are ignored
feature, labels = inputs
cos = self.cos_similarity(feature)
if training:
theta = tf.acos(tf.clip_by_value(cos, -1, 1))
cos_add = tf.cos(theta + self.margin)
mask = tf.cast(labels, dtype=cos_add.dtype)
logits = mask*cos_add + (1-mask)*cos
logits *= self.scale
return logits
else:
return cos
def get_config(self):
config = super().get_config().copy()
config.update({
'num_classes': self.num_classes,
'margin': self.margin,
'scale': self.scale
})
return config
class AdaCos(tf.keras.layers.Layer):
def __init__(self, num_classes, **kwargs):
super().__init__(**kwargs)
self.num_classes = num_classes
self.cos_similarity = CosineSimilarity(num_classes)
self.scale = tf.Variable(tf.sqrt(2.0)*tf.math.log(num_classes - 1.0),
trainable=False)
def call(self, inputs, training):
# In inference, labels are ignored
feature, labels = inputs
cos = self.cos_similarity(feature)
if training:
mask = tf.cast(labels, dtype=cos.dtype)
# Collect cosine values at only false labels
B = (1 - mask)*tf.exp(self.scale*cos)
B_avg = tf.reduce_mean(tf.reduce_sum(B, axis=-1), axis=0)
theta = tf.acos(tf.clip_by_value(cos, -1, 1))
# Collect cosine at true labels
theta_true = tf.reduce_sum(mask*theta, axis=-1)
# get median (=50-percentile)
theta_med = tfp.stats.percentile(theta_true, q=50)
scale = tf.math.log(B_avg) / tf.cos(tf.minimum(np.pi/4, theta_med))
scale = tf.stop_gradient(scale)
logits = scale*cos
self.scale.assign(scale)
return logits
else:
return cos
def get_config(self):
config = super().get_config().copy()
config.update({
'num_classes': self.num_classes
})
return config
class CircleLoss(tf.keras.layers.Layer):
"""
Implementation of https://arxiv.org/abs/2002.10857 (pair-level label version)
"""
def __init__(self, margin=0.25, scale=256, **kwargs):
"""
Args
margin: a float value, margin for the true label (default 0.25)
scale: a float value, final scale value,
stated as gamma in the original paper (default 256)
Returns:
a tf.keras.layers.Layer object, outputs logit values of each class
In the original paper, margin and scale (=gamma) are set depends on tasks
- Face recognition: m=0.25, scale=256 (default)
- Person re-identification: m=0.25, scale=256
- Fine-grained image retrieval: m=0.4, scale=64
"""
super().__init__(**kwargs)
self.margin = margin
self.scale = scale
self._Op = 1 + margin # O_positive
self._On = -margin # O_negative
self._Dp = 1 - margin # Delta_positive
self._Dn = margin # Delta_negative
def call(self, inputs, training):
feature, labels = inputs
x = tf.nn.l2_normalize(feature, axis=-1)
cos = tf.matmul(x, x, transpose_b=True) # (batch_size, batch_size)
if training:
# pairwise version
mask = tf.cast(labels, dtype=cos.dtype)
mask_p = tf.matmul(mask, mask, transpose_b=True)
mask_n = 1 - mask_p
mask_p = mask_p - tf.eye(mask_p.shape[0])
logits_p = - self.scale * tf.nn.relu(self._Op - cos) * (cos - self._Dp)
logits_n = self.scale * tf.nn.relu(cos - self._On) * (cos - self._Dn)
logits_p = tf.where(mask_p == 1, logits_p, -np.inf)
logits_n = tf.where(mask_n == 1, logits_n, -np.inf)
logsumexp_p = tf.reduce_logsumexp(logits_p, axis=-1)
logsumexp_n = tf.reduce_logsumexp(logits_n, axis=-1)
mask_p_row = tf.reduce_max(mask_p, axis=-1)
mask_n_row = tf.reduce_max(mask_n, axis=-1)
logsumexp_p = tf.where(mask_p_row == 1, logsumexp_p, 0)
logsumexp_n = tf.where(mask_n_row == 1, logsumexp_n, 0)
losses = tf.nn.softplus(logsumexp_p + logsumexp_n)
mask_paired = mask_p_row*mask_n_row
losses = mask_paired * losses
return losses
else:
return cos
def get_config(self):
config = super().get_config().copy()
config.update({
'margin': self.margin,
'scale': self.scale
})
return config
class CircleLossCL(tf.keras.layers.Layer):
"""
Implementation of https://arxiv.org/abs/2002.10857 (class-level label version)
"""
def __init__(self, num_classes, margin=0.25, scale=256, **kwargs):
"""
Args
num_classes: an int value, number of target classes
margin: a float value, margin for the true label (default 0.25)
scale: a float value, final scale value,
stated as gamma in the original paper (default 256)
Returns:
a tf.keras.layers.Layer object, outputs logit values of each class
In the original paper, margin and scale (=gamma) are set depends on tasks
- Face recognition: m=0.25, scale=256 (default)
- Person re-identification: m=0.25, scale=256
- Fine-grained image retrieval: m=0.4, scale=64
"""
super().__init__(**kwargs)
self.num_classes = num_classes
self.margin = margin
self.scale = scale
self._Op = 1 + margin # O_positive
self._On = -margin # O_negative
self._Dp = 1 - margin # Delta_positive
self._Dn = margin # Delta_negative
self.cos_similarity = CosineSimilarity(num_classes)
def call(self, inputs, training):
feature, labels = inputs
cos = self.cos_similarity(feature)
if training:
# class-lebel version
mask = tf.cast(labels, dtype=cos.dtype)
alpha_p = tf.nn.relu(self._Op - cos)
alpha_n = tf.nn.relu(cos - self._On)
logits_p = self.scale*alpha_p*(cos - self._Dp)
logits_n = self.scale*alpha_n*(cos - self._Dn)
logits = mask*logits_p + (1-mask)*logits_n
return logits
else:
return cos
def get_config(self):
config = super().get_config().copy()
config.update({
'num_classes': self.num_classes,
'margin': self.margin,
'scale': self.scale
})
return config
def get_margin(margin):
if (margin == "ArcMargin"):
return ArcMargin(n_classes = config["n_classes"], m = 0.1, s = 32)
elif (margin == "ArcFace"):
return ArcFace(num_classes = config["n_classes"], margin = 0.1, scale = 32)
elif (margin == "AdaCos"):
return AdaCos(num_classes = config["n_classes"])
elif (margin == "CircleLossCL"):
return CircleLossCL(num_classes = config["n_classes"], margin = 0.6, scale = 32)
elif (margin == "CosFace"):
return ArcFace(num_classes = config["n_classes"], margin = -0.1, scale = 32)
def get_backbone(backbone, x):
if (hasattr(efn, backbone)):
return GeMPooling(p = 3.0)(getattr(efn, backbone)(weights = "noisy-student", include_top = False)(x))
elif hasattr(tf.keras.applications, backbone):
return GeMPooling(p = 3.0)(getattr(tf.keras.applications, backbone)(weights = "imagenet", include_top = False)(x))
elif hasattr(vit, backbone):
return getattr(vit, backbone)(image_size = (config["image_size"][0], config["image_size"][1]),
pretrained=True,
include_top=False,
pretrained_top=False)(x)
elif "eff" in backbone:
return hub.KerasLayer("gs://cloud-tpu-checkpoints/efficientnet/v2/hub/"+backbone+"-21k-ft1k/feature-vector", trainable=True)(x)
def model_factory(backbone, image_size, embedding_dimensions, margin):
x = Input(shape = (*image_size,), name = 'input')
label = Input(shape = (), name = 'label')
headModel = get_backbone(backbone, x)
headModel = Dense(embedding_dimensions, activation = "linear")(headModel)
headModel = BatchNormalization()(headModel)
headModel = PReLU()(headModel)
headModel = get_margin(margin = margin)([headModel, label])
output = Softmax(dtype='float32')(headModel)
model = tf.keras.models.Model(inputs = [x, label], outputs = [output])
return model
def get_lr_callback(plot=False):
LR_START = config["lr"] * (config["batch_size"] / 256)
LR_MAX = 5 * LR_START
LR_MIN = LR_START/10
LR_RAMPUP_EPOCHS = 5
LR_SUSTAIN_EPOCHS = 0
LR_EXP_DECAY = 0.8
def lrfn(epoch):
if epoch < LR_RAMPUP_EPOCHS:
lr = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START
elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS:
lr = LR_MAX
else:
lr = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN
return lr
if plot:
epochs = list(range(config["epochs"]))
learning_rates = [lrfn(x) for x in epochs]
plt.plot(epochs, learning_rates)
print("Learning rate schedule: {:.3g} to {:.3g} to {:.3g}".format(learning_rates[0], max(learning_rates), learning_rates[-1]))
plt.show()
lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=False)
return lr_callback
class SaveModelCheckpoint(tf.keras.callbacks.Callback):
def __init__(self, path):
self.path = path
def on_epoch_end(self, epoch, logs={}):
self.model.save(os.path.join(self.path, "model_{}_{}_{}.h5".format(config["backbone"], config["margin"], (epoch + 1)%config["total_save"])))
df = pd.read_csv(config["encoded_csv_path"])
FILENAMES = []
for GCS_PATH in config["image_paths"]:
FILENAMES += tf.io.gfile.glob(GCS_PATH + '/train*.tfrec')
TRAINING_FILENAMES, VALIDATION_FILENAMES = train_test_split(FILENAMES,
test_size=config["valid_size"],
random_state=42)
training_groups = [int(re.compile(r"_([0-9]*)\.").search(filename).group(1)) for filename in TRAINING_FILENAMES]
validation_groups = [int(re.compile(r"_([0-9]*)\.").search(filename).group(1)) for filename in VALIDATION_FILENAMES]
n_trn_classes = df[df['group'].isin(training_groups)]['landmark_id_encode'].nunique()
n_val_classes = df[df['group'].isin(validation_groups)]['landmark_id_encode'].nunique()
print(f'The number of unique training classes is {n_trn_classes} of {config["n_classes"]} total classes')
print(f'The number of unique validation classes is {n_val_classes} of {config["n_classes"]} total classes')
STEPS_PER_EPOCH = count_data_items(TRAINING_FILENAMES) // config["batch_size"]
train_dataset = get_training_dataset(TRAINING_FILENAMES, ordered = False, do_aug = True)
valid_dataset = get_validation_dataset(VALIDATION_FILENAMES, ordered = True, prediction = False)
with strategy.scope():
optimizer = Adam(learning_rate = config["lr"])
model = model_factory(margin = config["margin"],
backbone = config["backbone"],
image_size = config["image_size"],
embedding_dimensions = config["embedding_dimensions"])
if (config["weight_path"]):
model.load_weights(config["weight_path"])
model.compile(optimizer = optimizer,
loss = [tf.keras.losses.CategoricalCrossentropy() if (config['margin'] != "ArcMargin") else tf.keras.losses.SparseCategoricalCrossentropy()],
metrics = [tf.keras.metrics.CategoricalAccuracy() if (config['margin'] != "ArcMargin") else tf.keras.metrics.SparseCategoricalAccuracy()])
checkpoint = SaveModelCheckpoint(path = config["save_path"])
lr_callback = get_lr_callback(plot = True)
H = model.fit(train_dataset,
steps_per_epoch = STEPS_PER_EPOCH,
epochs = config["epochs"],
callbacks = [checkpoint, lr_callback],
validation_data = valid_dataset,
initial_epoch = config["last_epoch"],
verbose = 1)
import gc
import os
import csv
import cv2
import sys
import copy
import math
import random
import shutil
import logging
import operator
import threading
import pydegensac
import numpy as np
import albumentations as A
import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.keras.layers import *
from scipy import spatial
from tqdm.notebook import *
from sklearn.metrics import *
try:
from vit_keras import vit
import efficientnet.tfkeras as efn
from ffyytt_tools.metric_learning.metric_learning_layers import *
except:
sys.path.append("../input/ffyytt-tools")
!pip install -qq vit-keras --no-index --find-links=file:///kaggle/input/2021-glr-lib
!pip install -qq efficientnet --no-index --find-links=file:///kaggle/input/2021-glr-lib
from vit_keras import vit
import efficientnet.tfkeras as efn
from ffyytt_tools.metric_learning.metric_learning_layers import *
model_paths = [
"../input/2021-glr-all-best-model/model_efficientnetv2-b3_1024_ArcMargin_0.h5",
"../input/2021-glr-all-best-model/model_efficientnetv2-m_1024_ArcMargin_0.h5",
"../input/2021-glr-all-best-model/model_EfficientNetB4_512_ArcMargin_0.h5",
"../input/2021-glr-all-best-model/model_EfficientNetB6_1024_AdaCos_0.h5",
"../input/2021-glr-craw-weights-13092021/model_vit_b32_AdaCos_0.h5",
"../input/2021-glr-model-temp/model_efficientnetv2-s_AdaCos_1.h5",
"../input/2021-glr-model-temp/model_vit_b32_AdaCos_1.h5",
]
backbones = [
"efficientnetv2-b3",
"efficientnetv2-m",
"EfficientNetB4",
"EfficientNetB6",
"vit_b32",
"efficientnetv2-s",
"vit_b32"
]
embedding_sizes = [
1024,
1024,
512,
1024,
1024,
1024,
1024
]
paddings = [
False,
False,
False,
True,
True,
True,
False,
]
margins = [
"ArcMargin",
"ArcMargin",
"ArcMargin",
"AdaCos",
"AdaCos",
"AdaCos",
"AdaCos",
]
image_sizes = [
(512, 512, 3),
(512, 512, 3),
(512, 512, 3),
(800, 800, 3),
(800, 800, 3),
(800, 800, 3),
(512, 512, 3),
]
ensemble_weight = np.array([1]*len(backbones))
augmentation = A.Compose([
A.HorizontalFlip(p = 0.5)
])
config = {
"n_workers": 4,
"batch_size": 32,
"n_classes": 81313,
"distance_batch": 512,
"NUM_PUBLIC_TRAIN_IMAGES": 1580470,
}
DEBUG = True
LOCAL = True
REMOVE_TOP_GLOBAL = 20
NUM_EMBEDDING_DIMENSIONS = max(embedding_sizes)
TOP_K = 3
NUM_TO_RERANK = 3
MAX_DISTANCE = 0.9
MAX_INLIER_SCORE = 50
HOMOGRAPHY_CONFIDENCE = 0.95
MAX_REPROJECTION_ERROR = 6.0
MAX_RANSAC_ITERATIONS = 900000
def get_margin(margin):
if (margin == "ArcMargin"):
return ArcMargin(n_classes = config["n_classes"], m = 0.1, s = 32)
elif (margin == "ArcFace"):
return ArcFace(num_classes = config["n_classes"], margin = 0.1, scale = 32)
elif (margin == "AdaCos"):
return AdaCos(num_classes = config["n_classes"])
elif (margin == "CircleLossCL"):
return CircleLossCL(num_classes = config["n_classes"], margin = 0.25, scale = 32)
elif (margin == "CosFace"):
return ArcFace(num_classes = config["n_classes"], margin = -0.1, scale = 32)
def get_backbone(backbone, x, image_size):
if (hasattr(efn, backbone)):
return GeMPooling(p = 3.0)(getattr(efn, backbone)(weights = None, include_top = False)(x))
elif hasattr(tf.keras.applications, backbone):
return GeMPooling(p = 3.0)(getattr(tf.keras.applications, backbone)(weights = None, include_top = False)(x))
elif hasattr(vit, backbone):
return getattr(vit, backbone)(image_size = (image_size[0], image_size[1]),
pretrained=False,
include_top=False,
pretrained_top=False)(x)
else:
return hub.KerasLayer("../input/efficientnetv2-tfhub-weight-files/tfhub_models/"+backbone+"/feature_vector", trainable=True)(x)
def model_factory(backbone, image_size, embedding_dimensions, margin):
x = Input(shape = (*image_size,), name = 'input')
label = Input(shape = (), name = 'label')
headModel = get_backbone(backbone, x, image_size)
headModel = Dense(embedding_dimensions, activation = "linear")(headModel)
headModel = BatchNormalization()(headModel)
headModel = PReLU()(headModel)
headModel = get_margin(margin = margin)([headModel, label])
output = Softmax(dtype='float32')(headModel)
model = tf.keras.models.Model(inputs = [x, label], outputs = [output])
return model
global_models = [None]*len(model_paths)
for model_index in trange(len(model_paths)):
model = model_factory(image_size = image_sizes[model_index],
margin = margins[model_index],
backbone = backbones[model_index],
embedding_dimensions = embedding_sizes[model_index])
model.load_weights(model_paths[model_index])
global_models[model_index] = tf.keras.models.Model(inputs = model.input[0],
outputs = model.layers[-4].output)
DELG_MODEL = tf.saved_model.load('../input/delg-saved-models/local_and_global')
DELG_IMAGE_SCALES = tf.convert_to_tensor([0.70710677, 1.0, 1.4142135])
DELG_SCORE_THRESHOLD = tf.constant(175.)
DELG_INPUT = ['input_image:0', 'input_scales:0', 'input_abs_thres:0']
# Local feature extraction:
LOCAL_FEATURE_NUM = tf.constant(1000)
LOCAL_MODEL = DELG_MODEL.prune(DELG_INPUT + ['input_max_feature_num:0'], ['boxes:0', 'features:0'])
def get_image_path(subset, image_id):
return os.path.join(DATASET_DIR, subset, image_id[0], image_id[1], image_id[2], '{}.jpg'.format(image_id))
def load_image_tensor(image_path):
image = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
image = augmentation(image=image)["image"]
#image = cv2.resize(image, (image_size[0], image_size[1]))
image = tf.convert_to_tensor(image)
return image
def extract_local_features(image_path):
image_tensor = load_image_tensor(image_path)
features = LOCAL_MODEL(image_tensor,
DELG_IMAGE_SCALES,
DELG_SCORE_THRESHOLD,
LOCAL_FEATURE_NUM)
keypoints = tf.divide(tf.add(tf.gather(features[0], [0, 1], axis=1),
tf.gather(features[0], [2, 3], axis=1)), 2.0).numpy()
descriptors = tf.nn.l2_normalize(features[1], axis=1, name='l2_normalization').numpy()
return keypoints, descriptors
def get_putative_matching_keypoints(test_keypoints,
test_descriptors,
train_keypoints,
train_descriptors,
max_distance=1.0):
train_descriptor_tree = spatial.cKDTree(train_descriptors)
_, matches = train_descriptor_tree.query(test_descriptors, distance_upper_bound=max_distance)
test_kp_count = test_keypoints.shape[0]
train_kp_count = train_keypoints.shape[0]
test_matching_keypoints = np.array([test_keypoints[i,] for i in range(test_kp_count) if matches[i] != train_kp_count])
train_matching_keypoints = np.array([train_keypoints[matches[i],] for i in range(test_kp_count) if matches[i] != train_kp_count])
return test_matching_keypoints, train_matching_keypoints
def get_num_inliers(test_keypoints, test_descriptors, train_keypoints, train_descriptors):
test_match_kp, train_match_kp = get_putative_matching_keypoints(test_keypoints,
test_descriptors,
train_keypoints,
train_descriptors,
MAX_DISTANCE)
if test_match_kp.shape[0] <= MAX_REPROJECTION_ERROR:
return 0
try:
_, mask = pydegensac.findHomography(test_match_kp,
train_match_kp,
MAX_REPROJECTION_ERROR,
HOMOGRAPHY_CONFIDENCE,
MAX_RANSAC_ITERATIONS)
except np.linalg.LinAlgError: # When det(H)=0, can't invert matrix.
return 0
return int(copy.deepcopy(mask).astype(np.float32).sum())
def get_total_score(num_inliers, global_score):
local_score = min(num_inliers, MAX_INLIER_SCORE) / MAX_INLIER_SCORE
return 0.5*local_score + global_score
def rescore_and_rerank_by_num_inliers(test_image_id, test_index, train_ids_and_scores):
try:
test_image_path = get_image_path('test', test_image_id)
test_keypoints, test_descriptors = extract_local_features(test_image_path)
for i in range(min(len(train_ids_and_scores[test_index]), NUM_TO_RERANK)):
train_image_id, global_score = train_ids_and_scores[test_index][i]
train_image_path = get_image_path('train', train_image_id)
train_keypoints, train_descriptors = extract_local_features(train_image_path)
num_inliers = get_num_inliers(test_keypoints,
test_descriptors,
train_keypoints,
train_descriptors)
total_score = get_total_score(num_inliers, global_score)
train_ids_and_scores[test_index][i] = (train_image_id, total_score)
train_ids_and_scores[test_index].sort(key=lambda x: x[1], reverse=True)
except Exception as e:
print(e)
def thread_call_rescore_and_rerank_by_num_inliers(test_indexs, test_ids, train_ids_and_scores):
threads = []
for test_index, test_id in zip(test_indexs, test_ids):
t = threading.Thread(target=rescore_and_rerank_by_num_inliers,
args=(test_id, test_index, train_ids_and_scores))
t.start()
threads.append(t)
for t in threads:
t.join()
def local_rerank(test_ids, train_ids_and_scores, test_remove):
list_test_id = []
list_test_index = []
for test_index, test_id in enumerate(tqdm(test_ids)):
if (test_id not in test_remove):
list_test_id.append(test_id)
list_test_index.append(test_index)
if (len(list_test_id) == config["n_workers"] or test_index == len(test_ids)-1):
thread_call_rescore_and_rerank_by_num_inliers(list_test_index, list_test_id, train_ids_and_scores)
list_test_id = []
list_test_index = []
return test_remove, train_ids_and_scores
def get_predictions(labelmap):
test_remove = {}
test_ids, global_test_remove, train_ids_and_scores = global_predictions()
test_remove = merge_dict(test_remove, global_test_remove, "global")
gc.collect()
if (LOCAL):
local_test_remove, train_ids_and_scores = local_rerank(test_ids, train_ids_and_scores, test_remove)
verification_predictions = get_prediction_map(test_ids, train_ids_and_scores, labelmap)
return verification_predictions, test_remove
def get_prediction_map(test_ids, train_ids_and_scores, labelmap):
prediction_map = dict()
for test_index, test_id in enumerate(test_ids):
aggregate_scores = {}
for index, (train_id, score) in enumerate(train_ids_and_scores[test_index][:TOP_K]):
label = labelmap[train_id]
if label not in aggregate_scores:
aggregate_scores[label] = 0
aggregate_scores[label] += score
label, score = max(aggregate_scores.items(), key=operator.itemgetter(1))
prediction_map[test_id] = {'score': score, 'class': label}
return prediction_map
def merge_dict(d, s, value):
for key in s:
if (key in d):
d[key] += "_" + value
else:
d[key] = value
return d
def my_glob(path, filetype):
ids = []
filepaths = []
for root, dirs, files in os.walk(path):
for file in files:
if file.endswith(filetype):
filepaths.append(os.path.join(root, file))
ids.append(file[:-4])
return ids, filepaths
def load_labelmap():
with open(TRAIN_LABELMAP_PATH, mode='r') as csv_file:
csv_reader = csv.DictReader(csv_file)
labelmap = {row['id']: row['landmark_id'] for row in csv_reader}
return labelmap
def save_submission_csv(predictions=None, test_remove = None):
if predictions is None:
shutil.copyfile(os.path.join(DATASET_DIR, 'sample_submission.csv'), 'submission.csv')
return True
with open('submission.csv', 'w') as submission_csv:
csv_writer = csv.DictWriter(submission_csv, fieldnames=['id', 'landmarks'])
csv_writer.writeheader()
for image_id, prediction in predictions.items():
if (image_id in test_remove):
label = prediction['class']
score = 0
csv_writer.writerow({'id': image_id, 'landmarks': f''})
else:
label = prediction['class']
score = prediction['score']
csv_writer.writerow({'id': image_id, 'landmarks': f'{label} {score}'})
INPUT_DIR = os.path.join('..', 'input')
DATASET_DIR = os.path.join(INPUT_DIR, 'landmark-recognition-2021')
TEST_IMAGE_DIR = os.path.join(DATASET_DIR, 'test')
TRAIN_IMAGE_DIR = os.path.join(DATASET_DIR, 'train')
TRAIN_LABELMAP_PATH = os.path.join(DATASET_DIR, 'train.csv')
labelmap = load_labelmap()
num_training_images = len(labelmap.keys())
print(f'Found {num_training_images} training images')
if (num_training_images == config["NUM_PUBLIC_TRAIN_IMAGES"]) and not DEBUG:
print("Copying sample submission")
save_submission_csv()
else:
if (num_training_images != config["NUM_PUBLIC_TRAIN_IMAGES"]):
DEBUG = False
verification_predictions, test_remove = get_predictions(labelmap)
save_submission_csv(verification_predictions, test_remove)