Merge pull request #74 from deepghs/dev/tagdefault

dev(narugo): use v3 swin as default tagger
deepghs · Mar 15, 2024 · fafef8f · fafef8f
2 parents 601dc13 + b3b85f2
commit fafef8f
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 69 deletions.
diff --git a/imgutils/tagging/wd14.py b/imgutils/tagging/wd14.py
@@ -42,6 +42,7 @@
     "ConvNext_v3": CONV_V3_MODEL_REPO,
     "ViT_v3": VIT_V3_MODEL_REPO,
 }
+_DEFAULT_MODEL_NAME = 'SwinV2_v3'
 
 
 def _version_support_check(model_name):
@@ -123,7 +124,7 @@ def _prepare_image_for_tagging(image: ImageTyping, target_size: int):
 
 def get_wd14_tags(
         image: ImageTyping,
-        model_name: str = 'ConvNextV2',
+        model_name: str = _DEFAULT_MODEL_NAME,
         general_threshold: float = 0.35,
         general_mcut_enabled: bool = False,
         character_threshold: float = 0.85,

diff --git a/test/tagging/test_wd14.py b/test/tagging/test_wd14.py
@@ -17,63 +17,62 @@ def _release_model_after_run():
 class TestTaggingWd14:
     def test_get_wd14_tags(self):
         rating, tags, chars = get_wd14_tags(get_testfile('6124220.jpg'))
+
         assert rating['general'] > 0.9
-        assert tags['cat'] >= 0.8
+        assert tags['cat_girl'] >= 0.8
         assert not chars
 
         rating, tags, chars = get_wd14_tags(get_testfile('6125785.jpg'))
-        assert 0.55 <= rating['general'] <= 0.65
-        assert 0.35 <= rating['sensitive'] <= 0.45
+        assert 0.6 <= rating['general'] <= 0.8
         assert tags['1girl'] >= 0.95
         assert chars['hu_tao_(genshin_impact)'] >= 0.95
 
     def test_wd14_tags_sample(self):
         rating, tags, chars = get_wd14_tags(get_testfile('nude_girl.png'))
+
         assert rating == pytest.approx({
-            'general': 0.0020540356636047363,
-            'sensitive': 0.0080718994140625,
-            'questionable': 0.003170192241668701,
-            'explicit': 0.984081506729126,
+            'general': 0.0006683468818664551,
+            'sensitive': 0.003294050693511963,
+            'questionable': 0.0007482171058654785,
+            'explicit': 0.9922184944152832
         }, abs=2e-2)
         assert tags == pytest.approx({
-            '1girl': 0.998561441898346, 'solo': 0.9918843507766724, 'long_hair': 0.9451607465744019,
-            'breasts': 0.9867608547210693, 'looking_at_viewer': 0.9200493693351746, 'blush': 0.8876285552978516,
-            'smile': 0.5031097531318665, 'bangs': 0.4979058504104614, 'large_breasts': 0.5059964656829834,
-            'navel': 0.9681310653686523, 'hair_between_eyes': 0.5816333293914795, 'medium_breasts': 0.36410677433013916,
-            'very_long_hair': 0.811715304851532, 'closed_mouth': 0.9338403940200806, 'nipples': 0.9715133905410767,
-            'purple_eyes': 0.9681202173233032, 'collarbone': 0.573296308517456, 'nude': 0.9568941593170166,
-            'red_hair': 0.9242303967475891, 'sweat': 0.8757796287536621, 'horns': 0.973071277141571,
-            'pussy': 0.9876313805580139, 'spread_legs': 0.9634276628494263, 'armpits': 0.9116500616073608,
-            'stomach': 0.6858262419700623, 'arms_up': 0.9398491978645325, 'completely_nude': 0.907513439655304,
-            'uncensored': 0.8703584671020508, 'pussy_juice': 0.6459053754806519,
-            'feet_out_of_frame': 0.3921701908111572, 'on_bed': 0.6049470901489258,
-            'arms_behind_head': 0.4758358597755432, 'breasts_apart': 0.38581883907318115,
-            'clitoris': 0.5746099948883057
+            '1girl': 0.9994567036628723, 'solo': 0.9867788553237915, 'long_hair': 0.9705560207366943,
+            'breasts': 0.9950063228607178, 'looking_at_viewer': 0.9309853315353394, 'blush': 0.9086592793464661,
+            'smile': 0.7154737710952759, 'navel': 0.9606291055679321, 'hair_between_eyes': 0.4996751546859741,
+            'closed_mouth': 0.7993873953819275, 'very_long_hair': 0.7326497435569763,
+            'medium_breasts': 0.7169027924537659, 'nipples': 0.9904205799102783, 'purple_eyes': 0.9592539668083191,
+            'thighs': 0.4377020001411438, 'sweat': 0.6950557827949524, 'red_hair': 0.9731366038322449,
+            'nude': 0.9811137318611145, 'lying': 0.5896710157394409, 'horns': 0.9796154499053955,
+            'pussy': 0.9834838509559631, 'spread_legs': 0.9527802467346191, 'stomach': 0.7781887054443359,
+            'on_back': 0.5563216805458069, 'armpits': 0.9518307447433472, 'arms_up': 0.8266783952713013,
+            'completely_nude': 0.9168736338615417, 'pillow': 0.5372565388679504, 'uncensored': 0.9515247344970703,
+            'pussy_juice': 0.6543970108032227, 'on_bed': 0.6051450371742249, 'hair_intakes': 0.8222305178642273,
+            'demon_horns': 0.4169325828552246, 'breasts_apart': 0.45593249797821045, 'clitoris': 0.4843749403953552
         }, abs=2e-2)
-        assert chars == pytest.approx({'surtr_(arknights)': 0.9942929744720459}, abs=2e-2)
+        assert chars == pytest.approx({'surtr_(arknights)': 0.9940401315689087}, abs=2e-2)
 
     def test_wd14_tags_sample_no_underline(self):
         rating, tags, chars = get_wd14_tags(get_testfile('nude_girl.png'), no_underline=True)
         assert rating == pytest.approx({
-            'general': 0.0020540356636047363,
-            'sensitive': 0.0080718994140625,
-            'questionable': 0.003170192241668701,
-            'explicit': 0.984081506729126,
+            'general': 0.0006683468818664551,
+            'sensitive': 0.003294050693511963,
+            'questionable': 0.0007482171058654785,
+            'explicit': 0.9922184944152832
         }, abs=2e-2)
         assert tags == pytest.approx({
-            '1girl': 0.998561441898346, 'solo': 0.9918843507766724, 'long hair': 0.9451607465744019,
-            'breasts': 0.9867608547210693, 'looking at viewer': 0.9200493693351746, 'blush': 0.8876285552978516,
-            'smile': 0.5031097531318665, 'bangs': 0.4979058504104614, 'large breasts': 0.5059964656829834,
-            'navel': 0.9681310653686523, 'hair between eyes': 0.5816333293914795, 'medium breasts': 0.36410677433013916,
-            'very long hair': 0.811715304851532, 'closed mouth': 0.9338403940200806, 'nipples': 0.9715133905410767,
-            'purple eyes': 0.9681202173233032, 'collarbone': 0.573296308517456, 'nude': 0.9568941593170166,
-            'red hair': 0.9242303967475891, 'sweat': 0.8757796287536621, 'horns': 0.973071277141571,
-            'pussy': 0.9876313805580139, 'spread legs': 0.9634276628494263, 'armpits': 0.9116500616073608,
-            'stomach': 0.6858262419700623, 'arms up': 0.9398491978645325, 'completely nude': 0.907513439655304,
-            'uncensored': 0.8703584671020508, 'pussy juice': 0.6459053754806519,
-            'feet out of frame': 0.3921701908111572, 'on bed': 0.6049470901489258,
-            'arms behind head': 0.4758358597755432, 'breasts apart': 0.38581883907318115,
-            'clitoris': 0.5746099948883057
+            '1girl': 0.9994567036628723, 'solo': 0.9867788553237915, 'long hair': 0.9705560207366943,
+            'breasts': 0.9950063228607178, 'looking at viewer': 0.9309853315353394, 'blush': 0.9086592793464661,
+            'smile': 0.7154737710952759, 'navel': 0.9606291055679321, 'hair between eyes': 0.4996751546859741,
+            'closed mouth': 0.7993873953819275, 'very long hair': 0.7326497435569763,
+            'medium breasts': 0.7169027924537659, 'nipples': 0.9904205799102783, 'purple eyes': 0.9592539668083191,
+            'thighs': 0.4377020001411438, 'sweat': 0.6950557827949524, 'red hair': 0.9731366038322449,
+            'nude': 0.9811137318611145, 'lying': 0.5896710157394409, 'horns': 0.9796154499053955,
+            'pussy': 0.9834838509559631, 'spread legs': 0.9527802467346191, 'stomach': 0.7781887054443359,
+            'on back': 0.5563216805458069, 'armpits': 0.9518307447433472, 'arms up': 0.8266783952713013,
+            'completely nude': 0.9168736338615417, 'pillow': 0.5372565388679504, 'uncensored': 0.9515247344970703,
+            'pussy juice': 0.6543970108032227, 'on bed': 0.6051450371742249, 'hair intakes': 0.8222305178642273,
+            'demon horns': 0.4169325828552246, 'breasts apart': 0.45593249797821045, 'clitoris': 0.4843749403953552
         }, abs=2e-2)
         assert chars == pytest.approx({'surtr (arknights)': 0.9942929744720459}, abs=2e-2)
 
@@ -84,43 +83,39 @@ def test_wd14_tags_sample_mcut(self):
             character_mcut_enabled=True,
         )
         assert rating == pytest.approx({
-            'general': 0.0020540356636047363,
-            'sensitive': 0.0080718994140625,
-            'questionable': 0.003170192241668701,
-            'explicit': 0.984081506729126,
+            'general': 0.0006683468818664551,
+            'sensitive': 0.003294050693511963,
+            'questionable': 0.0007482171058654785,
+            'explicit': 0.9922184944152832
         }, abs=2e-2)
         assert tags == pytest.approx({
-            '1girl': 0.998561441898346, 'solo': 0.9918843507766724, 'long_hair': 0.9451607465744019,
-            'breasts': 0.9867608547210693, 'looking_at_viewer': 0.9200493693351746, 'blush': 0.8876285552978516,
-            'navel': 0.9681310653686523, 'very_long_hair': 0.811715304851532, 'closed_mouth': 0.9338403940200806,
-            'nipples': 0.9715133905410767, 'purple_eyes': 0.9681202173233032, 'nude': 0.9568941593170166,
-            'red_hair': 0.9242303967475891, 'sweat': 0.8757796287536621, 'horns': 0.973071277141571,
-            'pussy': 0.9876313805580139, 'spread_legs': 0.9634276628494263, 'armpits': 0.9116500616073608,
-            'arms_up': 0.9398491978645325, 'completely_nude': 0.907513439655304, 'uncensored': 0.8703584671020508
+            '1girl': 0.9994567036628723, 'solo': 0.9867788553237915, 'long_hair': 0.9705560207366943,
+            'breasts': 0.9950063228607178, 'looking_at_viewer': 0.9309853315353394, 'blush': 0.9086592793464661,
+            'navel': 0.9606291055679321, 'nipples': 0.9904205799102783, 'purple_eyes': 0.9592539668083191,
+            'red_hair': 0.9731366038322449, 'nude': 0.9811137318611145, 'horns': 0.9796154499053955,
+            'pussy': 0.9834838509559631, 'spread_legs': 0.9527802467346191, 'armpits': 0.9518307447433472,
+            'completely_nude': 0.9168736338615417, 'uncensored': 0.9515247344970703
         }, abs=2e-2)
