diff --git a/docs/source/api_doc/validate/index.rst b/docs/source/api_doc/validate/index.rst index 14be53919cb..809b0fe6e01 100644 --- a/docs/source/api_doc/validate/index.rst +++ b/docs/source/api_doc/validate/index.rst @@ -16,5 +16,6 @@ imgutils.validate nsfw portrait rating + style_age teen truncate diff --git a/docs/source/api_doc/validate/style_age.plot.py b/docs/source/api_doc/validate/style_age.plot.py new file mode 100644 index 00000000000..16ca9373336 --- /dev/null +++ b/docs/source/api_doc/validate/style_age.plot.py @@ -0,0 +1,18 @@ +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('style_age', '1970s-', '*.jpg'))), + *natsorted(glob.glob(os.path.join('style_age', '1980s', '*.jpg'))), + *natsorted(glob.glob(os.path.join('style_age', '1990s', '*.jpg'))), + *natsorted(glob.glob(os.path.join('style_age', '2000s', '*.jpg'))), + *natsorted(glob.glob(os.path.join('style_age', '2010s', '*.jpg'))), + *natsorted(glob.glob(os.path.join('style_age', '2015s', '*.jpg'))), + *natsorted(glob.glob(os.path.join('style_age', '2020s', '*.jpg'))), + columns=4, figsize=(10, 26), + ) diff --git a/docs/source/api_doc/validate/style_age.rst b/docs/source/api_doc/validate/style_age.rst new file mode 100644 index 00000000000..f990f93d675 --- /dev/null +++ b/docs/source/api_doc/validate/style_age.rst @@ -0,0 +1,21 @@ +imgutils.validate.style_age +============================================= + +.. currentmodule:: imgutils.validate.style_age + +.. automodule:: imgutils.validate.style_age + + +anime_style_age_score +----------------------------- + +.. autofunction:: anime_style_age_score + + + +anime_style_age +----------------------------- + +.. autofunction:: anime_style_age + + diff --git a/docs/source/api_doc/validate/style_age/1970s-/1.jpg b/docs/source/api_doc/validate/style_age/1970s-/1.jpg new file mode 100644 index 00000000000..ae7fe0e9940 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1970s-/1.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1970s-/2.jpg b/docs/source/api_doc/validate/style_age/1970s-/2.jpg new file mode 100644 index 00000000000..95d4f92661d Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1970s-/2.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1970s-/3.jpg b/docs/source/api_doc/validate/style_age/1970s-/3.jpg new file mode 100644 index 00000000000..5579231e69d Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1970s-/3.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1970s-/4.jpg b/docs/source/api_doc/validate/style_age/1970s-/4.jpg new file mode 100644 index 00000000000..ffc999ee997 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1970s-/4.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1980s/5.jpg b/docs/source/api_doc/validate/style_age/1980s/5.jpg new file mode 100644 index 00000000000..6f21766ed6f Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1980s/5.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1980s/6.jpg b/docs/source/api_doc/validate/style_age/1980s/6.jpg new file mode 100644 index 00000000000..8036453e0ee Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1980s/6.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1980s/7.jpg b/docs/source/api_doc/validate/style_age/1980s/7.jpg new file mode 100644 index 00000000000..0d6e0d8e5a2 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1980s/7.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1980s/8.jpg b/docs/source/api_doc/validate/style_age/1980s/8.jpg new file mode 100644 index 00000000000..cc6fb2aafa7 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1980s/8.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1990s/10.jpg b/docs/source/api_doc/validate/style_age/1990s/10.jpg new file mode 100644 index 00000000000..44f767da566 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1990s/10.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1990s/11.jpg b/docs/source/api_doc/validate/style_age/1990s/11.jpg new file mode 100644 index 00000000000..49fa9e41406 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1990s/11.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1990s/12.jpg b/docs/source/api_doc/validate/style_age/1990s/12.jpg new file mode 100644 index 00000000000..a410f4822ad Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1990s/12.jpg differ diff --git a/docs/source/api_doc/validate/style_age/1990s/9.jpg b/docs/source/api_doc/validate/style_age/1990s/9.jpg new file mode 100644 index 00000000000..5d455e4ded9 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/1990s/9.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2000s/13.jpg b/docs/source/api_doc/validate/style_age/2000s/13.jpg new file mode 100644 index 00000000000..4f277acf789 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2000s/13.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2000s/14.jpg b/docs/source/api_doc/validate/style_age/2000s/14.jpg new file mode 100644 index 00000000000..23b8f97b373 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2000s/14.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2000s/15.jpg b/docs/source/api_doc/validate/style_age/2000s/15.jpg new file mode 100644 index 00000000000..467d440d482 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2000s/15.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2000s/16.jpg b/docs/source/api_doc/validate/style_age/2000s/16.jpg new file mode 100644 index 00000000000..f625b82a1d5 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2000s/16.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2010s/17.jpg b/docs/source/api_doc/validate/style_age/2010s/17.jpg new file mode 100644 index 00000000000..0c980089563 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2010s/17.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2010s/18.jpg b/docs/source/api_doc/validate/style_age/2010s/18.jpg new file mode 100644 index 00000000000..0e3fac308c3 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2010s/18.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2010s/19.jpg b/docs/source/api_doc/validate/style_age/2010s/19.jpg new file mode 100644 index 00000000000..c0b6b83c501 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2010s/19.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2010s/20.jpg b/docs/source/api_doc/validate/style_age/2010s/20.jpg new file mode 100644 index 00000000000..bafe3c060d2 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2010s/20.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2015s/21.jpg b/docs/source/api_doc/validate/style_age/2015s/21.jpg new file mode 100644 index 00000000000..24740aadf3d Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2015s/21.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2015s/22.jpg b/docs/source/api_doc/validate/style_age/2015s/22.jpg new file mode 100644 index 00000000000..4950eff7f2a Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2015s/22.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2015s/23.jpg b/docs/source/api_doc/validate/style_age/2015s/23.jpg new file mode 100644 index 00000000000..7fa2513c6d1 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2015s/23.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2015s/24.jpg b/docs/source/api_doc/validate/style_age/2015s/24.jpg new file mode 100644 index 00000000000..b6b7389d10f Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2015s/24.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2020s/25.jpg b/docs/source/api_doc/validate/style_age/2020s/25.jpg new file mode 100644 index 00000000000..84a136b779c Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2020s/25.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2020s/26.jpg b/docs/source/api_doc/validate/style_age/2020s/26.jpg new file mode 100644 index 00000000000..220e91c205f Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2020s/26.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2020s/27.jpg b/docs/source/api_doc/validate/style_age/2020s/27.jpg new file mode 100644 index 00000000000..bec3eb55cd3 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2020s/27.jpg differ diff --git a/docs/source/api_doc/validate/style_age/2020s/28.jpg b/docs/source/api_doc/validate/style_age/2020s/28.jpg new file mode 100644 index 00000000000..ac1ec0ea181 Binary files /dev/null and b/docs/source/api_doc/validate/style_age/2020s/28.jpg differ diff --git a/docs/source/api_doc/validate/style_age_benchmark.plot.py b/docs/source/api_doc/validate/style_age_benchmark.plot.py new file mode 100644 index 00000000000..49f4c3de02d --- /dev/null +++ b/docs/source/api_doc/validate/style_age_benchmark.plot.py @@ -0,0 +1,45 @@ +import os +import random + +from huggingface_hub import HfFileSystem + +from benchmark import BaseBenchmark, create_plot_cli +from imgutils.validate import anime_style_age + +hf_fs = HfFileSystem() + +_REPOSITORY = 'deepghs/anime_style_ages' +_MODEL_NAMES = [ + os.path.relpath(file, _REPOSITORY).split('/')[0] for file in + hf_fs.glob(f'{_REPOSITORY}/*/model.onnx') +] + + +class AnimeStyleAgeBenchmark(BaseBenchmark): + def __init__(self, model): + BaseBenchmark.__init__(self) + self.model = model + + def load(self): + from imgutils.validate.style_age import _open_anime_style_age_model + _ = _open_anime_style_age_model(self.model) + + def unload(self): + from imgutils.validate.style_age import _open_anime_style_age_model + _open_anime_style_age_model.cache_clear() + + def run(self): + image_file = random.choice(self.all_images) + _ = anime_style_age(image_file, self.model) + + +if __name__ == '__main__': + create_plot_cli( + [ + (name, AnimeStyleAgeBenchmark(name)) + for name in _MODEL_NAMES + ], + title='Benchmark for Anime Portrait Models', + run_times=10, + try_times=20, + )() diff --git a/imgutils/validate/__init__.py b/imgutils/validate/__init__.py index e02659787de..9dc5da249e5 100644 --- a/imgutils/validate/__init__.py +++ b/imgutils/validate/__init__.py @@ -9,5 +9,6 @@ from .nsfw import * from .portrait import * from .rating import * +from .style_age import * from .teen import * from .truncate import * diff --git a/imgutils/validate/style_age.py b/imgutils/validate/style_age.py new file mode 100644 index 00000000000..81a861b5ea7 --- /dev/null +++ b/imgutils/validate/style_age.py @@ -0,0 +1,178 @@ +""" +Overview: + A model for classifying anime style_age images into 7 classes + (``1970s-``, ``1980s``, ``1990s``, ``2000s``, ``2010s``, ``2015s``, ``2020s``). + + The following are sample images for testing. + + .. image:: style_age.plot.py.svg + :align: center + + This is an overall benchmark of all the style_age classification models: + + .. image:: style_age_benchmark.plot.py.svg + :align: center + + The models are hosted on + `huggingface - deepghs/anime_style_ages `_. +""" +import json +from functools import lru_cache +from typing import Tuple, Optional, Dict, List + +import numpy as np +from PIL import Image +from huggingface_hub import hf_hub_download + +from imgutils.data import rgb_encode, ImageTyping, load_image +from imgutils.utils import open_onnx_model + +__all__ = [ + 'anime_style_age_score', + 'anime_style_age', +] + +_DEFAULT_MODEL_NAME = 'mobilenetv3_v0_dist' + + +@lru_cache() +def _open_anime_style_age_model(model_name): + """ + Open the anime style age model. + + :param model_name: The model name. + :type model_name: str + :return: The ONNX model. + """ + return open_onnx_model(hf_hub_download( + f'deepghs/anime_style_ages', + f'{model_name}/model.onnx', + )) + + +@lru_cache() +def _get_anime_style_age_labels(model_name) -> List[str]: + """ + Get the labels for the anime style age model. + + :param model_name: The model name. + :type model_name: str + :return: The list of labels. + :rtype: List[str] + """ + with open(hf_hub_download( + f'deepghs/anime_style_ages', + f'{model_name}/meta.json', + ), 'r') as f: + return json.load(f)['labels'] + + +def _img_encode(image: Image.Image, size: Tuple[int, int] = (384, 384), + normalize: Optional[Tuple[float, float]] = (0.5, 0.5)): + """ + Encode the input image. + + :param image: The input image. + :type image: Image.Image + :param size: The desired size of the image. + :type size: Tuple[int, int] + :param normalize: Mean and standard deviation for normalization. Default is (0.5, 0.5). + :type normalize: Optional[Tuple[float, float]] + :return: The encoded image data. + :rtype: np.ndarray + """ + image = image.resize(size, Image.BILINEAR) + data = rgb_encode(image, order_='CHW') + + if normalize is not None: + mean_, std_ = normalize + mean = np.asarray([mean_]).reshape((-1, 1, 1)) + std = np.asarray([std_]).reshape((-1, 1, 1)) + data = (data - mean) / std + + return data.astype(np.float32) + + +def _raw_anime_style_age(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME): + """ + Perform raw anime style age processing on the input image. + + :param image: The input image. + :type image: ImageTyping + :param model_name: The model name. Default is 'mobilenetv3_v0_dist'. + :type model_name: str + :return: The processed image data. + :rtype: np.ndarray + """ + image = load_image(image, force_background='white', mode='RGB') + input_ = _img_encode(image)[None, ...] + output, = _open_anime_style_age_model(model_name).run(['output'], {'input': input_}) + return output + + +def anime_style_age_score(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME) -> Dict[str, float]: + """ + Get the scores for different types in an anime style age. + + :param image: The input image. + :type image: ImageTyping + :param model_name: The model name. Default is 'mobilenetv3_v0_dist'. + :type model_name: str + :return: A dictionary with type scores. + :rtype: Dict[str, float] + + Examples:: + >>> from imgutils.validate import anime_style_age_score + >>> + >>> anime_style_age_score('style_age/1970s-/1.jpg') + {'1970s-': 0.9805465340614319, '1980s': 8.761269782553427e-06, '1990s': 0.0005044879508204758, '2000s': 0.01569165475666523, '2010s': 0.002850610064342618, '2015s': 0.00037849770160391927, '2020s': 1.9434612113400362e-05} + >>> anime_style_age_score('style_age/1980s/5.jpg') + {'1970s-': 9.053497342392802e-05, '1980s': 0.9992554783821106, '1990s': 0.0006490182713605464, '2000s': 2.8857468805654207e-06, '2010s': 4.317252262353577e-07, '2015s': 6.314484721769986e-07, '2020s': 1.0750001138148946e-06} + >>> anime_style_age_score('style_age/1990s/9.jpg') + {'1970s-': 1.706833609205205e-05, '1980s': 0.00034479793976061046, '1990s': 0.9995512366294861, '2000s': 4.391363472677767e-05, '2010s': 1.4607510820496827e-05, '2015s': 2.0679690351244062e-05, '2020s': 7.661913514311891e-06} + >>> anime_style_age_score('style_age/2000s/13.jpg') + {'1970s-': 3.757471131393686e-05, '1980s': 3.0744897230761126e-05, '1990s': 2.76177470368566e-05, '2000s': 0.9996387958526611, '2010s': 9.160279296338558e-05, '2015s': 0.00013228354509919882, '2020s': 4.1361367038916796e-05} + >>> anime_style_age_score('style_age/2010s/17.jpg') + {'1970s-': 7.464057489414699e-06, '1980s': 3.2412899599876255e-05, '1990s': 5.703883653040975e-05, '2000s': 9.127358498517424e-05, '2010s': 0.9973921775817871, '2015s': 0.0022309015039354563, '2020s': 0.00018872201326303184} + >>> anime_style_age_score('style_age/2015s/21.jpg') + {'1970s-': 3.780902943617548e-06, '1980s': 1.422096920578042e-05, '1990s': 1.638929097680375e-05, '2000s': 2.152203023797483e-06, '2010s': 0.00028818511054851115, '2015s': 0.9996094107627869, '2020s': 6.58777353237383e-05} + >>> anime_style_age_score('style_age/2020s/25.jpg') + {'1970s-': 1.9200742826797068e-05, '1980s': 0.00017117452807724476, '1990s': 9.518441947875544e-05, '2000s': 2.885544381570071e-05, '2010s': 1.4389253010449465e-05, '2015s': 3.1696006772108376e-05, '2020s': 0.9996393918991089} + """ + output = _raw_anime_style_age(image, model_name) + values = dict(zip(_get_anime_style_age_labels(model_name), map(lambda x: x.item(), output[0]))) + return values + + +def anime_style_age(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME) -> Tuple[str, float]: + """ + Get the primary anime style age type and its score. + + :param image: The input image. + :type image: ImageTyping + :param model_name: The model name. Default is 'mobilenetv3_v0_dist'. + :type model_name: str + :return: A tuple with the primary type and its score. + :rtype: Tuple[str, float] + + Examples:: + >>> from imgutils.validate import anime_style_age + >>> + >>> anime_style_age('style_age/1970s-/1.jpg') + ('1970s-', 0.9805465340614319) + >>> anime_style_age('style_age/1980s/5.jpg') + ('1980s', 0.9992554783821106) + >>> anime_style_age('style_age/1990s/9.jpg') + ('1990s', 0.9995512366294861) + >>> anime_style_age('style_age/2000s/13.jpg') + ('2000s', 0.9996387958526611) + >>> anime_style_age('style_age/2010s/17.jpg') + ('2010s', 0.9973921775817871) + >>> anime_style_age('style_age/2015s/21.jpg') + ('2015s', 0.9996094107627869) + >>> anime_style_age('style_age/2020s/25.jpg') + ('2020s', 0.9996393918991089) + """ + output = _raw_anime_style_age(image, model_name)[0] + max_id = np.argmax(output) + return _get_anime_style_age_labels(model_name)[max_id], output[max_id].item() diff --git a/test/testfile/anime_style_age/1970s-/danbooru_2074100.jpg b/test/testfile/anime_style_age/1970s-/danbooru_2074100.jpg new file mode 100644 index 00000000000..ffc999ee997 Binary files /dev/null and b/test/testfile/anime_style_age/1970s-/danbooru_2074100.jpg differ diff --git a/test/testfile/anime_style_age/1970s-/danbooru_214923.jpg b/test/testfile/anime_style_age/1970s-/danbooru_214923.jpg new file mode 100644 index 00000000000..95d4f92661d Binary files /dev/null and b/test/testfile/anime_style_age/1970s-/danbooru_214923.jpg differ diff --git a/test/testfile/anime_style_age/1970s-/danbooru_2539537.jpg b/test/testfile/anime_style_age/1970s-/danbooru_2539537.jpg new file mode 100644 index 00000000000..ae7fe0e9940 Binary files /dev/null and b/test/testfile/anime_style_age/1970s-/danbooru_2539537.jpg differ diff --git a/test/testfile/anime_style_age/1970s-/danbooru_4439329.jpg b/test/testfile/anime_style_age/1970s-/danbooru_4439329.jpg new file mode 100644 index 00000000000..5579231e69d Binary files /dev/null and b/test/testfile/anime_style_age/1970s-/danbooru_4439329.jpg differ diff --git a/test/testfile/anime_style_age/1980s/danbooru_1385442.jpg b/test/testfile/anime_style_age/1980s/danbooru_1385442.jpg new file mode 100644 index 00000000000..cc6fb2aafa7 Binary files /dev/null and b/test/testfile/anime_style_age/1980s/danbooru_1385442.jpg differ diff --git a/test/testfile/anime_style_age/1980s/danbooru_150260.jpg b/test/testfile/anime_style_age/1980s/danbooru_150260.jpg new file mode 100644 index 00000000000..0d6e0d8e5a2 Binary files /dev/null and b/test/testfile/anime_style_age/1980s/danbooru_150260.jpg differ diff --git a/test/testfile/anime_style_age/1980s/danbooru_45736.jpg b/test/testfile/anime_style_age/1980s/danbooru_45736.jpg new file mode 100644 index 00000000000..8036453e0ee Binary files /dev/null and b/test/testfile/anime_style_age/1980s/danbooru_45736.jpg differ diff --git a/test/testfile/anime_style_age/1980s/danbooru_6419655.jpg b/test/testfile/anime_style_age/1980s/danbooru_6419655.jpg new file mode 100644 index 00000000000..6f21766ed6f Binary files /dev/null and b/test/testfile/anime_style_age/1980s/danbooru_6419655.jpg differ diff --git a/test/testfile/anime_style_age/1990s/danbooru_407370.jpg b/test/testfile/anime_style_age/1990s/danbooru_407370.jpg new file mode 100644 index 00000000000..49fa9e41406 Binary files /dev/null and b/test/testfile/anime_style_age/1990s/danbooru_407370.jpg differ diff --git a/test/testfile/anime_style_age/1990s/danbooru_4220228.jpg b/test/testfile/anime_style_age/1990s/danbooru_4220228.jpg new file mode 100644 index 00000000000..44f767da566 Binary files /dev/null and b/test/testfile/anime_style_age/1990s/danbooru_4220228.jpg differ diff --git a/test/testfile/anime_style_age/1990s/danbooru_747042.jpg b/test/testfile/anime_style_age/1990s/danbooru_747042.jpg new file mode 100644 index 00000000000..a410f4822ad Binary files /dev/null and b/test/testfile/anime_style_age/1990s/danbooru_747042.jpg differ diff --git a/test/testfile/anime_style_age/1990s/danbooru_883457.jpg b/test/testfile/anime_style_age/1990s/danbooru_883457.jpg new file mode 100644 index 00000000000..5d455e4ded9 Binary files /dev/null and b/test/testfile/anime_style_age/1990s/danbooru_883457.jpg differ diff --git a/test/testfile/anime_style_age/2000s/danbooru_580150.jpg b/test/testfile/anime_style_age/2000s/danbooru_580150.jpg new file mode 100644 index 00000000000..467d440d482 Binary files /dev/null and b/test/testfile/anime_style_age/2000s/danbooru_580150.jpg differ diff --git a/test/testfile/anime_style_age/2000s/danbooru_584135.jpg b/test/testfile/anime_style_age/2000s/danbooru_584135.jpg new file mode 100644 index 00000000000..4f277acf789 Binary files /dev/null and b/test/testfile/anime_style_age/2000s/danbooru_584135.jpg differ diff --git a/test/testfile/anime_style_age/2000s/danbooru_585417.jpg b/test/testfile/anime_style_age/2000s/danbooru_585417.jpg new file mode 100644 index 00000000000..f625b82a1d5 Binary files /dev/null and b/test/testfile/anime_style_age/2000s/danbooru_585417.jpg differ diff --git a/test/testfile/anime_style_age/2000s/danbooru_585755.jpg b/test/testfile/anime_style_age/2000s/danbooru_585755.jpg new file mode 100644 index 00000000000..23b8f97b373 Binary files /dev/null and b/test/testfile/anime_style_age/2000s/danbooru_585755.jpg differ diff --git a/test/testfile/anime_style_age/2010s/danbooru_1879381.jpg b/test/testfile/anime_style_age/2010s/danbooru_1879381.jpg new file mode 100644 index 00000000000..c0b6b83c501 Binary files /dev/null and b/test/testfile/anime_style_age/2010s/danbooru_1879381.jpg differ diff --git a/test/testfile/anime_style_age/2010s/danbooru_1884263.jpg b/test/testfile/anime_style_age/2010s/danbooru_1884263.jpg new file mode 100644 index 00000000000..0c980089563 Binary files /dev/null and b/test/testfile/anime_style_age/2010s/danbooru_1884263.jpg differ diff --git a/test/testfile/anime_style_age/2010s/danbooru_1886025.jpg b/test/testfile/anime_style_age/2010s/danbooru_1886025.jpg new file mode 100644 index 00000000000..0e3fac308c3 Binary files /dev/null and b/test/testfile/anime_style_age/2010s/danbooru_1886025.jpg differ diff --git a/test/testfile/anime_style_age/2010s/danbooru_1888985.jpg b/test/testfile/anime_style_age/2010s/danbooru_1888985.jpg new file mode 100644 index 00000000000..bafe3c060d2 Binary files /dev/null and b/test/testfile/anime_style_age/2010s/danbooru_1888985.jpg differ diff --git a/test/testfile/anime_style_age/2015s/danbooru_3724119.jpg b/test/testfile/anime_style_age/2015s/danbooru_3724119.jpg new file mode 100644 index 00000000000..b6b7389d10f Binary files /dev/null and b/test/testfile/anime_style_age/2015s/danbooru_3724119.jpg differ diff --git a/test/testfile/anime_style_age/2015s/danbooru_3726124.jpg b/test/testfile/anime_style_age/2015s/danbooru_3726124.jpg new file mode 100644 index 00000000000..24740aadf3d Binary files /dev/null and b/test/testfile/anime_style_age/2015s/danbooru_3726124.jpg differ diff --git a/test/testfile/anime_style_age/2015s/danbooru_3729031.jpg b/test/testfile/anime_style_age/2015s/danbooru_3729031.jpg new file mode 100644 index 00000000000..4950eff7f2a Binary files /dev/null and b/test/testfile/anime_style_age/2015s/danbooru_3729031.jpg differ diff --git a/test/testfile/anime_style_age/2015s/danbooru_3729662.jpg b/test/testfile/anime_style_age/2015s/danbooru_3729662.jpg new file mode 100644 index 00000000000..7fa2513c6d1 Binary files /dev/null and b/test/testfile/anime_style_age/2015s/danbooru_3729662.jpg differ diff --git a/test/testfile/anime_style_age/2020s/danbooru_6916723.jpg b/test/testfile/anime_style_age/2020s/danbooru_6916723.jpg new file mode 100644 index 00000000000..bec3eb55cd3 Binary files /dev/null and b/test/testfile/anime_style_age/2020s/danbooru_6916723.jpg differ diff --git a/test/testfile/anime_style_age/2020s/danbooru_6916916.jpg b/test/testfile/anime_style_age/2020s/danbooru_6916916.jpg new file mode 100644 index 00000000000..ac1ec0ea181 Binary files /dev/null and b/test/testfile/anime_style_age/2020s/danbooru_6916916.jpg differ diff --git a/test/testfile/anime_style_age/2020s/danbooru_6919611.jpg b/test/testfile/anime_style_age/2020s/danbooru_6919611.jpg new file mode 100644 index 00000000000..220e91c205f Binary files /dev/null and b/test/testfile/anime_style_age/2020s/danbooru_6919611.jpg differ diff --git a/test/testfile/anime_style_age/2020s/danbooru_6923562.jpg b/test/testfile/anime_style_age/2020s/danbooru_6923562.jpg new file mode 100644 index 00000000000..84a136b779c Binary files /dev/null and b/test/testfile/anime_style_age/2020s/danbooru_6923562.jpg differ diff --git a/test/validate/test_style_age.py b/test/validate/test_style_age.py new file mode 100644 index 00000000000..d4c32b019d4 --- /dev/null +++ b/test/validate/test_style_age.py @@ -0,0 +1,37 @@ +import glob +import os.path + +import pytest + +from imgutils.validate import anime_style_age +from imgutils.validate.style_age import _open_anime_style_age_model, anime_style_age_score +from test.testings import get_testfile + +_ROOT_DIR = get_testfile('anime_style_age') +_EXAMPLE_FILES = [ + (os.path.relpath(file, _ROOT_DIR), os.path.basename(os.path.dirname(file))) + for file in glob.glob(get_testfile('anime_style_age', '**', '*.jpg'), recursive=True) +] + + +@pytest.fixture(scope='module', autouse=True) +def _release_model_after_run(): + try: + yield + finally: + _open_anime_style_age_model.cache_clear() + + +@pytest.mark.unittest +class TestValidatePortrait: + @pytest.mark.parametrize(['image', 'label'], _EXAMPLE_FILES) + def test_anime_style_age(self, image, label): + image_file = get_testfile('anime_style_age', image) + tag, score = anime_style_age(image_file) + assert tag == label + + @pytest.mark.parametrize(['image', 'label'], _EXAMPLE_FILES) + def test_anime_style_age_score(self, image, label): + image_file = get_testfile('anime_style_age', image) + scores = anime_style_age_score(image_file) + assert scores[label] > 0.5