dev(narugo): use convnextv2 back

deepghs · Aug 30, 2023 · 4598eaa · 4598eaa
1 parent 81e8f63
commit 4598eaa
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/imgutils/tagging/wd14.py b/imgutils/tagging/wd14.py
@@ -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.

diff --git a/test/tagging/test_wd14.py b/test/tagging/test_wd14.py
@@ -22,7 +22,7 @@ def test_get_wd14_tags(self):
         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