-        assert chars == pytest.approx({'surtr_(arknights)': 0.9942929744720459}, abs=2e-2)
+        assert chars == pytest.approx({'surtr_(arknights)': 0.9940401315689087}, abs=2e-2)
 
     def test_wd14_tags_no_overlap(self):
         rating, tags, chars = get_wd14_tags(get_testfile('nude_girl.png'), drop_overlap=True)
-        # print(tags)
         assert rating == pytest.approx({
-            'general': 0.0020540356636047363,
-            'sensitive': 0.0080718994140625,
-            'questionable': 0.003170192241668701,
-            'explicit': 0.984081506729126,
+            'general': 0.0006683468818664551,
+            'sensitive': 0.003294050693511963,
+            'questionable': 0.0007482171058654785,
+            'explicit': 0.9922184944152832
         }, abs=2e-2)
         assert tags == pytest.approx({
-            '1girl': 0.998561441898346, 'solo': 0.9918843507766724, 'looking_at_viewer': 0.9200493693351746,
-            'blush': 0.8876285552978516, 'smile': 0.5031097531318665, 'bangs': 0.4979058504104614,
-            'large_breasts': 0.5059964656829834, 'navel': 0.9681310653686523, 'hair_between_eyes': 0.5816333293914795,
-            'medium_breasts': 0.36410677433013916, 'very_long_hair': 0.811715304851532,
-            'closed_mouth': 0.9338403940200806, 'nipples': 0.9715133905410767, 'purple_eyes': 0.9681202173233032,
-            'collarbone': 0.573296308517456, 'red_hair': 0.9242303967475891, 'sweat': 0.8757796287536621,
-            'horns': 0.973071277141571, 'spread_legs': 0.9634276628494263, 'armpits': 0.9116500616073608,
-            'stomach': 0.6858262419700623, 'arms_up': 0.9398491978645325, 'completely_nude': 0.907513439655304,
-            'uncensored': 0.8703584671020508, 'pussy_juice': 0.6459053754806519,
-            'feet_out_of_frame': 0.3921701908111572, 'on_bed': 0.6049470901489258,
-            'arms_behind_head': 0.4758358597755432, 'breasts_apart': 0.38581883907318115,
-            'clitoris': 0.5746099948883057
+            '1girl': 0.9994567036628723, 'solo': 0.9867788553237915, 'looking_at_viewer': 0.9309853315353394,
+            'blush': 0.9086592793464661, 'smile': 0.7154737710952759, 'navel': 0.9606291055679321,
+            'hair_between_eyes': 0.4996751546859741, 'closed_mouth': 0.7993873953819275,
+            'very_long_hair': 0.7326497435569763, 'medium_breasts': 0.7169027924537659, 'nipples': 0.9904205799102783,
+            'purple_eyes': 0.9592539668083191, 'thighs': 0.4377020001411438, 'sweat': 0.6950557827949524,
+            'red_hair': 0.9731366038322449, 'spread_legs': 0.9527802467346191, 'stomach': 0.7781887054443359,
+            'on_back': 0.5563216805458069, 'armpits': 0.9518307447433472, 'arms_up': 0.8266783952713013,
+            'completely_nude': 0.9168736338615417, 'pillow': 0.5372565388679504, 'uncensored': 0.9515247344970703,
+            'pussy_juice': 0.6543970108032227, 'on_bed': 0.6051450371742249, 'hair_intakes': 0.8222306370735168,
+            'demon_horns': 0.4169325828552246, 'breasts_apart': 0.45593249797821045, 'clitoris': 0.4843749403953552
         }, abs=2e-2)
-        assert chars == pytest.approx({'surtr_(arknights)': 0.9942929744720459}, abs=2e-2)
+        assert chars == pytest.approx({'surtr_(arknights)': 0.9940401315689087}, abs=2e-2)