iBioHash2024

Config

import os

import pickle

import random

import sklearn

import numpy as np

import pandas as pd

import tensorflow as tf

import matplotlib.pyplot as plt

from tqdm.notebook import *

from sklearn.neighbors import *

from sklearn.decomposition import *

from sklearn.model_selection import *

try:

import farmhash

import keras_cv_attention_models

import tensorflow_addons as tfa

except:

!pip install -qq cityhash

!pip install -qq tensorflow-addons

!pip install -qq keras-cv-attention-models

import farmhash

import keras_cv_attention_models

import tensorflow_addons as tfa

try:

tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection

except ValueError:

tpu = None

gpus = tf.config.experimental.list_logical_devices("GPU")

if tpu:

tf.config.experimental_connect_to_cluster(tpu)

tf.tpu.experimental.initialize_tpu_system(tpu)

strategy = tf.distribute.experimental.TPUStrategy(tpu)

tf.config.set_soft_device_placement(True)

print('Running on TPU ', tpu.master())

elif len(gpus) > 0:

strategy = tf.distribute.MirroredStrategy(gpus)

print('Running on ', len(gpus), ' GPU(s) ')

else:

strategy = tf.distribute.get_strategy()

print('Running on CPU')

print("Number of accelerators: ", strategy.num_replicas_in_sync)

AUTO = tf.data.experimental.AUTOTUNE

config = {

"seed": 1213,

"lr": 1e-5,

"epochs": 10,

"batch_size": 32 * strategy.num_replicas_in_sync,

"n_classes": 1000,

"image_size": [224, 224, 3],

"hashLength": 2048,

"data_paths": ['gs://kds-3f5ce2e56559d17c577c6d9f0c1acd2179da3bc0c1f87ae24d66f478', 'gs://kds-6d0cfb649153417b33f05a9913c59a358a3c03706a8fd6c5f523b8fc', 'gs://kds-0a7e8b42af254bd7711b5670f5aeb0eeda8d169841722ea59680d3d4', 'gs://kds-d27f1d2f32cff13b0ce97a6b912eb44a68e8ee3128504dcb52e23a3e', 'gs://kds-7a36360336e83a60d74c0196158fed9d6fa70cd342dc96383b7f6e2a', 'gs://kds-f99762a40174159b2d5b1caedff9d3e42758db9134c67405c68ce52d', 'gs://kds-af8ef10c715e13cb364f2655d5188aedfbb874b71e250d2773dade89', 'gs://kds-b0b80433ab5beabf70aca68659672322018f637039ae27fbd9ee9bf1'],

"save_path": "./",

"backbones": ["EfficientNetV2S", "beit.BeitV2BasePatch16"] #["EfficientNetV2M", "beit.BeitV2BasePatch16", "davit.DaViT_S"]

}

def seed_everything(seed):

random.seed(seed)

np.random.seed(seed)

tf.random.set_seed(seed)

seed_everything(config["seed"])

Data

def train_transform(image, label):

image = tf.expand_dims(image, axis=0)

if (tf.random.uniform([1]) < 0.5):

image = tf.image.random_flip_left_right(image)

if (tf.random.uniform([1]) < 0.2):

image = tf.image.random_brightness(image, 0.08)

if (tf.random.uniform([1]) < 0.3):

image = tfa.image.rotate(image, 1.0*tf.random.uniform([1]) - 0.5)

return tf.squeeze(image, axis=0), label

def margin_format(image, label):

return {'image': image, 'label': label}, label

def onehot(data, label):

label = tf.one_hot(label, tf.constant(config["n_classes"]), axis =-1, dtype=tf.float32)

return data, label

def decode_image(image_data):

image = tf.image.decode_jpeg(image_data, channels = 3)

image = tf.image.resize_with_pad(image, target_width = config["image_size"][0], target_height = config["image_size"][1])

return image

def read_tfrecord(example):

TFREC_FORMAT = {

"image": tf.io.FixedLenFeature([], tf.string),

'label': tf.io.FixedLenFeature([], tf.int64),

}

example = tf.io.parse_single_example(example, TFREC_FORMAT)

image = decode_image(example['image'])

label = tf.cast(example['label'], dtype = tf.int64)

return image, label

def load_dataset(filenames, ordered):

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 get_train_dataset(filenames):

