fix imports, cleanup

app.py

@@ -10,15 +10,14 @@ import os
 import glob
 from shutil import rmtree
 from pathlib import Path
-from subprocess import run
-import json
 
 import gradio as gr
 from huggingface_hub import hf_hub_download
+import matplotlib.image as mpimg
 from yolov5 import detect
 import numpy as np
 from tensorflow.keras import backend as K
-from utils import get_model, get_cfg, get_embeddings, get_comp_embeddings, get_test_embedding
+from utils import get_model, get_cfg, get_comp_embeddings, get_test_embedding, get_confidence
 
 
 # YOLOv5 parameters
@@ -31,14 +30,13 @@ working = Path(os.getcwd())
 modelbox = "yellowdolphin/happywhale-models"
 checkpoint_files = [hf_hub_download(modelbox, f'yolov5_l6_{yolo_input_size}_fold{x}.pt') for x in versions]
 image_root = working / 'images'
+yolo_source = str(image_root / 'testimage.jpg')
 
 
 # Individual identifier parameters
 max_distance = 0.865
 normalize_similarity = None  # test-train, None
-gamma = 0.4
-threshold = 0.09951 if (normalize_similarity == 'test-train') else 0.6 # 0.381
-knn = 300
+threshold = 0.09951 if (normalize_similarity == 'test-train') else 0.6  # 0.381
 rst_names = 'convnext_base_384_in22ft1k_colab220 efnv1b7_colab216 hub_efnv2xl_v73'.split()
 use_fold = {
     'efnv1b7_colab216': 4,
@@ -49,12 +47,12 @@ use_fold = {
     'efnv1b7_colab229': 9,
     'efnv1b6_colab217': 5,
     'efnv1b6_colab218': 6,
-    'hub_efnv2xl_colab221': 8, 
+    'hub_efnv2xl_colab221': 8,
     'hub_efnv2xl_v69': 2,
     'hub_efnv2xl_v73': 0,
     'efnv1b6_colab226': 2,
     'hub_efnv2l_v70': 3,
-    'hub_efnv2l_colab200': 2, 
+    'hub_efnv2l_colab200': 2,
     'hub_efnv2l_colab199': 1,
     'convnext_base_384_in22ft1k_v68': 0,
     'convnext_base_384_in22ft1k_colab220': 9,
@@ -73,7 +71,6 @@ def fast_yolo_crop(image):
 
     mpimg.imsave(yolo_source, image)
 
-    #print(f"\nInference on best {len(checkpoint_files[5:])} models with detect.py ...")
     detect.run(weights=checkpoint_files[4:],
                source=yolo_source,
                data='data/dataset.yaml',
@@ -87,7 +84,6 @@ def fast_yolo_crop(image):
                exist_ok=True,
                name=str(working / 'results_ensemble'))
 
-    #print(f"YOLOv5 inference finished in {(perf_counter() - t0) / 60:.2f} min")
     cropped = sorted(glob(f'{working}/results_ensemble/crops/*/{Path(yolo_source).name}'))
     assert len(cropped) == 1, f'{len(cropped)} maritime species detected'
     cropped = cropped[0]
@@ -102,7 +98,7 @@ comp_embeddings = get_comp_embeddings(emb_files, use_folds)
 # Preload embedding models, input sizes
 K.clear_session()
 embed_models, sizes = [], []
-for cfg_file, rst_file, npz_file in zip (cfg_files, rst_files, emb_files):
+for cfg_file, rst_file, npz_file in zip(cfg_files, rst_files, emb_files):
     cfg = get_cfg(cfg_file)
     assert cfg.FOLD_TO_RUN == use_fold[npz_file]
     cfg.pretrained = None  # avoid weight downloads
@@ -112,7 +108,7 @@ for cfg_file, rst_file, npz_file in zip (cfg_files, rst_files, emb_files):
     model, embed_model = get_model(cfg)
     model.load_weights(rst_file)
     print(f"\nWeights loaded from {rst_file}")
-    print(f"input_size {scaled_img.shape[:2]}, fold {cfg.FOLD_TO_RUN}, arch {cfg.arch_name}, ",
+    print(f"input_size {cfg.IMAGE_SIZE}, fold {cfg.FOLD_TO_RUN}, arch {cfg.arch_name}, ",
           f"DATASET {cfg.DATASET}, dropout_ps {cfg.dropout_ps}, subcenters {cfg.subcenters}")
     embed_models.append(embed_model)
 
@@ -122,13 +118,13 @@ def pred_fn(image, fake=False):
         x0, x1 = (int(f * image.shape[0]) for f in (0.2, 0.8))
         y0, y1 = (int(f * image.shape[1]) for f in (0.2, 0.8))
         cropped_image = image[x0:x1, y0:y1, :]
-        response_str = f"This looks like a common dolphin, but I have not seen this individual before (0.834 confidence).\n" \
+        response_str = "This looks like a common dolphin, but I have not seen this individual before (0.834 confidence).\n" \
                        "Go submit your photo on www.happywhale.com!"
         return cropped_image, response_str
 
     cropped_image, species = fast_yolo_crop(image)
-    test_embedding = get_test_embedding(embed_models, sizes)
-    
+    test_embedding = get_test_embedding(cropped_image, embed_models, sizes)
+
     cosine_similarity = np.dot(comp_embeddings, test_embedding[0]) / n_models
     cosine_distances = 1 - cosine_similarity
     normalized_distances = cosine_distances / max_distance
@@ -139,8 +135,8 @@ def pred_fn(image, fake=False):
     confidence = get_confidence(max_similarity, threshold)
 
     print(f"Similarities: {min_similarity:.4f} ... {max_similarity:.4f}")
-    print(f"Threshold:", threshold)
-    
+    print(f"Threshold: {threshold}")
+
     if max_similarity > threshold:
         response_str = f"This looks like a {species} I have seen before ({confidence:.3f} confidence).\n" \
                        "You might find its previous encounters on www.happywhale.com"
@@ -150,6 +146,7 @@ def pred_fn(image, fake=False):
 
     return cropped_image, response_str
 
+
 examples = [str(image_root / f'negative{i:03d}') for i in range(3)]
 
 demo = gr.Interface(fn=pred_fn, inputs="image", outputs=["image", "text"],

utils.py

@@ -1,14 +1,18 @@
 import math
+import json
 
 import numpy as np
 import tensorflow as tf
 import tfimm
-import efficientnet
 import efficientnet.tfkeras as efnv1
 import keras_efficientnet_v2 as efnv2
 import tensorflow_hub as hub
 
 
+embedding_size = 1024
+n_images = 51033 + 27956
+
+
 class DotDict(dict):
     """dot.notation access to dictionary attributes
 
@@ -19,7 +23,7 @@ class DotDict(dict):
     __setattr__ = dict.__setitem__
     __delattr__ = dict.__delitem__
 
-    
+
 def get_cfg(json_file):
     json_file = str(json_file)
     config_dict = json.load(open(json_file))
@@ -32,7 +36,7 @@ def get_embeddings(img, embed_model):
     return embeddings
 
 
-# Train embeddings have to be re-ordered: embeddings were concatenated (train, valid) 
+# Train embeddings have to be re-ordered: embeddings were concatenated (train, valid)
 # in the training notebook and the valid fold is different for each ensemble model.
 FOLDS = 10
 shards, n_total = [], 0
@@ -42,6 +46,7 @@ for fold in range(10):
     n_total += n_img
 assert n_total == 51033
 
+
 def get_train_idx(use_fold):
     "Return embedding index that restores the order of images in the tfrec files."
     train_folds = [i for i in range(10) if i % FOLDS != use_fold]
@@ -53,14 +58,12 @@ def get_train_idx(use_fold):
     for fold in folds:
         train_idx.append(shards[fold])
     train_idx = np.concatenate(train_idx)
-    
+
     return np.argsort(train_idx)
 
 
 def get_comp_embeddings(emb_files, use_folds):
     "Load embeddings for competition images [n_images, embedding_size]"
-    embedding_size = 1024
-    n_images = 51033 + 27956
     comp_embeddings = []
 
     for npz_file, use_fold in zip(emb_files, use_folds):
@@ -68,7 +71,7 @@ def get_comp_embeddings(emb_files, use_folds):
         d = np.load(str(npz_file))
         comp_train_emb = d['train']
         comp_test_emb = d['test']
-        
+
         # Restore original order of comp_train_emb, targets (use targets as fingerprint-check)
         comp_train_idx = get_train_idx(use_fold)
         comp_train_emb = comp_train_emb[comp_train_idx, :]
@@ -85,22 +88,20 @@ def get_comp_embeddings(emb_files, use_folds):
     return np.concatenate(comp_embeddings, axis=1)
 
 
-def get_test_embedding(embed_models, sizes):
-    test_embedding, similarities = [], []
+def get_test_embedding(image, embed_models, sizes):
+    test_embedding = []
 
     for embed_model, size in zip(embed_models, sizes):
         # Get model input
-        scaled_img = tf.image.resize(img, size)
-        scaled_img = tf.cast(scaled_img, tf.float32) / 255.0
-        #print("test image normalized and resized to", scaled_img.shape[:2])
+        scaled_image = tf.image.resize(image, size)
+        scaled_image = tf.cast(scaled_image, tf.float32) / 255.0
 
         # Get embedding for test image
-        test_emb = get_embeddings(scaled_img, embed_model)  # shape: [1, embedding_size]
+        test_emb = get_embeddings(scaled_image, embed_model)  # shape: [1, embedding_size]
         assert test_emb.shape == (1, embedding_size)
 
         # Normalize embeddings
         test_emb_norm = np.linalg.norm(test_emb, axis=1)
-        #print("test_emb norm: ", test_emb_norm[0])
         test_emb /= test_emb_norm[:, None]
 
         test_embedding.append(test_emb)
@@ -108,6 +109,23 @@ def get_test_embedding(embed_models, sizes):
     return np.concatenate(test_embedding, axis=1)  # [1, embedding_size]
 
 
+def p2logit(x):
+    return np.log(x / (1 - x))
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+
+def get_confidence(similarity, threshold):
+    "Calculate confidence in known/unknown prediction"
+    if similarity <= 0:
+        return 0
+    logit_sim = p2logit(similarity)
+    logit_threshold = p2logit(threshold)
+    return sigmoid(abs(logit_sim - logit_threshold))
+
+
 class ArcMarginProductSubCenter(tf.keras.layers.Layer):
     '''
     Implements large margin arc distance.
@@ -204,56 +222,39 @@ TFHUB = {
 
 def get_model(cfg):
     aux_arcface = False  # Chris Deotte suggested this
-    if cfg.head == 'arcface2':
-        head = ArcMarginPenaltyLogists
-    elif cfg.head == 'arcface':
+    if cfg.head == 'arcface':
         head = ArcMarginProductSubCenter
-    elif cfg.head == 'addface':
-        head = AddMarginProductSubCenter
     else:
         assert False, "INVALID HEAD"
 
     if cfg.adaptive_margin:
-        # define adaptive margins depending on class frequencies (dynamic margins)
-        df = pd.read_csv(f'{project_dir}/train.csv')
-        fewness = df['individual_id'].value_counts().sort_index() ** (-1/4)
-        fewness -= fewness.min()
-        fewness /= fewness.max() - fewness.min()
-        adaptive_margin = cfg.margin_min + fewness * (cfg.margin_max - cfg.margin_min)
-
-        # align margins with targets
-        splits_path = '/kaggle/input/happywhale-splits'
-        with open (f'{splits_path}/individual_ids.json', "r") as f:
-            target_encodings = json.loads(f.read())  # individual_id: index
-        individual_ids = pd.Series(target_encodings).sort_values().index.values
-        adaptive_margin = adaptive_margin.loc[individual_ids].values.astype(np.float32)
-        
+        raise NotImplementedError
+
     if cfg.arch_name.startswith('efnv1'):
-        EFN = {'efnv1b0': efnv1.EfficientNetB0, 'efnv1b1': efnv1.EfficientNetB1, 
+        EFN = {'efnv1b0': efnv1.EfficientNetB0, 'efnv1b1': efnv1.EfficientNetB1,
                'efnv1b2': efnv1.EfficientNetB2, 'efnv1b3': efnv1.EfficientNetB3,
-               'efnv1b4': efnv1.EfficientNetB4, 'efnv1b5': efnv1.EfficientNetB5, 
+               'efnv1b4': efnv1.EfficientNetB4, 'efnv1b5': efnv1.EfficientNetB5,
                'efnv1b6': efnv1.EfficientNetB6, 'efnv1b7': efnv1.EfficientNetB7}
 
     if cfg.arch_name.startswith('efnv2'):
         EFN = {'efnv2s': efnv2.EfficientNetV2S, 'efnv2m': efnv2.EfficientNetV2M,
                'efnv2l': efnv2.EfficientNetV2L, 'efnv2xl': efnv2.EfficientNetV2XL}
 
-
-    with strategy.scope():
+    with tf.distribute.get_strategy().scope():
 
         margin = head(
-            n_classes = cfg.N_CLASSES,
-            s = 30,
-            m = adaptive_margin if cfg.adaptive_margin else 0.3,
-            k = cfg.subcenters or 1,
-            easy_margin = False,
-            name=f'head/{cfg.head}', 
+            n_classes=cfg.N_CLASSES,
+            s=30,
+            m=0.3,
+            k=cfg.subcenters or 1,
+            easy_margin=False,
+            name=f'head/{cfg.head}',
             dtype='float32')
 
-        inp = tf.keras.layers.Input(shape = [*cfg.IMAGE_SIZE, 3], name = 'inp1')
-        label = tf.keras.layers.Input(shape = (), name = 'inp2')
+        inp = tf.keras.layers.Input(shape=[*cfg.IMAGE_SIZE, 3], name='inp1')
+        label = tf.keras.layers.Input(shape=(), name='inp2')
         if aux_arcface:
-            label2 = tf.keras.layers.Input(shape = (), name = 'inp3')
+            label2 = tf.keras.layers.Input(shape=(), name='inp3')
 
         if cfg.arch_name.startswith('efnv1'):
             x = EFN[cfg.arch_name](weights=cfg.pretrained, include_top=False)(inp)
@@ -263,8 +264,6 @@ def get_model(cfg):
                 embed = tf.keras.layers.Flatten()(x)
                 embed = tf.keras.layers.Dropout(0.1)(embed)
                 embed = tf.keras.layers.Dense(1024)(embed)
-            elif cfg.pool == 'gem':
-                embed = GeMPoolingLayer(train_p=True)(x)
             elif cfg.pool == 'concat':
                 embed = tf.keras.layers.concatenate([tf.keras.layers.GlobalAveragePooling2D()(x),
                                                      tf.keras.layers.GlobalAveragePooling2D()(x)])
@@ -272,7 +271,7 @@ def get_model(cfg):
                 embed = tf.keras.layers.GlobalMaxPooling2D()(x)
             else:
                 embed = tf.keras.layers.GlobalAveragePooling2D()(x)
-            
+
         elif cfg.arch_name.startswith('efnv2'):
             x = EFN[cfg.arch_name](input_shape=(None, None, 3), num_classes=0,
                                    pretrained=cfg.pretrained)(inp)
@@ -282,8 +281,6 @@ def get_model(cfg):
                 embed = tf.keras.layers.Flatten()(x)
                 embed = tf.keras.layers.Dropout(0.1)(embed)
                 embed = tf.keras.layers.Dense(1024)(embed)
-            elif cfg.pool == 'gem':
-                embed = GeMPoolingLayer(train_p=True)(x)
             elif cfg.pool == 'concat':
                 embed = tf.keras.layers.concatenate([tf.keras.layers.GlobalAveragePooling2D()(x),
                                                      tf.keras.layers.GlobalAveragePooling2D()(x)])
@@ -297,15 +294,11 @@ def get_model(cfg):
             url = TFHUB[cfg.arch_name]
             model = hub.KerasLayer(url, trainable=True)
             embed = model(inp)
-            #print(f"{cfg.arch_name} from tfhub")
             assert cfg.pool in [None, False, 'avg', ''], 'tfhub model, no custom pooling supported!'
-            
+
         elif cfg.arch_name in tfimm.list_models(pretrained="timm"):
-            #print(f"{cfg.arch_name} from tfimm")
-            #embed = tfimm.create_model(cfg.arch_name, pretrained="timm", nb_classes=0)(inp)
             embed = tfimm.create_model(cfg.arch_name, pretrained=None, nb_classes=0)(inp)
-            # create_model(nb_classes=0) includes pooling as last layer
-        
+
         if len(cfg.dropout_ps) > 0:
             # Chris Deotte posted model code without Dropout/FC1 after pooling
             embed = tf.keras.layers.Dropout(cfg.dropout_ps[0])(embed)
@@ -314,32 +307,17 @@ def get_model(cfg):
         x = margin([embed, label])
 
         output = tf.keras.layers.Softmax(dtype='float32', name='arc' if cfg.aux_loss else None)(x)
-        
-        if cfg.aux_loss and aux_arcface:
-            # Use 2nd arcface head for species (aux loss)
-            head2 = ArcMarginProductSubCenter
-            margin2 = head(
-                n_classes = cfg.n_species,
-                s = 30,
-                m = 0.3,
-                k = 1,
-                easy_margin = False,
-                name=f'auxhead/{cfg.head}', 
-                dtype='float32')
-            aux_features = margin2([embed, label2])
-            aux_output = tf.keras.layers.Softmax(dtype='float32', name='aux')(aux_features)
 
-        elif cfg.aux_loss:
+        if cfg.aux_loss:
             aux_features = tf.keras.layers.Dense(cfg.n_species)(embed)
             aux_output = tf.keras.layers.Softmax(dtype='float32', name='aux')(aux_features)
         inputs = [inp, label, label2] if (cfg.aux_loss and aux_arcface) else [inp, label]
         outputs = (output, aux_output) if cfg.aux_loss else [output]
-        
+
         model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
         embed_model = tf.keras.models.Model(inputs=inp, outputs=embed)
-        
-        opt = tf.keras.optimizers.Adam(learning_rate=cfg.LR)
+
         if cfg.FREEZE_BATCH_NORM:
-            freeze_BN(model)
-        
+            raise NotImplementedError
+
         return model, embed_model