Unverified Commit c1d159e5 authored by narugo1992's avatar narugo1992 Committed by GitHub
Browse files

Merge pull request #39 from deepghs/dev/wd14 

dev(narugo): add benchmark for moat
parents 41ec7fe1 4598eaa6

docs/source/api_doc/tagging/tagging_benchmark.plot.py

+1 −0
Original line number Diff line number Diff line
@@ -58,6 +58,7 @@ if __name__ == '__main__': 
            ('wd14-convnext', Wd14Benchmark("ConvNext")),

            ('wd14-convnextv2', Wd14Benchmark("ConvNextV2")),

            ('wd14-vit', Wd14Benchmark("ViT")),

            ('wd14-moat', Wd14Benchmark("MOAT")),

            ('mldanbooru', MLDanbooruBenchmark()),

        ],

        title='Benchmark for Tagging Models',

docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg

+487 −505

File changed.

Preview size limit exceeded, changes collapsed.

imgutils/tagging/wd14.py

+4 −2
Original line number Diff line number Diff line
@@ -48,6 +48,7 @@ SWIN_MODEL_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2" 
CONV_MODEL_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"

CONV2_MODEL_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"

VIT_MODEL_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"

MOAT_MODEL_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"

MODEL_FILENAME = "model.onnx"

LABEL_FILENAME = "selected_tags.csv"
@@ -55,7 +56,8 @@ MODEL_NAMES = { 
    "SwinV2": SWIN_MODEL_REPO,

    "ConvNext": CONV_MODEL_REPO,

    "ConvNextV2": CONV2_MODEL_REPO,

    "ViT": VIT_MODEL_REPO

    "ViT": VIT_MODEL_REPO,

    "MOAT": MOAT_MODEL_REPO,

}
@@ -89,7 +91,7 @@ def get_wd14_tags(image: ImageTyping, model_name: str = "ConvNextV2", 


    :param image: Image to tagging.

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

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

        ``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.