dataset = load_dataset(filenames, ordered = False).repeat().shuffle(config["seed"])

dataset = dataset.map(train_transform, num_parallel_calls = AUTO)

dataset = dataset.map(onehot, num_parallel_calls = AUTO)

dataset = dataset.map(margin_format, num_parallel_calls = AUTO)

dataset = dataset.batch(config["batch_size"])

dataset = dataset.prefetch(AUTO)

return dataset

def get_valid_dataset(filenames):

dataset = load_dataset(filenames, ordered = True)

dataset = dataset.map(onehot, num_parallel_calls = AUTO)

dataset = dataset.map(margin_format, num_parallel_calls = AUTO)

dataset = dataset.batch(config["batch_size"])

dataset = dataset.prefetch(AUTO)

return dataset

def get_test_dataset(filenames):

dataset = load_dataset(filenames, ordered = True)

dataset = dataset.map(margin_format, num_parallel_calls = AUTO)

dataset = dataset.batch(config["batch_size"])

dataset = dataset.prefetch(AUTO)

return dataset

Model

class Margin(tf.keras.layers.Layer):

def __init__(self, num_classes, margin = 0.1, scale=32, **kwargs):

super().__init__(**kwargs)

self.scale = scale

self.margin = margin

self.num_classes = num_classes

def build(self, input_shape):

self.W = self.add_weight(shape=(self.num_classes, input_shape[0][-1]), initializer='glorot_uniform', trainable=True)

def build_hash(self, x, scale = 10.0):

return tf.nn.sigmoid(scale*tf.nn.l2_normalize(x, axis = 1))

def hamming_distance(self, feature):

x = self.build_hash(feature, self.scale)

w = self.build_hash(self.W, self.scale)

x = tf.tile(tf.expand_dims(feature, 2), [1, 1, self.W.shape[0]])

w = tf.transpose(self.W)

return tf.clip_by_value(tf.reduce_sum(tf.math.abs(x - w), axis = 1), 1e-4, 48)

def logits_hamming(self, feature, labels):

distance = self.hamming_distance(feature)

mr = tf.random.normal(shape = tf.shape(distance), mean = self.margin, stddev = 0.1*self.margin)

distance_add = distance + mr

mask = tf.cast(labels, dtype=distance.dtype)

logits = self.scale/(mask*distance_add + (1-mask)*distance)

return logits

def distance(self, feature):

x = tf.nn.l2_normalize(feature, axis=1)

w = tf.nn.l2_normalize(self.W, axis=1)

x = tf.tile(tf.expand_dims(x, 2), [1, 1, self.W.shape[0]])

w = tf.transpose(w)

return tf.reduce_sum( tf.math.pow( x - w, 2 ), axis=1)

def logits_distance(self, feature, labels):

distance = self.distance(feature)

mr = tf.random.normal(shape = tf.shape(distance), mean = self.margin, stddev = 0.1*self.margin)

distance_add = distance + mr

mask = tf.cast(labels, dtype=distance.dtype)

logits = self.scale/(mask*distance_add + (1-mask)*distance)

return logits

def cosine(self, feature):

x = tf.nn.l2_normalize(feature, axis=1)

w = tf.nn.l2_normalize(self.W, axis=1)

cos = tf.matmul(x, tf.transpose(w))

return cos

def logits_cosine(self, feature, labels):

cosine = self.cosine(feature)

mr = tf.random.normal(shape = tf.shape(cosine), mean = self.margin, stddev = 0.1*self.margin)

theta = tf.acos(tf.clip_by_value(cosine, -1, 1))

cosine_add = tf.math.cos(theta + mr)

mask = tf.cast(labels, dtype=cosine.dtype)

logits = mask*cosine_add + (1-mask)*cosine

return logits*self.scale

def call(self, inputs, training):

feature, labels = inputs

if training:

logits = self.logits_cosine(feature, labels)

else:

logits = self.cosine(feature)

return logits

def get_config(self):

config = super().get_config().copy()

config.update({

'scale': self.scale,

'margin': self.margin,

'num_classes': self.num_classes,

})

return config

def get_backbone(backbone_name, x):

if hasattr(tf.keras.applications, backbone_name):

headModel = tf.keras.layers.Lambda(lambda data: tf.keras.applications.imagenet_utils.preprocess_input(tf.cast(data, tf.float32), mode="tf"))(x)

