ccip.py

import json
import os.path
from functools import lru_cache
from typing import Union, List

import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download, HfFileSystem

try:
    from typing import Literal
except (ModuleNotFoundError, ImportError):
    from typing_extensions import Literal

from imgutils.data import MultiImagesTyping, load_images, ImageTyping
from imgutils.utils import open_onnx_model

hf_fs = HfFileSystem()


def _normalize(data, mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)):
    mean, std = np.asarray(mean), np.asarray(std)
    return (data - mean[:, None, None]) / std[:, None, None]


def _preprocess_image(image: Image.Image, size: int = 384):
    image = image.resize((size, size), resample=Image.BILINEAR)
    # noinspection PyTypeChecker
    data = np.array(image).transpose(2, 0, 1).astype(np.float32) / 255.0
    data = _normalize(data)

    return data


@lru_cache()
def _open_feat_model(model):
    return open_onnx_model(hf_hub_download(
        f'deepghs/ccip_onnx',
        f'{model}/model_feat.onnx',
    ))


@lru_cache()
def _open_metric_model(model):
    return open_onnx_model(hf_hub_download(
        f'deepghs/ccip_onnx',
        f'{model}/model_metrics.onnx',
    ))


@lru_cache()
def _open_metrics(model):
    with open(hf_hub_download(f'deepghs/ccip_onnx', f'{model}/metrics.json'), 'r') as f:
        return json.load(f)


@lru_cache()
def _open_cluster_metrics(model):
    with open(hf_hub_download(f'deepghs/ccip_onnx', f'{model}/cluster.json'), 'r') as f:
        return json.load(f)


_VALID_MODEL_NAMES = [
    os.path.basename(os.path.dirname(file)) for file in
    hf_fs.glob('deepghs/ccip_onnx/*/model.ckpt')
]
_DEFAULT_MODEL_NAMES = 'ccip-caformer-24-randaug-pruned'


def ccip_extract_feature(image: ImageTyping, size: int = 384, model: str = _DEFAULT_MODEL_NAMES):
    """
    Extracts the feature vector of the character from the given anime image.

    :param image: The anime image containing a single character.
    :type image: ImageTyping

    :param size: The size of the input image to be used for feature extraction. (default: ``384``)
    :type size: int

    :param model: The name of the model to use for feature extraction. (default: ``ccip-caformer-24-randaug-pruned``)
                  The available model names are: ``ccip-caformer-24-randaug-pruned``,
                  ``ccip-caformer-6-randaug-pruned_fp32``, ``ccip-caformer-5_fp32``.
    :type model: str

    :return: The feature vector of the character.
    :rtype: numpy.ndarray

    Examples::
        >>> from imgutils.metrics import ccip_extract_feature
        >>>
        >>> feat = ccip_extract_feature('ccip/1.jpg')
        >>> feat.shape, feat.dtype
        ((768,), dtype('float32'))
    """
    return ccip_batch_extract_features([image], size, model)[0]


def ccip_batch_extract_features(images: MultiImagesTyping, size: int = 384, model: str = _DEFAULT_MODEL_NAMES):
    """
    Extracts the feature vectors of multiple images using the specified model.

    :param images: The input images from which to extract the feature vectors.
    :type images: MultiImagesTyping

    :param size: The size of the input image to be used for feature extraction. (default: ``384``)
    :type size: int

    :param model: The name of the model to use for feature extraction. (default: ``ccip-caformer-24-randaug-pruned``)
                  The available model names are: ``ccip-caformer-24-randaug-pruned``,
                  ``ccip-caformer-6-randaug-pruned_fp32``, ``ccip-caformer-5_fp32``.
    :type model: str

    :return: The feature vectors of the input images.
    :rtype: numpy.ndarray

    Examples::
        >>> from imgutils.metrics import ccip_batch_extract_features
        >>>
        >>> feat = ccip_batch_extract_features(['ccip/1.jpg', 'ccip/2.jpg', 'ccip/6.jpg'])
        >>> feat.shape, feat.dtype
        ((3, 768), dtype('float32'))
    """
    images = load_images(images, mode='RGB')
    data = np.stack([_preprocess_image(item, size=size) for item in images]).astype(np.float32)
    output, = _open_feat_model(model).run(['output'], {'input': data})
    return output


_FeatureOrImage = Union[ImageTyping, np.ndarray]


def _p_feature(x: _FeatureOrImage, size: int = 384, model: str = _DEFAULT_MODEL_NAMES):
    if isinstance(x, np.ndarray):  # if feature
        return x
    else:  # is image or path
        return ccip_extract_feature(x, size, model)


