app.py

# If TF version is not understood by tfimm requirements, try this:
#try:
#    import tfimm
#except ModuleNotFoundError:
#    !pip install --no-deps tfimm timm
#    import timm
#    import tfimm

import os
import glob
from shutil import rmtree
from pathlib import Path

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_comp_embeddings, get_test_embedding, get_confidence


# YOLOv5 parameters
yolo_input_size = 384
versions = ('2_v108', '4_v109', '0_int6', '1_v110', '3_v111')
score_thr = 0.025
iou_thr = 0.6
max_det = 1
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
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,
    'efnv1b7_colab225': 1,
    'efnv1b7_colab197': 0,
    'efnv1b7_colab227': 5,
    'efnv1b7_v72': 6,
    'efnv1b7_colab229': 9,
    'efnv1b6_colab217': 5,
    'efnv1b6_colab218': 6,
    '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_colab199': 1,
    'convnext_base_384_in22ft1k_v68': 0,
    'convnext_base_384_in22ft1k_colab220': 9,
    'convnext_base_384_in22ft1k_colab201': 3,  # new
}
cfg_files = [hf_hub_download(modelbox, f'{x}_config.json') for x in rst_names]
emb_files = [hf_hub_download(modelbox, f'{x}_emb.npz') for x in rst_names]
rst_files = [hf_hub_download(modelbox, f'{x}.h5') for x in rst_names]
use_folds = [use_fold[x] for x in rst_names]
n_models = len(rst_names)


def fast_yolo_crop(image):
    rmtree(working / 'labels', ignore_errors=True)
    rmtree(working / 'results_ensemble', ignore_errors=True)

    mpimg.imsave(yolo_source, image)

    detect.run(weights=checkpoint_files[4:],
               source=yolo_source,
               data='data/dataset.yaml',
               imgsz=yolo_input_size,
               conf_thres=score_thr,
               iou_thres=iou_thr,
               max_det=max_det,
               save_txt=False,
               save_conf=False,
               save_crop=True,
               exist_ok=True,
               name=str(working / 'results_ensemble'))

    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]
    species = Path(cropped).parent.name
    cropped_image = mpimg.imread(cropped)
    return cropped_image, species.replace('_', ' ')


# Preload embeddings for known individuals
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, use_fold in zip(cfg_files, rst_files, use_folds):
    cfg = get_cfg(cfg_file)
    assert cfg.FOLD_TO_RUN == use_fold
    cfg.pretrained = None  # avoid weight downloads
    if isinstance(cfg.IMAGE_SIZE, int):
        cfg.IMAGE_SIZE = (cfg.IMAGE_SIZE, cfg.IMAGE_SIZE)
    sizes.append(cfg.IMAGE_SIZE)
    model, embed_model = get_model(cfg)
    model.load_weights(rst_file)
    print(f"\nWeights loaded from {rst_file}")
    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)


def pred_fn(image, fake=False):
    if fake:
        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 = "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(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
    normalized_similarities = 1 - normalized_distances

    min_similarity = normalized_similarities.min()
    max_similarity = normalized_similarities.max()
    confidence = get_confidence(max_similarity, threshold)

    print(f"Similarities: {min_similarity:.4f} ... {max_similarity:.4f}")
    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"
    else:
        response_str = f"This looks like a {species}, but I have not seen this individual before ({confidence:.3f} confidence).\n" \
                       "Go submit your photo on www.happywhale.com!"

    return cropped_image, response_str


examples = [str(image_root / f'negative{i:03d}.jpg') for i in range(3)]

demo = gr.Interface(fn=pred_fn, inputs="image", outputs=["image", "text"],
                    examples=examples)
demo.launch()