return tf.keras.layers.GlobalAveragePooling2D()(getattr(tf.keras.applications, backbone_name)(weights = "imagenet", include_top = False)(headModel))

else:

backbone = getattr(getattr(keras_cv_attention_models, backbone_name.split(".")[0]), backbone_name.split(".")[1])(num_classes=0)

headModel = tf.keras.layers.Lambda(lambda data: tf.keras.applications.imagenet_utils.preprocess_input(tf.cast(data, tf.float32), mode="torch"))(x)

backbone.trainable = True

if "beit" in backbone_name:

return backbone(headModel)

return tf.keras.layers.GlobalAveragePooling2D()(backbone(headModel))

def model_factory(backbones, n_classes):

image = tf.keras.layers.Input(shape = (None, None, 3), dtype=tf.uint8, name = 'image')

label = tf.keras.layers.Input(shape = (), name = 'label', dtype = tf.int64)

features = [get_backbone(backbone, image) for backbone in backbones]

headModel = tf.keras.layers.Concatenate(name = "concat")(features)

headModel = tf.keras.layers.Dense(config["hashLength"], activation = "linear", name = "feature")(headModel)

margin = Margin(num_classes = n_classes, name = "margin")([headModel, label])

output = tf.keras.layers.Softmax(dtype=tf.float32, name = "output")(margin)

model = tf.keras.models.Model(inputs = [image, label], outputs = [output])

return model

Callbacks

%%writefile generate_submit_csv.py

import argparse

import numpy as np

parser = argparse.ArgumentParser()

parser.add_argument('--query', type=str, default='submit_query.csv', help='48-bit hash code file for query set.')

parser.add_argument('--gallery', type=str, default='submit_gallery.csv', help='48-bit hash code file for gallery set.')

parser.add_argument('--submit', type=str, default='submit.csv', help='Final submit csv file.')

parser.add_argument('--k', type=int, default=100, help='Topk == 100.')

args = parser.parse_args()

k = args.k

query_code_path = args.query

gallery_code_path = args.gallery

submit_path = args.submit

query_code = []

query_images = []

with open(query_code_path, 'r') as f:

f.readline()

for line in f:

image, code = line[:-1].split(',')

code = [-1. if i == 0 else 1. for i in list(map(float, list(code[1:-1])))]

query_code.append(code)

query_images.append(image)

query_code = np.array(query_code)

gallery_code = []

gallery_images = []

with open(gallery_code_path, 'r') as f:

f.readline()

for line in f:

image, code = line[:-1].split(',')

code = [-1. if i == 0 else 1. for i in list(map(float, list(code[1:-1])))]

gallery_code.append(code)

gallery_images.append(image[:-4])

gallery_code = np.array(gallery_code)

with open(submit_path, 'w') as f:

f.write('Id,Predicted\n')

for i, q in enumerate(query_code):

hamming_dist = q.shape[0] - np.dot(q, gallery_code.T)

index = np.argsort(hamming_dist)[:k]

retrieval_images = [gallery_images[j] for j in index]

f.write(query_images[i] + ',')

f.write(' '.join(retrieval_images) + '\n')

print(f'\r writing {i + 1}/{query_code.shape[0]}', end='')

class Evaluation(tf.keras.callbacks.Callback):

def __init__(self, g_data, q_data, v_g_data):

self.g_data = g_data

self.q_data = q_data

self.v_g_data = v_g_data

def bitdigest(self, digest):

return [str(x) for i, x in enumerate(digest.tolist())]

def to_id(self, ids):

return [str(x) + ".jpg" for x in ids]

def process_vgallery(self, model, epoch, steps = None):

feature, vg_id = model.predict(self.v_g_data, verbose = 1, steps = steps)

feature = np.array(tf.nn.l2_normalize(feature, axis=1)).astype(np.float16)

return feature, vg_id

# vg_id.tofile(f"{config['save_path']}VGid_{epoch}")

# np.array(feature).astype(np.float16).tofile(f"{config['save_path']}VGfeature_{epoch}")

def process_gallery(self, model, epoch, steps = None):

feature, g_id = model.predict(self.g_data, verbose = 1, steps = steps)

