diff --git a/docs/source/api_doc/metrics/dbaesthetic.rst b/docs/source/api_doc/metrics/dbaesthetic.rst
new file mode 100644
index 00000000000..1d1f8d51277
--- /dev/null
+++ b/docs/source/api_doc/metrics/dbaesthetic.rst
@@ -0,0 +1,15 @@
+imgutils.metrics.dbaesthetic
+=================================
+
+.. currentmodule:: imgutils.metrics.dbaesthetic
+
+.. automodule:: imgutils.metrics.dbaesthetic
+
+
+anime_dbaesthetic
+--------------------------------------
+
+.. autofunction:: anime_dbaesthetic
+
+
+
diff --git a/docs/source/api_doc/metrics/dbaesthetic/best.jpg b/docs/source/api_doc/metrics/dbaesthetic/best.jpg
new file mode 100644
index 00000000000..c88c9aa705b
Binary files /dev/null and b/docs/source/api_doc/metrics/dbaesthetic/best.jpg differ
diff --git a/docs/source/api_doc/metrics/dbaesthetic/good.jpg b/docs/source/api_doc/metrics/dbaesthetic/good.jpg
new file mode 100644
index 00000000000..124b8a0f757
Binary files /dev/null and b/docs/source/api_doc/metrics/dbaesthetic/good.jpg differ
diff --git a/docs/source/api_doc/metrics/dbaesthetic/great.jpg b/docs/source/api_doc/metrics/dbaesthetic/great.jpg
new file mode 100644
index 00000000000..487b1f8da09
Binary files /dev/null and b/docs/source/api_doc/metrics/dbaesthetic/great.jpg differ
diff --git a/docs/source/api_doc/metrics/dbaesthetic/low.jpg b/docs/source/api_doc/metrics/dbaesthetic/low.jpg
new file mode 100644
index 00000000000..5dd4104460f
Binary files /dev/null and b/docs/source/api_doc/metrics/dbaesthetic/low.jpg differ
diff --git a/docs/source/api_doc/metrics/dbaesthetic/masterpiece.jpg b/docs/source/api_doc/metrics/dbaesthetic/masterpiece.jpg
new file mode 100644
index 00000000000..a3108fad030
Binary files /dev/null and b/docs/source/api_doc/metrics/dbaesthetic/masterpiece.jpg differ
diff --git a/docs/source/api_doc/metrics/dbaesthetic/normal.jpg b/docs/source/api_doc/metrics/dbaesthetic/normal.jpg
new file mode 100644
index 00000000000..2dc844350f4
Binary files /dev/null and b/docs/source/api_doc/metrics/dbaesthetic/normal.jpg differ
diff --git a/docs/source/api_doc/metrics/dbaesthetic/worst.jpg b/docs/source/api_doc/metrics/dbaesthetic/worst.jpg
new file mode 100644
index 00000000000..e75a6666436
Binary files /dev/null and b/docs/source/api_doc/metrics/dbaesthetic/worst.jpg differ
diff --git a/docs/source/api_doc/metrics/dbaesthetic_benchmark.plot.py b/docs/source/api_doc/metrics/dbaesthetic_benchmark.plot.py
new file mode 100644
index 00000000000..1cf70c3e7d1
--- /dev/null
+++ b/docs/source/api_doc/metrics/dbaesthetic_benchmark.plot.py
@@ -0,0 +1,38 @@
+import random
+
+from benchmark import BaseBenchmark, create_plot_cli
+from imgutils.metrics import anime_dbaesthetic
+from imgutils.metrics.dbaesthetic import _MODEL
+
+_MODEL_NAMES = _MODEL.classifier.model_names
+
+
+class DBAestheticBenchmark(BaseBenchmark):
+ def __init__(self, model_name: str):
+ BaseBenchmark.__init__(self)
+ self.model_name = model_name
+
+ def load(self):
+ from imgutils.metrics.dbaesthetic import _MODEL
+ _MODEL.classifier._open_model(self.model_name)
+ _MODEL._get_xy_samples(self.model_name)
+
+ def unload(self):
+ from imgutils.metrics.dbaesthetic import _MODEL
+ _MODEL.clear()
+
+ def run(self):
+ image_file = random.choice(self.all_images)
+ _ = anime_dbaesthetic(image_file)
+
+
+if __name__ == '__main__':
+ create_plot_cli(
+ [
+ (model_name, DBAestheticBenchmark(model_name))
+ for model_name in _MODEL_NAMES
+ ],
+ title='Benchmark for Danbooru-Based Aesthetic Models',
+ run_times=10,
+ try_times=20,
+ )()
diff --git a/docs/source/api_doc/metrics/dbaesthetic_benchmark.plot.py.svg b/docs/source/api_doc/metrics/dbaesthetic_benchmark.plot.py.svg
new file mode 100644
index 00000000000..8caacd0526f
--- /dev/null
+++ b/docs/source/api_doc/metrics/dbaesthetic_benchmark.plot.py.svg
@@ -0,0 +1,2499 @@
+
+
+
diff --git a/docs/source/api_doc/metrics/dbaesthetic_full.plot.py b/docs/source/api_doc/metrics/dbaesthetic_full.plot.py
new file mode 100644
index 00000000000..3b76566548c
--- /dev/null
+++ b/docs/source/api_doc/metrics/dbaesthetic_full.plot.py
@@ -0,0 +1,19 @@
+import os.path
+
+from imgutils.metrics import anime_dbaesthetic
+from imgutils.metrics.dbaesthetic import _LABELS
+from plot import image_plot
+
+if __name__ == '__main__':
+ items = []
+ for label in _LABELS:
+ image = os.path.join('dbaesthetic', f'{label}.jpg')
+ l, p, s = anime_dbaesthetic(image, fmt=('label', 'percentile', 'score'))
+ display_name = f'{l}\nScore: {s:.2f}/{6.0:.2f}\nPercentile: {p:.3f}'
+ items.append((image, display_name))
+
+ image_plot(
+ *items,
+ columns=4,
+ figsize=(11, 9),
+ )
diff --git a/docs/source/api_doc/metrics/dbaesthetic_full.plot.py.svg b/docs/source/api_doc/metrics/dbaesthetic_full.plot.py.svg
new file mode 100644
index 00000000000..2cc79331faa
--- /dev/null
+++ b/docs/source/api_doc/metrics/dbaesthetic_full.plot.py.svg
@@ -0,0 +1,1157 @@
+
+
+
diff --git a/docs/source/api_doc/metrics/index.rst b/docs/source/api_doc/metrics/index.rst
index 400f0ce5621..b59bd11ed3b 100644
--- a/docs/source/api_doc/metrics/index.rst
+++ b/docs/source/api_doc/metrics/index.rst
@@ -11,5 +11,6 @@ imgutils.metrics
aesthetic
ccip
+ dbaesthetic
lpips
psnr_
diff --git a/imgutils/metrics/__init__.py b/imgutils/metrics/__init__.py
index 474561cd1bc..487e7679c60 100644
--- a/imgutils/metrics/__init__.py
+++ b/imgutils/metrics/__init__.py
@@ -4,5 +4,6 @@
"""
from .aesthetic import *
from .ccip import *
+from .dbaesthetic import *
from .lpips import *
from .psnr_ import *
diff --git a/imgutils/metrics/aesthetic.py b/imgutils/metrics/aesthetic.py
index fb3fa9556a8..81ccc712074 100644
--- a/imgutils/metrics/aesthetic.py
+++ b/imgutils/metrics/aesthetic.py
@@ -10,14 +10,20 @@
.. image:: aesthetic_benchmark.plot.py.svg
:align: center
+
+ .. warning::
+ These model is deprecated due to the poor effectiveness.
+ Please use `imgutils.metrics.aesthetic.anime_dbaesthetic` for better evaluation.
"""
from functools import lru_cache
import cv2
import numpy as np
from PIL import Image
+from deprecation import deprecated
from huggingface_hub import hf_hub_download
+from ..config.meta import __VERSION__
from ..data import ImageTyping, load_image
from ..utils import open_onnx_model
@@ -47,6 +53,8 @@ def _preprocess(image: Image.Image):
return img_input[np.newaxis, :]
+@deprecated(deprecated_in='0.4.2', removed_in='1.0.0', current_version=__VERSION__,
+ details='Deprecated due to the low effectiveness.')
def get_aesthetic_score(image: ImageTyping):
"""
Overview:
diff --git a/imgutils/metrics/dbaesthetic.py b/imgutils/metrics/dbaesthetic.py
new file mode 100644
index 00000000000..7b221e88dd5
--- /dev/null
+++ b/imgutils/metrics/dbaesthetic.py
@@ -0,0 +1,261 @@
+"""
+Overview:
+ A tool for assessing the aesthetic quality of anime images using a pre-trained model,
+ based on danbooru dataset and metadata analysis result of
+ `HakuBooru `_ by
+ `KohakuBlueleaf `_.
+
+ .. image:: dbaesthetic_full.plot.py.svg
+ :align: center
+
+ This is an overall benchmark of all the operations in aesthetic models:
+
+ .. image:: dbaesthetic_benchmark.plot.py.svg
+ :align: center
+
+"""
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+from imgutils.data import ImageTyping
+from imgutils.generic import ClassifyModel
+
+__all__ = [
+ 'anime_dbaesthetic',
+]
+
+_DEFAULT_MODEL_NAME = 'swinv2pv3_v0_448_ls0.2_x'
+_REPO_ID = 'deepghs/anime_aesthetic'
+_LABELS = ["worst", "low", "normal", "good", "great", "best", "masterpiece"]
+_DEFAULT_LABEL_MAPPING = {
+ 'masterpiece': 0.95,
+ 'best': 0.85,
+ 'great': 0.75,
+ 'good': 0.5,
+ 'normal': 0.25,
+ 'low': 0.1,
+ 'worst': 0.0,
+}
+
+
+def _value_replace(v, mapping):
+ """
+ Replaces values in a data structure using a mapping dictionary.
+
+ :param v: The input data structure.
+ :type v: Any
+ :param mapping: A dictionary mapping values to replacement values.
+ :type mapping: Dict
+ :return: The modified data structure.
+ :rtype: Any
+ """
+ if isinstance(v, (list, tuple)):
+ return type(v)([_value_replace(vitem, mapping) for vitem in v])
+ elif isinstance(v, dict):
+ return type(v)({key: _value_replace(value, mapping) for key, value in v.items()})
+ else:
+ try:
+ _ = hash(v)
+ except TypeError: # pragma: no cover
+ return v
+ else:
+ return mapping.get(v, v)
+
+
+class AestheticModel:
+ """
+ A model for assessing the aesthetic quality of anime images.
+ """
+
+ def __init__(self, repo_id: str):
+ """
+ Initializes an AestheticModel instance.
+
+ :param repo_id: The repository ID of the aesthetic assessment model.
+ :type repo_id: str
+ """
+ self.repo_id = repo_id
+ self.classifier = ClassifyModel(repo_id)
+ self.cached_samples: Dict[str, Tuple] = {}
+
+ def get_aesthetic_score(self, image: ImageTyping, model_name: str) -> Tuple[float, Dict[str, float]]:
+ """
+ Calculates the aesthetic score and confidence for an anime image.
+
+ :param image: The input anime image.
+ :type image: ImageTyping
+ :param model_name: The name of the aesthetic assessment model to use.
+ :type model_name: str
+ :return: A tuple containing the aesthetic score and confidence.
+ :rtype: Tuple[float, Dict[str, float]]
+ """
+ scores = self.classifier.predict_score(image, model_name)
+ return sum(scores[label] * i for i, label in enumerate(_LABELS)), scores
+
+ def _get_xy_samples(self, model_name: str):
+ """
+ Retrieves cached samples for aesthetic assessment.
+
+ :param model_name: The name of the aesthetic assessment model.
+ :type model_name: str
+ :return: Cached samples for aesthetic assessment.
+ :rtype: Tuple[Tuple[np.ndarray, float, float], Tuple[np.ndarray, float, float]]
+ """
+ if model_name not in self.cached_samples:
+ stacked = np.load(hf_hub_download(
+ repo_id=self.repo_id,
+ repo_type='model',
+ filename=f'{model_name}/samples.npz',
+ ))['arr_0']
+ x, y = stacked[0], stacked[1]
+ self.cached_samples[model_name] = ((x, x.min(), x.max()), (y, y.min(), y.max()))
+ return self.cached_samples[model_name]
+
+ def score_to_percentile(self, score: float, model_name: str) -> float:
+ """
+ Converts an aesthetic score to a percentile rank.
+
+ :param score: The aesthetic score.
+ :type score: float
+ :param model_name: The name of the aesthetic assessment model to use.
+ :type model_name: str
+ :return: The percentile rank corresponding to the given score.
+ :rtype: float
+ """
+ (x, x_min, x_max), (y, y_min, y_max) = self._get_xy_samples(model_name)
+ idx = np.searchsorted(x, np.clip(score, a_min=x_min, a_max=x_max))
+ if idx < x.shape[0] - 1:
+ x0, y0 = x[idx], y[idx]
+ x1, y1 = x[idx + 1], y[idx + 1]
+ if np.isclose(x1, x0):
+ return y[idx]
+ else:
+ return np.clip((score - x0) / (x1 - x0) * (y1 - y0) + y0, a_min=y_min, a_max=y_max)
+ else:
+ return y[idx]
+
+ @classmethod
+ def percentile_to_label(cls, percentile: float, mapping: Optional[Dict[str, float]] = None) -> str:
+ """
+ Converts a percentile rank to an aesthetic label.
+
+ :param percentile: The percentile rank.
+ :type percentile: float
+ :param mapping: A dictionary mapping labels to percentile thresholds.
+ :type mapping: Optional[Dict[str, float]]
+ :return: The aesthetic label corresponding to the given percentile rank.
+ :rtype: str
+ """
+ mapping = mapping or _DEFAULT_LABEL_MAPPING
+ for label, threshold in sorted(mapping.items(), key=lambda x: (-x[1], x[0])):
+ if percentile >= threshold:
+ return label
+ else:
+ raise ValueError(f'No label for unknown percentile {percentile:.3f}.') # pragma: no cover
+
+ def get_aesthetic(self, image: ImageTyping, model_name: str, fmt=('label', 'percentile')):
+ """
+ Analyzes the aesthetic quality of an anime image and returns the results in the specified format.
+
+ :param image: The input anime image.
+ :type image: ImageTyping
+ :param model_name: The name of the aesthetic assessment model to use.
+ :type model_name: str
+ :param fmt: The format of the output.
+ :type fmt: Tuple[str, ...]
+ :return: A dictionary containing the aesthetic assessment results.
+ :rtype: Dict[str, float]
+ """
+ score, confidence = self.get_aesthetic_score(image, model_name)
+ percentile = self.score_to_percentile(score, model_name)
+ label = self.percentile_to_label(percentile)
+ return _value_replace(
+ v=fmt,
+ mapping={
+ 'label': label,
+ 'percentile': percentile,
+ 'score': score,
+ 'confidence': confidence,
+ }
+ )
+
+ def clear(self):
+ """
+ Clears the internal state of the AestheticModel instance.
+ """
+ self.classifier.clear()
+ self.cached_samples.clear()
+
+
+_MODEL = AestheticModel(_REPO_ID)
+
+
+def anime_dbaesthetic(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME,
+ fmt=('label', 'percentile')):
+ """
+ Analyzes the aesthetic quality of an anime image using a pre-trained model.
+
+ :param image: The input anime image.
+ :type image: ImageTyping
+ :param model_name: The name of the aesthetic assessment model to use. Default is _DEFAULT_MODEL_NAME.
+ :type model_name: str
+ :param fmt: The format of the output. Default is ('label', 'percentile').
+ :type fmt: Tuple[str, ...]
+ :return: A dictionary containing the aesthetic assessment results.
+ :rtype: Dict[str, float]
+
+ Examples::
+ >>> from imgutils.metrics import anime_dbaesthetic
+ >>>
+ >>> anime_dbaesthetic('masterpiece.jpg')
+ ('masterpiece', 0.9831666690063624)
+ >>> anime_dbaesthetic('best.jpg')
+ ('best', 0.8810615667538594)
+ >>> anime_dbaesthetic('great.jpg')
+ ('great', 0.8225559148288356)
+ >>> anime_dbaesthetic('good.jpg')
+ ('good', 0.591020403706702)
+ >>> anime_dbaesthetic('normal.jpg')
+ ('normal', 0.2888798940585766)
+ >>> anime_dbaesthetic('low.jpg')
+ ('low', 0.243279223969715)
+ >>> anime_dbaesthetic('worst.jpg')
+ ('worst', 0.005268185993767627)
+
+ * Custom format
+
+ >>> anime_dbaesthetic('masterpiece.jpg', fmt=('label', 'percentile', 'score'))
+ ('masterpiece', 0.9831666690063624, 5.275707557797432)
+ >>> anime_dbaesthetic('best.jpg', fmt=('label', 'percentile', 'score'))
+ ('best', 0.8810615667538594, 4.7977807857096195)
+ >>> anime_dbaesthetic('great.jpg', fmt=('label', 'percentile', 'score'))
+ ('great', 0.8225559148288356, 4.56098810210824)
+ >>> anime_dbaesthetic('good.jpg', fmt=('label', 'percentile', 'score'))
+ ('good', 0.591020403706702, 3.670568235218525)
+ >>> anime_dbaesthetic('normal.jpg', fmt=('label', 'percentile', 'score'))
+ ('normal', 0.2888798940585766, 2.1677918508648872)
+ >>> anime_dbaesthetic('low.jpg', fmt=('label', 'percentile', 'score'))
+ ('low', 0.243279223969715, 1.9305131509900093)
+ >>> anime_dbaesthetic('worst.jpg', fmt=('label', 'percentile', 'score'))
+ ('worst', 0.005268185993767627, 0.6085879728198051)
+
+ * Get confidence
+
+ >>> anime_dbaesthetic('masterpiece.jpg', fmt='confidence')
+ {'masterpiece': 0.6834832429885864, 'best': 0.16141420602798462, 'great': 0.05435194447636604, 'good': 0.025083942338824272, 'normal': 0.024000568315386772, 'low': 0.027076328173279762, 'worst': 0.024589713662862778}
+ >>> anime_dbaesthetic('best.jpg', fmt='confidence')
+ {'masterpiece': 0.3757021427154541, 'best': 0.3451208472251892, 'great': 0.1511985808610916, 'good': 0.04740551486611366, 'normal': 0.02172713913023472, 'low': 0.027498546987771988, 'worst': 0.03134724497795105}
+ >>> anime_dbaesthetic('great.jpg', fmt='confidence')
+ {'masterpiece': 0.39281174540519714, 'best': 0.22457796335220337, 'great': 0.15563568472862244, 'good': 0.10796019434928894, 'normal': 0.047730278223752975, 'low': 0.0393439345061779, 'worst': 0.031940147280693054}
+ >>> anime_dbaesthetic('good.jpg', fmt='confidence')
+ {'masterpiece': 0.13832266628742218, 'best': 0.20687267184257507, 'great': 0.2509062886238098, 'good': 0.1644320785999298, 'normal': 0.11332042515277863, 'low': 0.08270663768053055, 'worst': 0.043439216911792755}
+ >>> anime_dbaesthetic('normal.jpg', fmt='confidence')
+ {'masterpiece': 0.033693961799144745, 'best': 0.03375888615846634, 'great': 0.050045162439346313, 'good': 0.16734018921852112, 'normal': 0.4311050772666931, 'low': 0.23242227733135223, 'worst': 0.05163438618183136}
+ >>> anime_dbaesthetic('low.jpg', fmt='confidence')
+ {'masterpiece': 0.012833272106945515, 'best': 0.01619996316730976, 'great': 0.03074900433421135, 'good': 0.1396280825138092, 'normal': 0.5038207173347473, 'low': 0.22299200296401978, 'worst': 0.07377689331769943}
+ >>> anime_dbaesthetic('worst.jpg', fmt='confidence')
+ {'masterpiece': 0.02854202501475811, 'best': 0.026677291840314865, 'great': 0.02838410809636116, 'good': 0.026617199182510376, 'normal': 0.02508518099784851, 'low': 0.06039097160100937, 'worst': 0.8043031692504883}
+ """
+ return _MODEL.get_aesthetic(image, model_name, fmt)
diff --git a/test/metrics/test_dbaesthetic.py b/test/metrics/test_dbaesthetic.py
new file mode 100644
index 00000000000..0d3388a6b5d
--- /dev/null
+++ b/test/metrics/test_dbaesthetic.py
@@ -0,0 +1,44 @@
+import pytest
+
+from imgutils.metrics import anime_dbaesthetic
+from imgutils.metrics.dbaesthetic import _MODEL, _LABELS
+from test.testings import get_testfile
+
+
+@pytest.fixture(scope='module', autouse=True)
+def _release_model():
+ try:
+ yield
+ finally:
+ _MODEL.clear()
+
+
+@pytest.mark.unittest
+class TestMetricsDBAesthetic:
+ @pytest.mark.parametrize(['label', 'image'], [(v, f'{v}.jpg') for v in _LABELS])
+ def test_anime_dbaesthetic(self, label, image):
+ image_file = get_testfile('dbaesthetic', image)
+ assert anime_dbaesthetic(image_file, fmt='label') == label
+
+ @pytest.mark.parametrize(['label', 'image'], [(v, f'{v}.jpg') for v in _LABELS])
+ def test_anime_dbaesthetic_default(self, label, image):
+ image_file = get_testfile('dbaesthetic', image)
+ label_, percentile = anime_dbaesthetic(image_file)
+ assert label_ == label
+
+ @pytest.mark.parametrize(['label', 'image'], [(v, f'{v}.jpg') for v in _LABELS])
+ def test_anime_dbaesthetic_list(self, label, image):
+ image_file = get_testfile('dbaesthetic', image)
+ r = anime_dbaesthetic(image_file, fmt=['label', 'percentile'])
+ assert isinstance(r, list)
+ assert len(r) == 2
+ label_, percentile = r
+ assert label_ == label
+
+ @pytest.mark.parametrize(['label', 'image'], [(v, f'{v}.jpg') for v in _LABELS])
+ def test_anime_dbaesthetic_dict(self, label, image):
+ image_file = get_testfile('dbaesthetic', image)
+ r = anime_dbaesthetic(image_file, fmt={'label': 'label', 'score': 'percentile', 'conf': 'confidence'})
+ assert r['label'] == label
+ assert isinstance(r['conf'], dict)
+ assert len(r['conf']) == 7
diff --git a/test/testfile/dbaesthetic/best.jpg b/test/testfile/dbaesthetic/best.jpg
new file mode 100644
index 00000000000..c88c9aa705b
Binary files /dev/null and b/test/testfile/dbaesthetic/best.jpg differ
diff --git a/test/testfile/dbaesthetic/good.jpg b/test/testfile/dbaesthetic/good.jpg
new file mode 100644
index 00000000000..124b8a0f757
Binary files /dev/null and b/test/testfile/dbaesthetic/good.jpg differ
diff --git a/test/testfile/dbaesthetic/great.jpg b/test/testfile/dbaesthetic/great.jpg
new file mode 100644
index 00000000000..487b1f8da09
Binary files /dev/null and b/test/testfile/dbaesthetic/great.jpg differ
diff --git a/test/testfile/dbaesthetic/low.jpg b/test/testfile/dbaesthetic/low.jpg
new file mode 100644
index 00000000000..5dd4104460f
Binary files /dev/null and b/test/testfile/dbaesthetic/low.jpg differ
diff --git a/test/testfile/dbaesthetic/masterpiece.jpg b/test/testfile/dbaesthetic/masterpiece.jpg
new file mode 100644
index 00000000000..a3108fad030
Binary files /dev/null and b/test/testfile/dbaesthetic/masterpiece.jpg differ
diff --git a/test/testfile/dbaesthetic/normal.jpg b/test/testfile/dbaesthetic/normal.jpg
new file mode 100644
index 00000000000..2dc844350f4
Binary files /dev/null and b/test/testfile/dbaesthetic/normal.jpg differ
diff --git a/test/testfile/dbaesthetic/worst.jpg b/test/testfile/dbaesthetic/worst.jpg
new file mode 100644
index 00000000000..e75a6666436
Binary files /dev/null and b/test/testfile/dbaesthetic/worst.jpg differ