def ccip_default_threshold(model: str = _DEFAULT_MODEL_NAMES) -> float:
    """
    Retrieves the default threshold value obtained from model metrics in the Hugging Face model repository.

    :param model: The name of the model to use for feature extraction. (default: ``ccip-caformer-24-randaug-pruned``)
                  The available model names are: ``ccip-caformer-24-randaug-pruned``,
                  ``ccip-caformer-6-randaug-pruned_fp32``, ``ccip-caformer-5_fp32``.
    :type model: str

    :return: The default threshold value obtained from model metrics.
    :rtype: float

    Examples::
        >>> from imgutils.metrics import ccip_default_threshold
        >>>
        >>> ccip_default_threshold()
        0.17847511429108218
        >>> ccip_default_threshold('ccip-caformer-6-randaug-pruned_fp32')
        0.1951224011983088
        >>> ccip_default_threshold('ccip-caformer-5_fp32')
        0.18397327797685215
    """
    return _open_metrics(model)['threshold']


def ccip_difference(x: _FeatureOrImage, y: _FeatureOrImage,
                    size: int = 384, model: str = _DEFAULT_MODEL_NAMES) -> float:
    """
    Calculates the difference value between two anime characters based on their images or feature vectors.

    :param x: The image or feature vector of the first anime character.
    :type x: Union[ImageTyping, np.ndarray]

    :param y: The image or feature vector of the second anime character.
    :type y: Union[ImageTyping, np.ndarray]

    :param size: The size of the input image to be used for feature extraction. (default: ``384``)
    :type size: int

    :param model: The name of the model to use for feature extraction. (default: ``ccip-caformer-24-randaug-pruned``)
                  The available model names are: ``ccip-caformer-24-randaug-pruned``,
                  ``ccip-caformer-6-randaug-pruned_fp32``, ``ccip-caformer-5_fp32``.
    :type model: str

    :return: The difference value between the two anime characters.
    :rtype: float

    Examples::
        >>> from imgutils.metrics import ccip_difference
        >>>
        >>> ccip_difference('ccip/1.jpg', 'ccip/2.jpg')  # same character
        0.16583099961280823
        >>>
        >>> # different characters
        >>> ccip_difference('ccip/1.jpg', 'ccip/6.jpg')
        0.42947039008140564
        >>> ccip_difference('ccip/1.jpg', 'ccip/7.jpg')
        0.4037521779537201
        >>> ccip_difference('ccip/2.jpg', 'ccip/6.jpg')
        0.4371533691883087
        >>> ccip_difference('ccip/2.jpg', 'ccip/7.jpg')
        0.40748104453086853
        >>> ccip_difference('ccip/6.jpg', 'ccip/7.jpg')
        0.392294704914093
    """
    return ccip_batch_differences([x, y], size, model)[0, 1].item()


def ccip_batch_differences(images: List[_FeatureOrImage],
                           size: int = 384, model: str = _DEFAULT_MODEL_NAMES) -> np.ndarray:
    """
    Calculates the pairwise differences between a given list of images or feature vectors representing anime characters.

    :param images: The list of images or feature vectors representing anime characters.
    :type images: List[Union[ImageTyping, np.ndarray]]

    :param size: The size of the input image to be used for feature extraction. (default: ``384``)
    :type size: int

    :param model: The name of the model to use for feature extraction. (default: ``ccip-caformer-24-randaug-pruned``)
                  The available model names are: ``ccip-caformer-24-randaug-pruned``,
                  ``ccip-caformer-6-randaug-pruned_fp32``, ``ccip-caformer-5_fp32``.
    :type model: str

    :return: The matrix of pairwise differences between the given images or feature vectors.
    :rtype: np.ndarray

    Examples::
        >>> from imgutils.metrics import ccip_batch_differences
        >>>
        >>> ccip_batch_differences(['ccip/1.jpg', 'ccip/2.jpg', 'ccip/6.jpg', 'ccip/7.jpg'])
        array([[6.5350548e-08, 1.6583106e-01, 4.2947042e-01, 4.0375218e-01],
               [1.6583106e-01, 9.8025822e-08, 4.3715334e-01, 4.0748104e-01],
               [4.2947042e-01, 4.3715334e-01, 3.2675274e-08, 3.9229470e-01],
               [4.0375218e-01, 4.0748104e-01, 3.9229470e-01, 6.5350548e-08]],
              dtype=float32)
    """
    input_ = np.stack([_p_feature(img, size, model) for img in images]).astype(np.float32)
    output, = _open_metric_model(model).run(['output'], {'input': input_})
    return output