Loading imgutils/tagging/wd14.py +2 −1 Original line number Diff line number Diff line Loading @@ -42,6 +42,7 @@ MODEL_NAMES = { "ConvNext_v3": CONV_V3_MODEL_REPO, "ViT_v3": VIT_V3_MODEL_REPO, } _DEFAULT_MODEL_NAME = 'SwinV2_v3' def _version_support_check(model_name): Loading Loading @@ -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, Loading test/tagging/test_wd14.py +63 −68 Original line number Diff line number Diff line Loading @@ -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) Loading @@ -84,43 +83,39 @@ class TestTaggingWd14: 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) Loading
imgutils/tagging/wd14.py +2 −1 Original line number Diff line number Diff line Loading @@ -42,6 +42,7 @@ MODEL_NAMES = { "ConvNext_v3": CONV_V3_MODEL_REPO, "ViT_v3": VIT_V3_MODEL_REPO, } _DEFAULT_MODEL_NAME = 'SwinV2_v3' def _version_support_check(model_name): Loading Loading @@ -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, Loading
test/tagging/test_wd14.py +63 −68 Original line number Diff line number Diff line Loading @@ -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) Loading @@ -84,43 +83,39 @@ class TestTaggingWd14: 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)
