Commit 4598eaa6 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): use convnextv2 back

parent 81e8f63d

imgutils/tagging/wd14.py

+2 −2
Original line number Diff line number Diff line
@@ -82,7 +82,7 @@ def _get_wd14_labels() -> Tuple[List[str], List[int], List[int], List[int]]: 
    return tag_names, rating_indexes, general_indexes, character_indexes





def get_wd14_tags(image: ImageTyping, model_name: str = "MOAT",

def get_wd14_tags(image: ImageTyping, model_name: str = "ConvNextV2",

                  general_threshold: float = 0.35, character_threshold: float = 0.85):

    """

    Overview:
@@ -91,7 +91,7 @@ def get_wd14_tags(image: ImageTyping, model_name: str = "MOAT", 


    :param image: Image to tagging.

    :param model_name: Name of the mode, should be one of the \

        ``SwinV2``, ``ConvNext``, ``ConvNextV2``, ``ViT`` or ``MOAT``, default is ``MOAT``.

        ``SwinV2``, ``ConvNext``, ``ConvNextV2``, ``ViT`` or ``MOAT``, default is ``ConvNextV2``.

    :param general_threshold: Threshold for default tags, default is ``0.35``.

    :param character_threshold: Threshold for character tags, default is ``0.85``.

    :return: Tagging results for levels, features and characters.

test/tagging/test_wd14.py

+2 −2
Original line number Diff line number Diff line
@@ -22,7 +22,7 @@ class TestTaggingWd14: 
        assert not chars



        rating, tags, chars = get_wd14_tags(get_testfile('6125785.jpg'))

        assert 0.55 <= rating['general']

        assert 0.20 <= rating['sensitive']

        assert 0.55 <= rating['general'] <= 0.65

        assert 0.35 <= rating['sensitive'] <= 0.45

        assert tags['1girl'] >= 0.95

        assert chars['hu_tao_(genshin_impact)'] >= 0.95