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

Merge pull request #141 from deepghs/dev/deepgelbooru 

dev(narugo): add support for deepgelbooru
parents f5555612 671a4817

docs/source/api_doc/tagging/deepgelbooru.rst

0 → 100644
+14 −0
Original line number Diff line number Diff line 
imgutils.tagging.deepgelbooru

====================================



.. currentmodule:: imgutils.tagging.deepgelbooru



.. automodule:: imgutils.tagging.deepgelbooru





get_deepgelbooru_tags

------------------------



.. autofunction:: get_deepgelbooru_tags

docs/source/api_doc/tagging/index.rst

+1 −0
Original line number Diff line number Diff line
@@ -12,6 +12,7 @@ imgutils.tagging 
    mldanbooru

    wd14

    deepdanbooru

    deepgelbooru

    format

    overlap

    blacklist

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

+20 −1
Original line number Diff line number Diff line 
import random



from benchmark import BaseBenchmark, create_plot_cli

from imgutils.tagging import get_deepdanbooru_tags, get_wd14_tags, get_mldanbooru_tags

from imgutils.tagging import get_deepdanbooru_tags, get_wd14_tags, get_mldanbooru_tags, get_deepgelbooru_tags





class DeepdanbooruBenchmark(BaseBenchmark):
@@ -18,6 +18,24 @@ class DeepdanbooruBenchmark(BaseBenchmark): 
        _ = get_deepdanbooru_tags(image_file)





class DeepgelbooruBenchmark(BaseBenchmark):

    def load(self):

        from imgutils.tagging.deepgelbooru import _open_tags, _open_model, _open_preprocessor

        _ = _open_tags()

        _ = _open_model()

        _ = _open_preprocessor



    def unload(self):

        from imgutils.tagging.deepgelbooru import _open_tags, _open_model, _open_preprocessor

        _open_tags.cache_clear()

        _open_model.cache_clear()

        _open_preprocessor.cache_clear()



    def run(self):

        image_file = random.choice(self.all_images)

        _ = get_deepgelbooru_tags(image_file)





class Wd14Benchmark(BaseBenchmark):

    def __init__(self, model):

        BaseBenchmark.__init__(self)
@@ -54,6 +72,7 @@ if __name__ == '__main__': 
    create_plot_cli(

        [

            ('deepdanbooru', DeepdanbooruBenchmark()),

            ('deepgelbooru', DeepgelbooruBenchmark()),

            ('wd14-swinv2', Wd14Benchmark("SwinV2")),

            ('wd14-convnext', Wd14Benchmark("ConvNext")),

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

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

+717 −639

File changed.

Preview size limit exceeded, changes collapsed.

imgutils/tagging/__init__.py

+1 −0
Original line number Diff line number Diff line
@@ -11,6 +11,7 @@ Overview: 
from .blacklist import is_blacklisted, drop_blacklisted_tags

from .character import is_basic_character_tag, drop_basic_character_tags

from .deepdanbooru import get_deepdanbooru_tags

from .deepgelbooru import get_deepgelbooru_tags

from .format import tags_to_text, add_underline, remove_underline

from .match import tag_match_suffix, tag_match_prefix, tag_match_full

from .mldanbooru import get_mldanbooru_tags