feature = np.array(tf.nn.l2_normalize(feature, axis=1)).astype(np.float16)

return feature, g_id

# g_id.tofile(f"{config['save_path']}Gid_{epoch}")

# np.array(feature).astype(np.float16).tofile(f"{config['save_path']}Gfeature_{epoch}")

def process_query(self, model, epoch, steps = None):

feature, q_id = model.predict(self.q_data, verbose = 1, steps = steps)

feature = np.array(tf.nn.l2_normalize(feature, axis=1)).astype(np.float16)

return feature, q_id

# q_id.tofile(f"{config['save_path']}Qid_{epoch}")

# np.array(feature).astype(np.float16).tofile(f"{config['save_path']}Qfeature_{epoch}")

def gen_sub(self, epoch):

return os.popen(f"python3 generate_submit_csv.py --gallery {config['save_path']}G_{epoch}.csv --query {config['save_path']}Q_{epoch}.csv --submit {config['save_path']}S_{epoch}.csv").read()

def on_epoch_end(self, epoch, logs={}):

model = tf.keras.models.Model(inputs = self.model.inputs,

outputs = [self.model.get_layer('feature').output, self.model.inputs[1]])

with open(f"{config['save_path']}data_{epoch}", 'wb') as handle:

pickle.dump([self.process_query(model, epoch), self.process_gallery(model, epoch), self.process_vgallery(model, epoch)], handle, protocol=pickle.HIGHEST_PROTOCOL)

# self.gen_sub(epoch)

class SaveModel(tf.keras.callbacks.Callback):

def __init__(self, path):

self.path = path

def on_epoch_end(self, epoch, logs={}):

pass

# self.model.save(self.path + "model.keras")

Run

DATA_FILENAMES = []

TEST_G_FILENAMES = []

TEST_Q_FILENAMES = []

for gcs_path in config["data_paths"]:

DATA_FILENAMES += tf.io.gfile.glob(gcs_path + '/*BIO*.tfrec')

TEST_G_FILENAMES += tf.io.gfile.glob(gcs_path + '/FGVC11*Test_G*.tfrec')

TEST_Q_FILENAMES += tf.io.gfile.glob(gcs_path + '/FGVC11*Test_Q*.tfrec')

TRAINING_FILENAMES, VALIDATION_FILENAMES = train_test_split(DATA_FILENAMES, test_size=0.02, random_state = config["seed"])

sample_dataset = get_test_dataset(random.choices(DATA_FILENAMES, k = 1))

train_dataset = get_train_dataset(TRAINING_FILENAMES)

valid_dataset = get_valid_dataset(VALIDATION_FILENAMES)

valid_G_dataset = get_test_dataset(VALIDATION_FILENAMES)

test_G_dataset = get_test_dataset(sorted(TEST_G_FILENAMES))

test_Q_dataset = get_test_dataset(sorted(TEST_Q_FILENAMES))

evaluation = Evaluation(test_G_dataset, test_Q_dataset, valid_G_dataset)

counter = 0

plt.figure()

f, axarr = plt.subplots(4,4)

for data in train_dataset:

for image in data[0]['image']:

axarr[counter//4, counter%4].imshow(image/256)

axarr[counter//4, counter%4].axis('off')

counter += 1

if (counter == 16):

break

if (counter == 16):

break

plt.show()

with strategy.scope():

model = model_factory(backbones = config["backbones"],

n_classes = config["n_classes"])

optimizer = tf.keras.optimizers.Adam(learning_rate = config["lr"])

model.compile(optimizer = optimizer,

loss = [tf.keras.losses.CategoricalCrossentropy()],

metrics = [tf.keras.metrics.CategoricalAccuracy(name = "ACC@1"),

tf.keras.metrics.TopKCategoricalAccuracy(k = 10, name = "ACC@10"),

tf.keras.metrics.TopKCategoricalAccuracy(k = 50, name = "ACC@50"),

])

savemodel = SaveModel(path = config['save_path'])

# if os.path.isfile("model.keras"):

# model = tf.keras.models.load_model("model.keras", safe_mode=False, custom_objects={"Hash": Hash, "Margin": Margin})

H = model.fit(train_dataset, verbose = 1,

validation_data = valid_dataset,

callbacks = [savemodel, evaluation],

steps_per_epoch = 1000,

epochs = config["epochs"])