"""

Calculates the Frechet Inception Distance between two distributions, using chosen feature extractor model.

RadImageNet Model source: https://github.com/BMEII-AI/RadImageNet

RadImageNet InceptionV3 weights (original, broken since 11.07.2023): https://drive.google.com/file/d/1p0q9AhG3rufIaaUE1jc2okpS8sdwN6PU

RadImageNet InceptionV3 weights (for medigan, updated link 11.07.2023): https://drive.google.com/drive/folders/1lGFiS8_a5y28l4f8zpc7fklwzPJC-gZv

Usage:

    python fid.py dir1 dir2 

"""

import argparse

import os

import cv2

import numpy as np

import tensorflow as tf

import tensorflow_gan as tfgan

import wget

from tensorflow.keras.applications import InceptionV3

from tensorflow.keras.applications.inception_v3 import preprocess_input

img_size = 299

batch_size = 64

num_batches = 1

RADIMAGENET_URL = "https://drive.google.com/uc?id=1uvJHLG1K71Qzl7Km4JMpNOwE7iTjN8g9"

RADIMAGENET_WEIGHTS = "RadImageNet-InceptionV3_notop.h5"

IMAGENET_TFHUB_URL = "https://tfhub.dev/tensorflow/tfgan/eval/inception/1"

def parse_args() -> argparse.Namespace:

    parser = argparse.ArgumentParser(

        description="Calculates the Frechet Inception Distance between two distributions using RadImageNet model."

    )

    parser.add_argument(

        "dataset_path_1",

        type=str,

        help="Path to images from first dataset",

    )

    parser.add_argument(

        "dataset_path_2",

        type=str,

        help="Path to images from second dataset",

    )

    parser.add_argument(

        "--model",

        type=str,

        default="imagenet",

        help="Use RadImageNet feature extractor for FID calculation",

    )

    parser.add_argument(

        "--lower_bound",

        action="store_true",

        help="Calculate lower bound of FID using the 50/50 split of images from dataset_path_1",

    )

    parser.add_argument(

        "--normalize_images",

        action="store_true",

        help="Normalize images from both datasources using min and max of each sample",

    )

    args = parser.parse_args()

    return args

def load_images(directory, normalize=False, split=False, limit=None):

    """

    Loads images from the given directory.

    If split is True, then half of the images is loaded to one array and the other half to another.

    """

    if split:

        subset_1 = []

        subset_2 = []

    else:

        images = []

    for count, filename in enumerate(os.listdir(directory)):

        if filename.lower().endswith((".png", ".jpg", ".jpeg")):

            img = cv2.imread(os.path.join(directory, filename))

            img = cv2.resize(img, (img_size, img_size), interpolation=cv2.INTER_LINEAR)

            if normalize:

                img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX)

            if len(img.shape) > 2 and img.shape[2] == 4:

                img = img[:, :, :3]

            if len(img.shape) == 2:

                img = np.stack([img] * 3, axis=2)

            if split:

                if count % 2 == 0:

                    subset_1.append(img)

                else:

                    subset_2.append(img)

            else:

                images.append(img)

        if count == limit:

            break

    if split:

        subset_1 = preprocess_input(np.array(subset_1))

        subset_2 = preprocess_input(np.array(subset_2))

        return subset_1, subset_2

    else:

        images = preprocess_input(np.array(images))

        return images

def check_model_weights(model_name):

    """

    Checks if the model weights are available and download them if not.

    """

    model_weights_path = None

    if model_name == "radimagenet":

        model_weights_path = RADIMAGENET_WEIGHTS

        if not os.path.exists(RADIMAGENET_WEIGHTS):

            print("Downloading RadImageNet InceptionV3 model:")

            wget.download(

                RADIMAGENET_URL,

                model_weights_path,

            )

            print("\n")

        return model_weights_path

def _radimagenet_fn(images):

    """

    Get RadImageNet inception v3 model

    """

    model = InceptionV3(

        weights=RADIMAGENET_WEIGHTS,

        input_shape=(img_size, img_size, 3),

        include_top=False,

        pooling="avg",

    )

    output = model(images)

    output = tf.nest.map_structure(tf.keras.layers.Flatten(), output)

    return output

def get_classifier_fn(model_name="imagenet"):

    """

    Get model as TF function for optimized inference.

    """

    check_model_weights(model_name)

    if model_name == "radimagenet":

        return _radimagenet_fn

    elif model_name == "imagenet":

        return tfgan.eval.classifier_fn_from_tfhub(IMAGENET_TFHUB_URL, "pool_3", True)

    else:

        raise ValueError("Model {} not recognized".format(model_name))

def calculate_fid(

    directory_1,

    directory_2,

    model_name,

    lower_bound=False,

    normalize_images=False,

):

    """

    Calculates the Frechet Inception Distance between two distributions using chosen feature extractor model.

    """

    limit = min(len(os.listdir(directory_1)), len(os.listdir(directory_2)))

    if lower_bound:

        images_1, images_2 = load_images(directory_1, split=True, limit=limit)

    else:

        images_1 = load_images(directory_1, limit=limit, normalize=normalize_images)

        images_2 = load_images(directory_2, limit=limit, normalize=normalize_images)

    fid = tfgan.eval.frechet_classifier_distance(

        images_1, images_2, get_classifier_fn(model_name)

    )

    return fid

if __name__ == "__main__":

    args = parse_args()

    directory_1 = args.dataset_path_1

    directory_2 = args.dataset_path_2

    lower_bound = args.lower_bound

    normalize_images = args.normalize_images

    model_name = args.model

    fid = calculate_fid(

        directory_1=directory_1,

        directory_2=directory_2,

        model_name=model_name,

        lower_bound=lower_bound,

        normalize_images=normalize_images,

    )

    if lower_bound:

        print("Lower bound FID {}: {}".format(model_name, fid))

    else:

        print("FID {}: {}".format(model_name, fid))