diff --git a/docs/source/api_doc/validate/dbrating.plot.py b/docs/source/api_doc/validate/dbrating.plot.py new file mode 100644 index 00000000000..8cb5838f7f8 --- /dev/null +++ b/docs/source/api_doc/validate/dbrating.plot.py @@ -0,0 +1,15 @@ +import glob +import os.path + +from natsort import natsorted + +from plot import image_plot + +if __name__ == '__main__': + image_plot( + *natsorted(glob.glob(os.path.join('dbrating', 'general', '*.jpg'))), + *natsorted(glob.glob(os.path.join('dbrating', 'sensitive', '*.jpg'))), + *natsorted(glob.glob(os.path.join('dbrating', 'questionable', '*.jpg'))), + *natsorted(glob.glob(os.path.join('dbrating', 'explicit', '*.jpg'))), + columns=4, figsize=(10, 15), + ) diff --git a/docs/source/api_doc/validate/dbrating/explicit/danbooru_1222061.jpg b/docs/source/api_doc/validate/dbrating/explicit/danbooru_1222061.jpg new file mode 100644 index 00000000000..8fe903172bf Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/explicit/danbooru_1222061.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/explicit/danbooru_4540707.jpg b/docs/source/api_doc/validate/dbrating/explicit/danbooru_4540707.jpg new file mode 100644 index 00000000000..85c6fe58c9f Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/explicit/danbooru_4540707.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/explicit/danbooru_5574035.jpg b/docs/source/api_doc/validate/dbrating/explicit/danbooru_5574035.jpg new file mode 100644 index 00000000000..0cd7fbf1f6d Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/explicit/danbooru_5574035.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/explicit/danbooru_5638872.jpg b/docs/source/api_doc/validate/dbrating/explicit/danbooru_5638872.jpg new file mode 100644 index 00000000000..7b92a1a9f27 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/explicit/danbooru_5638872.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/general/danbooru_1770285.jpg b/docs/source/api_doc/validate/dbrating/general/danbooru_1770285.jpg new file mode 100644 index 00000000000..e6aa4879db7 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/general/danbooru_1770285.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/general/danbooru_3992500.jpg b/docs/source/api_doc/validate/dbrating/general/danbooru_3992500.jpg new file mode 100644 index 00000000000..b6b0b4005df Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/general/danbooru_3992500.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/general/danbooru_4419371.jpg b/docs/source/api_doc/validate/dbrating/general/danbooru_4419371.jpg new file mode 100644 index 00000000000..2655bf9d150 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/general/danbooru_4419371.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/general/danbooru_5414874.jpg b/docs/source/api_doc/validate/dbrating/general/danbooru_5414874.jpg new file mode 100644 index 00000000000..4b507540ad9 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/general/danbooru_5414874.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/general/danbooru_6627908.jpg b/docs/source/api_doc/validate/dbrating/general/danbooru_6627908.jpg new file mode 100644 index 00000000000..541e356b392 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/general/danbooru_6627908.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/general/danbooru_6788719.jpg b/docs/source/api_doc/validate/dbrating/general/danbooru_6788719.jpg new file mode 100644 index 00000000000..110a61dd9a7 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/general/danbooru_6788719.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/questionable/danbooru_1055724.jpg b/docs/source/api_doc/validate/dbrating/questionable/danbooru_1055724.jpg new file mode 100644 index 00000000000..d83c679ec61 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/questionable/danbooru_1055724.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/questionable/danbooru_2400240.jpg b/docs/source/api_doc/validate/dbrating/questionable/danbooru_2400240.jpg new file mode 100644 index 00000000000..d268ae4c8d9 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/questionable/danbooru_2400240.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/questionable/danbooru_3268458.jpg b/docs/source/api_doc/validate/dbrating/questionable/danbooru_3268458.jpg new file mode 100644 index 00000000000..3735adf15d3 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/questionable/danbooru_3268458.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/questionable/danbooru_5542691.jpg b/docs/source/api_doc/validate/dbrating/questionable/danbooru_5542691.jpg new file mode 100644 index 00000000000..fd824c9ed46 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/questionable/danbooru_5542691.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/questionable/danbooru_6844379.jpg b/docs/source/api_doc/validate/dbrating/questionable/danbooru_6844379.jpg new file mode 100644 index 00000000000..2fb21e5d0f1 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/questionable/danbooru_6844379.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/questionable/danbooru_873806.jpg b/docs/source/api_doc/validate/dbrating/questionable/danbooru_873806.jpg new file mode 100644 index 00000000000..5ccfc7a4db4 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/questionable/danbooru_873806.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/sensitive/danbooru_1559062.jpg b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_1559062.jpg new file mode 100644 index 00000000000..a0d60cb06bf Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_1559062.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/sensitive/danbooru_3267697.jpg b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_3267697.jpg new file mode 100644 index 00000000000..1d4bff1b5ae Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_3267697.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/sensitive/danbooru_4116416.jpg b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_4116416.jpg new file mode 100644 index 00000000000..45354146964 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_4116416.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/sensitive/danbooru_4654826.jpg b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_4654826.jpg new file mode 100644 index 00000000000..d74db0b0556 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_4654826.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/sensitive/danbooru_5557434.jpg b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_5557434.jpg new file mode 100644 index 00000000000..34139344163 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_5557434.jpg differ diff --git a/docs/source/api_doc/validate/dbrating/sensitive/danbooru_6249139.jpg b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_6249139.jpg new file mode 100644 index 00000000000..f04c7e67610 Binary files /dev/null and b/docs/source/api_doc/validate/dbrating/sensitive/danbooru_6249139.jpg differ diff --git a/docs/source/api_doc/validate/dbrating_benchmark.plot.py b/docs/source/api_doc/validate/dbrating_benchmark.plot.py new file mode 100644 index 00000000000..7574ebdd1a9 --- /dev/null +++ b/docs/source/api_doc/validate/dbrating_benchmark.plot.py @@ -0,0 +1,36 @@ +import random + +from benchmark import BaseBenchmark, create_plot_cli +from imgutils.generic.classify import _open_models_for_repo_id +from imgutils.validate import anime_dbrating +from imgutils.validate.dbrating import _REPO_ID + +_MODEL_NAMES = _open_models_for_repo_id(_REPO_ID).model_names + + +class AnimeDBRatingCharacterBenchmark(BaseBenchmark): + def __init__(self, model): + BaseBenchmark.__init__(self) + self.model = model + + def load(self): + _open_models_for_repo_id(_REPO_ID)._open_model(self.model) + + def unload(self): + _open_models_for_repo_id(_REPO_ID).clear() + + def run(self): + image_file = random.choice(self.all_images) + _ = anime_dbrating(image_file, self.model) + + +if __name__ == '__main__': + create_plot_cli( + [ + (name, AnimeDBRatingCharacterBenchmark(name)) + for name in _MODEL_NAMES + ], + title='Benchmark for Danbooru Rating Models', + run_times=10, + try_times=20, + )() diff --git a/imgutils/validate/__init__.py b/imgutils/validate/__init__.py index 0c5cda4824b..83879e7ae66 100644 --- a/imgutils/validate/__init__.py +++ b/imgutils/validate/__init__.py @@ -7,6 +7,7 @@ from .classify import * from .color import * from .completeness import * +from .dbrating import * from .monochrome import * from .nsfw import * from .portrait import * diff --git a/imgutils/validate/dbrating.py b/imgutils/validate/dbrating.py new file mode 100644 index 00000000000..33b5b9d655c --- /dev/null +++ b/imgutils/validate/dbrating.py @@ -0,0 +1,76 @@ +""" +Overview: + A model for rating anime images into 4 classes (``general``, ``sensitive``, ``questionable`` and ``explicit``), + based on danbooru rating system. + + The following are sample images for testing. + + .. collapse:: The following are sample images for testing. (WARNING: NSFW!!!) + + .. image:: dbrating.plot.py.svg + :align: center + + This is an overall benchmark of all the rating validation models: + + .. image:: dbrating_benchmark.plot.py.svg + :align: center + + The models are hosted on + `huggingface - deepghs/anime_dbrating `_. + + .. note:: + This model is based on danbooru rating system, trained with 1.2 million images. + If you need 3-level rating prediction, use :func:`imgutils.validate.rating.anime_rating`. + + .. note:: + Please note that the classification of ``general``, ``sensitive``, ``questionable`` and ``explicit`` types + does not have clear boundaries, making it challenging to clean the training data. As a result, + there is no strict ground truth for the rating classification problem. The judgment functionality + provided by the current module is intended as a quick and rough estimation. + + **If you require an accurate filtering or judgment function specifically for R-18 images, + it is recommended to consider using object detection-based methods**, + such as using :func:`imgutils.detect.censor.detect_censors` to detect sensitive regions as the basis for judgment. +""" +from typing import Tuple, Dict + +from ..data import ImageTyping +from ..generic import classify_predict, classify_predict_score + +__all__ = [ + 'anime_dbrating_score', + 'anime_dbrating', +] + +_DEFAULT_MODEL_NAME = 'mobilenetv3_large_100_v0_ls0.2' +_REPO_ID = 'deepghs/anime_dbrating' + + +def anime_dbrating_score(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME) -> Dict[str, float]: + """ + Overview: + Predict the rating of the given image, return the score with as a dict object. + + :param image: Image to rating. + :param model_name: Model to use. Default is ``mobilenetv3_large_100_v0_ls0.2``. + All available models are listed on the benchmark plot above. + If you need better accuracy, just set this to ``caformer_s36_v0_ls0.2``. + :return: A dict with ratings and scores. + + """ + return classify_predict_score(image, _REPO_ID, model_name) + + +def anime_dbrating(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME) -> Tuple[str, float]: + """ + Overview: + Predict the rating of the given image, return the class and its score. + + :param image: Image to rating. + :param model_name: Model to use. Default is ``mobilenetv3_large_100_v0_ls0.2``. All available models are listed + on the benchmark plot above. If you need better accuracy, just set this to ``caformer_s36_v0_ls0.2``. + :return: A tuple contains the rating and its score. + + + """ + return classify_predict(image, _REPO_ID, model_name) diff --git a/imgutils/validate/rating.py b/imgutils/validate/rating.py index 675f996b2cf..06f7fbd6d7f 100644 --- a/imgutils/validate/rating.py +++ b/imgutils/validate/rating.py @@ -1,6 +1,6 @@ """ Overview: - A model for rating anime images into 4 classes (``safe``, ``r15`` and ``r18``). + A model for rating anime images into 3 classes (``safe``, ``r15`` and ``r18``), based on sankaku rating system. The following are sample images for testing. diff --git a/test/testfile/anime_dbrating/explicit/danbooru_1222061.jpg b/test/testfile/anime_dbrating/explicit/danbooru_1222061.jpg new file mode 100644 index 00000000000..8fe903172bf Binary files /dev/null and b/test/testfile/anime_dbrating/explicit/danbooru_1222061.jpg differ diff --git a/test/testfile/anime_dbrating/explicit/danbooru_1728047.jpg b/test/testfile/anime_dbrating/explicit/danbooru_1728047.jpg new file mode 100644 index 00000000000..733e29985c6 Binary files /dev/null and b/test/testfile/anime_dbrating/explicit/danbooru_1728047.jpg differ diff --git a/test/testfile/anime_dbrating/explicit/danbooru_4540707.jpg b/test/testfile/anime_dbrating/explicit/danbooru_4540707.jpg new file mode 100644 index 00000000000..85c6fe58c9f Binary files /dev/null and b/test/testfile/anime_dbrating/explicit/danbooru_4540707.jpg differ diff --git a/test/testfile/anime_dbrating/explicit/danbooru_4835284.jpg b/test/testfile/anime_dbrating/explicit/danbooru_4835284.jpg new file mode 100644 index 00000000000..01ca6852513 Binary files /dev/null and b/test/testfile/anime_dbrating/explicit/danbooru_4835284.jpg differ diff --git a/test/testfile/anime_dbrating/explicit/danbooru_5574035.jpg b/test/testfile/anime_dbrating/explicit/danbooru_5574035.jpg new file mode 100644 index 00000000000..0cd7fbf1f6d Binary files /dev/null and b/test/testfile/anime_dbrating/explicit/danbooru_5574035.jpg differ diff --git a/test/testfile/anime_dbrating/explicit/danbooru_5638872.jpg b/test/testfile/anime_dbrating/explicit/danbooru_5638872.jpg new file mode 100644 index 00000000000..7b92a1a9f27 Binary files /dev/null and b/test/testfile/anime_dbrating/explicit/danbooru_5638872.jpg differ diff --git a/test/testfile/anime_dbrating/general/danbooru_1770285.jpg b/test/testfile/anime_dbrating/general/danbooru_1770285.jpg new file mode 100644 index 00000000000..e6aa4879db7 Binary files /dev/null and b/test/testfile/anime_dbrating/general/danbooru_1770285.jpg differ diff --git a/test/testfile/anime_dbrating/general/danbooru_3992500.jpg b/test/testfile/anime_dbrating/general/danbooru_3992500.jpg new file mode 100644 index 00000000000..b6b0b4005df Binary files /dev/null and b/test/testfile/anime_dbrating/general/danbooru_3992500.jpg differ diff --git a/test/testfile/anime_dbrating/general/danbooru_4419371.jpg b/test/testfile/anime_dbrating/general/danbooru_4419371.jpg new file mode 100644 index 00000000000..2655bf9d150 Binary files /dev/null and b/test/testfile/anime_dbrating/general/danbooru_4419371.jpg differ diff --git a/test/testfile/anime_dbrating/general/danbooru_5414874.jpg b/test/testfile/anime_dbrating/general/danbooru_5414874.jpg new file mode 100644 index 00000000000..4b507540ad9 Binary files /dev/null and b/test/testfile/anime_dbrating/general/danbooru_5414874.jpg differ diff --git a/test/testfile/anime_dbrating/general/danbooru_6627908.jpg b/test/testfile/anime_dbrating/general/danbooru_6627908.jpg new file mode 100644 index 00000000000..541e356b392 Binary files /dev/null and b/test/testfile/anime_dbrating/general/danbooru_6627908.jpg differ diff --git a/test/testfile/anime_dbrating/general/danbooru_6788719.jpg b/test/testfile/anime_dbrating/general/danbooru_6788719.jpg new file mode 100644 index 00000000000..110a61dd9a7 Binary files /dev/null and b/test/testfile/anime_dbrating/general/danbooru_6788719.jpg differ diff --git a/test/testfile/anime_dbrating/questionable/danbooru_1055724.jpg b/test/testfile/anime_dbrating/questionable/danbooru_1055724.jpg new file mode 100644 index 00000000000..d83c679ec61 Binary files /dev/null and b/test/testfile/anime_dbrating/questionable/danbooru_1055724.jpg differ diff --git a/test/testfile/anime_dbrating/questionable/danbooru_2400240.jpg b/test/testfile/anime_dbrating/questionable/danbooru_2400240.jpg new file mode 100644 index 00000000000..d268ae4c8d9 Binary files /dev/null and b/test/testfile/anime_dbrating/questionable/danbooru_2400240.jpg differ diff --git a/test/testfile/anime_dbrating/questionable/danbooru_3268458.jpg b/test/testfile/anime_dbrating/questionable/danbooru_3268458.jpg new file mode 100644 index 00000000000..3735adf15d3 Binary files /dev/null and b/test/testfile/anime_dbrating/questionable/danbooru_3268458.jpg differ diff --git a/test/testfile/anime_dbrating/questionable/danbooru_5542691.jpg b/test/testfile/anime_dbrating/questionable/danbooru_5542691.jpg new file mode 100644 index 00000000000..fd824c9ed46 Binary files /dev/null and b/test/testfile/anime_dbrating/questionable/danbooru_5542691.jpg differ diff --git a/test/testfile/anime_dbrating/questionable/danbooru_6844379.jpg b/test/testfile/anime_dbrating/questionable/danbooru_6844379.jpg new file mode 100644 index 00000000000..2fb21e5d0f1 Binary files /dev/null and b/test/testfile/anime_dbrating/questionable/danbooru_6844379.jpg differ diff --git a/test/testfile/anime_dbrating/questionable/danbooru_873806.jpg b/test/testfile/anime_dbrating/questionable/danbooru_873806.jpg new file mode 100644 index 00000000000..5ccfc7a4db4 Binary files /dev/null and b/test/testfile/anime_dbrating/questionable/danbooru_873806.jpg differ diff --git a/test/testfile/anime_dbrating/sensitive/danbooru_1559062.jpg b/test/testfile/anime_dbrating/sensitive/danbooru_1559062.jpg new file mode 100644 index 00000000000..a0d60cb06bf Binary files /dev/null and b/test/testfile/anime_dbrating/sensitive/danbooru_1559062.jpg differ diff --git a/test/testfile/anime_dbrating/sensitive/danbooru_3267697.jpg b/test/testfile/anime_dbrating/sensitive/danbooru_3267697.jpg new file mode 100644 index 00000000000..1d4bff1b5ae Binary files /dev/null and b/test/testfile/anime_dbrating/sensitive/danbooru_3267697.jpg differ diff --git a/test/testfile/anime_dbrating/sensitive/danbooru_4116416.jpg b/test/testfile/anime_dbrating/sensitive/danbooru_4116416.jpg new file mode 100644 index 00000000000..45354146964 Binary files /dev/null and b/test/testfile/anime_dbrating/sensitive/danbooru_4116416.jpg differ diff --git a/test/testfile/anime_dbrating/sensitive/danbooru_4654826.jpg b/test/testfile/anime_dbrating/sensitive/danbooru_4654826.jpg new file mode 100644 index 00000000000..d74db0b0556 Binary files /dev/null and b/test/testfile/anime_dbrating/sensitive/danbooru_4654826.jpg differ diff --git a/test/testfile/anime_dbrating/sensitive/danbooru_5557434.jpg b/test/testfile/anime_dbrating/sensitive/danbooru_5557434.jpg new file mode 100644 index 00000000000..34139344163 Binary files /dev/null and b/test/testfile/anime_dbrating/sensitive/danbooru_5557434.jpg differ diff --git a/test/testfile/anime_dbrating/sensitive/danbooru_6249139.jpg b/test/testfile/anime_dbrating/sensitive/danbooru_6249139.jpg new file mode 100644 index 00000000000..f04c7e67610 Binary files /dev/null and b/test/testfile/anime_dbrating/sensitive/danbooru_6249139.jpg differ diff --git a/test/validate/test_dbrating.py b/test/validate/test_dbrating.py new file mode 100644 index 00000000000..f24b3b670c7 --- /dev/null +++ b/test/validate/test_dbrating.py @@ -0,0 +1,38 @@ +import glob +import os.path + +import pytest + +from imgutils.generic.classify import _open_models_for_repo_id +from imgutils.validate import anime_dbrating +from imgutils.validate.dbrating import anime_dbrating_score, _REPO_ID +from test.testings import get_testfile + +_ROOT_DIR = get_testfile('anime_dbrating') +_EXAMPLE_FILES = [ + (os.path.relpath(file, _ROOT_DIR), os.path.basename(os.path.dirname(file))) + for file in glob.glob(get_testfile(_ROOT_DIR, '**', '*.jpg'), recursive=True) +] + + +@pytest.fixture(scope='module', autouse=True) +def _release_model_after_run(): + try: + yield + finally: + _open_models_for_repo_id(_REPO_ID).clear() + + +@pytest.mark.unittest +class TestValidatedbrating: + @pytest.mark.parametrize(['image', 'label'], _EXAMPLE_FILES) + def test_anime_dbrating(self, image, label): + image_file = get_testfile('anime_dbrating', image) + tag, score = anime_dbrating(image_file) + assert tag == label + + @pytest.mark.parametrize(['image', 'label'], _EXAMPLE_FILES) + def test_anime_dbrating_score(self, image, label): + image_file = get_testfile('anime_dbrating', image) + scores = anime_dbrating_score(image_file) + assert scores[label] > 0.5