Commit 9b8cf3d0 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): add unittest for imgutils.validate.rating

parent 01d99b1d

imgutils/validate/rating.py

+55 −0
Original line number Diff line number Diff line
@@ -98,6 +98,34 @@ def anime_rating_score(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME 
    :param model_name: Model to use. Default is ``mobilenetv3_sce_dist``. All available models are listed

        on the benchmark plot above. If you need better accuracy, just set this to ``caformer_s36_plus``.

    :return: A dict with ratings and scores.



    Example::

        >>> from imgutils.validate import anime_rating_score

        >>>

        >>> anime_rating_score('rating/safe/1.jpg')

        {'safe': 0.9999998807907104, 'r15': 2.5863172936624323e-08, 'r18': 6.480062353375615e-08}

        >>> anime_rating_score('rating/safe/2.jpg')

        {'safe': 0.9924363493919373, 'r15': 0.007255776319652796, 'r18': 0.0003077814180869609}

        >>> anime_rating_score('rating/safe/3.jpg')

        {'safe': 0.996969997882843, 'r15': 0.0030054834205657244, 'r18': 2.4601260520284995e-05}

        >>> anime_rating_score('rating/safe/4.jpg')

        {'safe': 0.9966891407966614, 'r15': 0.003293127752840519, 'r18': 1.770909148035571e-05}

        >>> anime_rating_score('rating/r15/5.jpg')

        {'safe': 0.00025384966284036636, 'r15': 0.9996721744537354, 'r18': 7.399192691082135e-05}

        >>> anime_rating_score('rating/r15/6.jpg')

        {'safe': 7.973351603141055e-05, 'r15': 0.9998563528060913, 'r18': 6.391309580067173e-05}

        >>> anime_rating_score('rating/r15/7.jpg')

        {'safe': 0.0018681309884414077, 'r15': 0.9827859997749329, 'r18': 0.015345841646194458}

        >>> anime_rating_score('rating/r15/8.jpg')

        {'safe': 0.013710384257137775, 'r15': 0.8339558839797974, 'r18': 0.15233369171619415}

        >>> anime_rating_score('rating/r18/9.jpg')

        {'safe': 3.951323833462084e-06, 'r15': 0.00029566374723799527, 'r18': 0.9997004270553589}

        >>> anime_rating_score('rating/r18/10.jpg')

        {'safe': 0.00018434497178532183, 'r15': 4.568440272123553e-05, 'r18': 0.9997699856758118}

        >>> anime_rating_score('rating/r18/11.jpg')

        {'safe': 9.11225129129889e-07, 'r15': 5.051862899563275e-05, 'r18': 0.9999485015869141}

        >>> anime_rating_score('rating/r18/12.jpg')

        {'safe': 6.902020231791539e-06, 'r15': 0.0005639699520543218, 'r18': 0.9994290471076965}

    """

    output = _raw_anime_rating(image, model_name)

    values = dict(zip(_open_anime_rating_labels(model_name), map(lambda x: x.item(), output[0])))
@@ -114,6 +142,33 @@ def anime_rating(image: ImageTyping, model_name: str = _DEFAULT_MODEL_NAME) -> T 
        on the benchmark plot above. If you need better accuracy, just set this to ``caformer_s36_plus``.

    :return: A tuple contains the rating and its score.



    Examples::

        >>> from imgutils.validate import anime_rating

        >>>

        >>> anime_rating('rating/safe/1.jpg')

        ('safe', 0.9999998807907104)

        >>> anime_rating('rating/safe/2.jpg')

        ('safe', 0.9924363493919373)

        >>> anime_rating('rating/safe/3.jpg')

        ('safe', 0.996969997882843)

        >>> anime_rating('rating/safe/4.jpg')

        ('safe', 0.9966891407966614)

        >>> anime_rating('rating/r15/5.jpg')

        ('r15', 0.9996721744537354)

        >>> anime_rating('rating/r15/6.jpg')

        ('r15', 0.9998563528060913)

        >>> anime_rating('rating/r15/7.jpg')

        ('r15', 0.9827859997749329)

        >>> anime_rating('rating/r15/8.jpg')

        ('r15', 0.8339558839797974)

        >>> anime_rating('rating/r18/9.jpg')

        ('r18', 0.9997004270553589)

        >>> anime_rating('rating/r18/10.jpg')

        ('r18', 0.9997699856758118)

        >>> anime_rating('rating/r18/11.jpg')

        ('r18', 0.9999485015869141)

        >>> anime_rating('rating/r18/12.jpg')

        ('r18', 0.9994290471076965)

    """

    output = _raw_anime_rating(image, model_name)[0]

    max_id = np.argmax(output)

test/validate/test_rating.py

0 → 100644
+37 −0
Original line number Diff line number Diff line 
import glob

import os.path



import pytest



from imgutils.validate import anime_rating

from imgutils.validate.rating import _open_anime_rating_model, anime_rating_score

from test.testings import get_testfile



_ROOT_DIR = get_testfile('rating')

_EXAMPLE_FILES = [

    (os.path.relpath(file, _ROOT_DIR), os.path.basename(os.path.dirname(file)))

    for file in glob.glob(get_testfile('rating', '**', '*.jpg'), recursive=True)

]





@pytest.fixture(scope='module', autouse=True)

def _release_model_after_run():

    try:

        yield

    finally:

        _open_anime_rating_model.cache_clear()





@pytest.mark.unittest

class TestValidateRating:

    @pytest.mark.parametrize(['image', 'label'], _EXAMPLE_FILES)

    def test_anime_rating(self, image, label):

        image_file = get_testfile('rating', image)

        tag, score = anime_rating(image_file)

        assert tag == label



    @pytest.mark.parametrize(['image', 'label'], _EXAMPLE_FILES)

    def test_anime_rating_score(self, image, label):

        image_file = get_testfile('rating', image)

        scores = anime_rating_score(image_file)

        assert scores[label] > 0.5