diff --git a/.github/workflows/doc.yml b/.github/workflows/doc.yml index 3417574ba98..2df2b2a4dab 100644 --- a/.github/workflows/doc.yml +++ b/.github/workflows/doc.yml @@ -67,6 +67,7 @@ jobs: ENV_PROD: 'true' PLANTUML_HOST: http://localhost:18080 HF_TOKEN: ${{ secrets.HF_TOKEN }} + CI: 'true' run: | plantumlcli -c make docs @@ -145,6 +146,7 @@ jobs: ENV_PROD: 'true' PLANTUML_HOST: http://localhost:18080 HF_TOKEN: ${{ secrets.HF_TOKEN }} + CI: 'true' run: | git fetch --all --tags git branch -av diff --git a/Makefile b/Makefile index 49fe1f7d5a9..044dbff10c9 100644 --- a/Makefile +++ b/Makefile @@ -41,7 +41,7 @@ unittest: --cov="${RANGE_SRC_DIR}" \ $(if ${MIN_COVERAGE},--cov-fail-under=${MIN_COVERAGE},) \ $(if ${WORKERS},-n ${WORKERS},) \ - --reruns 8 --reruns-delay 2 --only-rerun '(OSError|Timeout|HTTPError.*429|HTTPError.*502|HTTPError.*504|check your connection|429 error)' + --reruns 8 --reruns-delay 2 --only-rerun '(OSError|Timeout|HTTPError.*429|HTTPError.*502|HTTPError.*504|check your connection|429 error|CAS service error)' docs: $(MAKE) -C "${DOC_DIR}" build diff --git a/docs/source/_libs/benchmark.py b/docs/source/_libs/benchmark.py index 7f7bf3c4430..d2d02a54f92 100644 --- a/docs/source/_libs/benchmark.py +++ b/docs/source/_libs/benchmark.py @@ -34,6 +34,9 @@ def load(self): def unload(self): raise NotImplementedError + def after_unload(self): + pass + def run(self): raise NotImplementedError @@ -60,6 +63,7 @@ def _record(name): self.unload() _record('') + self.after_unload() mems = np.array([mem for _, mem, _ in logs]) mems -= mems[0] diff --git a/docs/source/api_doc/tagging/index.rst b/docs/source/api_doc/tagging/index.rst index 34c51a733c8..b6213451a3f 100644 --- a/docs/source/api_doc/tagging/index.rst +++ b/docs/source/api_doc/tagging/index.rst @@ -12,6 +12,7 @@ imgutils.tagging mldanbooru wd14 camie + pixai deepdanbooru deepgelbooru format diff --git a/docs/source/api_doc/tagging/pixai.rst b/docs/source/api_doc/tagging/pixai.rst new file mode 100644 index 00000000000..b87c782d432 --- /dev/null +++ b/docs/source/api_doc/tagging/pixai.rst @@ -0,0 +1,15 @@ +imgutils.tagging.pixai +==================================== + +.. currentmodule:: imgutils.tagging.pixai + +.. automodule:: imgutils.tagging.pixai + + +get_pixai_tags +---------------------- + +.. autofunction:: get_pixai_tags + + + diff --git a/docs/source/api_doc/tagging/tagging_benchmark.plot.py b/docs/source/api_doc/tagging/tagging_benchmark.plot.py index e110be999bc..b81ebf47b90 100644 --- a/docs/source/api_doc/tagging/tagging_benchmark.plot.py +++ b/docs/source/api_doc/tagging/tagging_benchmark.plot.py @@ -1,25 +1,36 @@ +import os import random +from hfutils.cache import delete_cache + from benchmark import BaseBenchmark, create_plot_cli from imgutils.tagging import get_deepdanbooru_tags, get_wd14_tags, get_mldanbooru_tags, get_deepgelbooru_tags, \ - get_camie_tags + get_camie_tags, get_pixai_tags + +class CleanModelStorageBenchmark(BaseBenchmark): + def after_unload(self): + if os.environ.get('CI'): + delete_cache() -class DeepdanbooruBenchmark(BaseBenchmark): + +class DeepdanbooruBenchmark(CleanModelStorageBenchmark): def load(self): - from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model + from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model, _get_deepdanbooru_labels _ = _get_deepdanbooru_model() + _ = _get_deepdanbooru_labels def unload(self): - from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model + from imgutils.tagging.deepdanbooru import _get_deepdanbooru_model, _get_deepdanbooru_labels _get_deepdanbooru_model.cache_clear() + _get_deepdanbooru_labels.cache_clear() def run(self): image_file = random.choice(self.all_images) _ = get_deepdanbooru_tags(image_file) -class DeepgelbooruBenchmark(BaseBenchmark): +class DeepgelbooruBenchmark(CleanModelStorageBenchmark): def load(self): from imgutils.tagging.deepgelbooru import _open_tags, _open_model, _open_preprocessor _ = _open_tags() @@ -37,39 +48,43 @@ def run(self): _ = get_deepgelbooru_tags(image_file) -class Wd14Benchmark(BaseBenchmark): +class Wd14Benchmark(CleanModelStorageBenchmark): def __init__(self, model): BaseBenchmark.__init__(self) self.model = model def load(self): - from imgutils.tagging.wd14 import _get_wd14_model + from imgutils.tagging.wd14 import _get_wd14_model, _get_wd14_labels _ = _get_wd14_model(self.model) + _ = _get_wd14_labels(self.model) def unload(self): - from imgutils.tagging.wd14 import _get_wd14_model + from imgutils.tagging.wd14 import _get_wd14_model, _get_wd14_labels _get_wd14_model.cache_clear() + _get_wd14_labels.cache_clear() def run(self): image_file = random.choice(self.all_images) _ = get_wd14_tags(image_file, model_name=self.model) -class MLDanbooruBenchmark(BaseBenchmark): +class MLDanbooruBenchmark(CleanModelStorageBenchmark): def load(self): - from imgutils.tagging.mldanbooru import _open_mldanbooru_model + from imgutils.tagging.mldanbooru import _open_mldanbooru_model, _get_mldanbooru_labels _ = _open_mldanbooru_model() + _ = _get_mldanbooru_labels() def unload(self): - from imgutils.tagging.mldanbooru import _open_mldanbooru_model + from imgutils.tagging.mldanbooru import _open_mldanbooru_model, _get_mldanbooru_labels _open_mldanbooru_model.cache_clear() + _get_mldanbooru_labels.cache_clear() def run(self): image_file = random.choice(self.all_images) _ = get_mldanbooru_tags(image_file) -class CamieBenchmark(BaseBenchmark): +class CamieBenchmark(CleanModelStorageBenchmark): def __init__(self, model_name): BaseBenchmark.__init__(self) self.model_name = model_name @@ -95,6 +110,32 @@ def run(self): _ = get_camie_tags(image_file, model_name=self.model_name) +class PixAITaggerBenchmark(CleanModelStorageBenchmark): + def __init__(self, model_name: str): + BaseBenchmark.__init__(self) + self.model_name = model_name + + def load(self): + from imgutils.tagging.pixai import _open_tags, _open_preprocess, _open_onnx_model, \ + _open_default_category_thresholds + _ = _open_tags(self.model_name) + _ = _open_preprocess(self.model_name) + _ = _open_onnx_model(self.model_name) + _ = _open_default_category_thresholds(self.model_name) + + def unload(self): + from imgutils.tagging.pixai import _open_tags, _open_preprocess, _open_onnx_model, \ + _open_default_category_thresholds + _open_tags.cache_clear() + _open_preprocess.cache_clear() + _open_onnx_model.cache_clear() + _open_default_category_thresholds.cache_clear() + + def run(self): + image_file = random.choice(self.all_images) + _ = get_pixai_tags(image_file, model_name=self.model_name) + + if __name__ == '__main__': create_plot_cli( [ @@ -113,6 +154,7 @@ def run(self): ('mldanbooru', MLDanbooruBenchmark()), ('camie-initial', CamieBenchmark('initial')), ('camie-refined', CamieBenchmark('refined')), + ('pixai-tagger-v0.9', PixAITaggerBenchmark('v0.9')), ], title='Benchmark for Tagging Models', run_times=10, diff --git a/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg b/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg index d86cc8a7681..071703ca40f 100644 --- a/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg +++ b/docs/source/api_doc/tagging/tagging_benchmark.plot.py.svg @@ -6,7 +6,7 @@ - 2025-03-29T17:54:03.129689 + 2025-09-09T10:30:22.529066 image/svg+xml @@ -44,16 +44,16 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - - + @@ -147,11 +147,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -259,11 +259,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -331,11 +331,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -376,11 +376,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -429,11 +429,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -469,11 +469,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -515,11 +515,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -566,11 +566,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -597,11 +597,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -657,11 +657,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -708,11 +708,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -751,11 +751,11 @@ z +" clip-path="url(#p643477681f)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + @@ -800,23 +800,23 @@ z - + - - + - + - + - + - + - + - + - + - + - + - + @@ -990,18 +990,18 @@ L 407.106562 207.931672 - + - + - + @@ -1014,18 +1014,18 @@ L 407.106562 155.509411 - + - + - + @@ -1228,244 +1228,260 @@ z - + - + - + - + - + - + - + - + - + - + - + - + - + - + - +L 392.144972 343.743542 +" clip-path="url(#p643477681f)" style="fill: none; stroke: #9467bd; stroke-width: 1.5; stroke-linecap: square"/> + + + - - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -2083,17 +2127,17 @@ z - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2106,17 +2150,17 @@ L 481.128153 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2129,17 +2173,17 @@ L 506.064138 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2148,17 +2192,17 @@ L 531.000123 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2167,17 +2211,17 @@ L 555.936108 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2186,17 +2230,17 @@ L 580.872093 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2205,17 +2249,17 @@ L 605.808078 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2224,17 +2268,17 @@ L 630.744062 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2243,17 +2287,17 @@ L 655.680047 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2262,17 +2306,17 @@ L 680.616032 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2281,17 +2325,17 @@ L 705.552017 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2300,17 +2344,17 @@ L 730.488002 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2320,17 +2364,17 @@ L 755.423987 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2347,17 +2391,17 @@ L 780.359972 63.816 - + +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square"/> - + - + - + @@ -2369,19 +2413,19 @@ L 795.321562 365.198455 - - + + - + - + - + - + @@ -2390,19 +2434,19 @@ L 795.321562 304.027578 - - + + - + - + - + - + @@ -2411,19 +2455,19 @@ L 795.321562 242.856702 - - + + - + - + - + - + @@ -2432,19 +2476,19 @@ L 795.321562 181.685826 - - + + - + - + - + - + @@ -2452,7 +2496,49 @@ L 795.321562 120.51495 - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -2501,245 +2587,261 @@ z - + +L 506.064138 323.925781 +L 531.000123 331.615286 +L 555.936108 340.210985 +L 580.872093 340.190535 +L 605.808078 340.107476 +L 630.744062 340.010794 +L 655.680047 340.111491 +L 680.616032 340.202517 +L 705.552017 340.10741 +L 730.488002 340.186755 +L 755.423987 340.153149 +L 780.359972 364.535717 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 312.355719 +L 531.000123 330.985934 +L 555.936108 340.176385 +L 580.872093 340.110622 +L 605.808078 340.2114 +L 630.744062 340.034903 +L 655.680047 340.035876 +L 680.616032 339.831517 +L 705.552017 340.146997 +L 730.488002 340.081009 +L 755.423987 340.094924 +L 780.359972 364.324706 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #ff7f0e; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 328.39418 +L 531.000123 315.955263 +L 555.936108 323.72465 +L 580.872093 324.152369 +L 605.808078 324.352908 +L 630.744062 324.133757 +L 655.680047 324.3429 +L 680.616032 324.320158 +L 705.552017 323.981704 +L 730.488002 324.284143 +L 755.423987 324.440268 +L 780.359972 364.560292 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #2ca02c; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 338.709806 +L 531.000123 319.685448 +L 555.936108 325.512944 +L 580.872093 325.621093 +L 605.808078 325.455635 +L 630.744062 325.708953 +L 655.680047 325.517984 +L 680.616032 325.445645 +L 705.552017 325.749403 +L 730.488002 325.541119 +L 755.423987 325.365627 +L 780.359972 364.162058 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #d62728; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 340.066557 +L 531.000123 318.245721 +L 555.936108 323.965042 +L 580.872093 324.097989 +L 605.808078 324.061633 +L 630.744062 323.967815 +L 655.680047 323.990237 +L 680.616032 323.932075 +L 705.552017 323.850727 +L 730.488002 323.510312 +L 755.423987 323.777315 +L 780.359972 364.522371 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #9467bd; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 343.877845 +L 531.000123 309.889321 +L 555.936108 315.446474 +L 580.872093 315.443192 +L 605.808078 315.498092 +L 630.744062 315.36341 +L 655.680047 315.462227 +L 680.616032 315.529157 +L 705.552017 315.385682 +L 730.488002 315.486813 +L 755.423987 315.33375 +L 780.359972 364.862459 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #8c564b; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 342.547758 +L 531.000123 308.245654 +L 555.936108 314.380257 +L 580.872093 314.243591 +L 605.808078 314.261339 +L 630.744062 314.355696 +L 655.680047 314.100247 +L 680.616032 314.417976 +L 705.552017 314.250372 +L 730.488002 314.311167 +L 755.423987 314.418134 +L 780.359972 364.522389 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #e377c2; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 327.082866 +L 531.000123 316.54407 +L 555.936108 322.706101 +L 580.872093 323.081296 +L 605.808078 322.735183 +L 630.744062 322.940496 +L 655.680047 323.101827 +L 680.616032 323.059728 +L 705.552017 323.142008 +L 730.488002 322.747859 +L 755.423987 322.845423 +L 780.359972 364.608142 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #7f7f7f; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 345.624355 +L 531.000123 316.347401 +L 555.936108 322.844031 +L 580.872093 322.786863 +L 605.808078 322.811832 +L 630.744062 322.758347 +L 655.680047 322.852199 +L 680.616032 322.821984 +L 705.552017 322.816273 +L 730.488002 322.711109 +L 755.423987 322.650508 +L 780.359972 364.973536 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #bcbd22; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 345.329644 +L 531.000123 321.672699 +L 555.936108 327.15374 +L 580.872093 327.492074 +L 605.808078 327.475968 +L 630.744062 327.459347 +L 655.680047 327.150824 +L 680.616032 327.506319 +L 705.552017 327.409542 +L 730.488002 327.227013 +L 755.423987 327.450393 +L 780.359972 364.830651 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #17becf; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 326.246943 +L 531.000123 152.950723 +L 555.936108 159.989088 +L 580.872093 160.738507 +L 605.808078 160.845645 +L 630.744062 160.727759 +L 655.680047 160.625849 +L 680.616032 160.919706 +L 705.552017 160.24756 +L 730.488002 159.491678 +L 755.423987 160.933349 +L 780.359972 364.588247 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 311.460964 +L 531.000123 147.586551 +L 555.936108 157.673699 +L 580.872093 159.377328 +L 605.808078 159.546118 +L 630.744062 159.866107 +L 655.680047 159.623068 +L 680.616032 160.040511 +L 705.552017 158.762889 +L 730.488002 158.629986 +L 755.423987 159.839427 +L 780.359972 364.168451 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #ff7f0e; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 319.509081 +L 531.000123 306.458904 +L 555.936108 312.310242 +L 580.872093 312.575749 +L 605.808078 312.749693 +L 630.744062 312.838507 +L 655.680047 312.551562 +L 680.616032 312.611093 +L 705.552017 312.746332 +L 730.488002 312.690496 +L 755.423987 312.192692 +L 780.359972 364.336582 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #2ca02c; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 305.895076 +L 531.000123 269.960015 +L 555.936108 325.584458 +L 580.872093 325.566958 +L 605.808078 325.533879 +L 630.744062 325.671947 +L 655.680047 325.620817 +L 680.616032 325.554297 +L 705.552017 325.618952 +L 730.488002 325.49596 +L 755.423987 325.600414 +L 780.359972 363.368163 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #d62728; stroke-width: 1.5; stroke-linecap: square"/> - + +L 506.064138 282.253284 +L 531.000123 236.299229 +L 555.936108 321.384693 +L 580.872093 320.696932 +L 605.808078 322.453292 +L 630.744062 323.36962 +L 655.680047 322.907081 +L 680.616032 323.081481 +L 705.552017 322.907283 +L 730.488002 323.179844 +L 755.423987 323.49364 +L 780.359972 361.990201 +" clip-path="url(#p3b71c1a520)" style="fill: none; stroke: #9467bd; stroke-width: 1.5; stroke-linecap: square"/> + + + - + @@ -2844,25 +2946,25 @@ z - - + - + @@ -2879,13 +2981,13 @@ L 648.704375 76.914437 - + - + @@ -2902,13 +3004,13 @@ L 648.704375 91.592562 - + - + @@ -2924,13 +3026,13 @@ L 648.704375 106.270687 - + - + @@ -2948,13 +3050,13 @@ L 648.704375 120.948812 - + - + @@ -2974,13 +3076,13 @@ L 648.704375 135.626937 - + - + @@ -2993,13 +3095,13 @@ L 648.704375 150.305062 - + - + @@ -3013,13 +3115,13 @@ L 648.704375 164.983187 - + - + @@ -3036,13 +3138,13 @@ L 648.704375 179.661312 - + - + @@ -3056,13 +3158,13 @@ L 648.704375 194.339437 - + - + @@ -3081,13 +3183,13 @@ L 648.704375 209.017562 - + - + @@ -3114,13 +3216,13 @@ L 648.704375 223.695687 - + - + @@ -3149,13 +3251,13 @@ L 648.704375 238.373812 - + - + @@ -3170,13 +3272,13 @@ L 648.704375 253.051937 - + - + @@ -3194,13 +3296,13 @@ L 648.704375 267.730062 - + - + @@ -3218,9 +3320,37 @@ L 648.704375 282.408187 + + + + + + + + + + + + + + + + + + + + + + + + + - + @@ -3305,10 +3435,10 @@ z - + - + diff --git a/imgutils/generic/multilabel_timm.py b/imgutils/generic/multilabel_timm.py index 294394a0714..b1e3a3e7623 100644 --- a/imgutils/generic/multilabel_timm.py +++ b/imgutils/generic/multilabel_timm.py @@ -29,6 +29,7 @@ from typing import Optional, Literal, Dict, Any, Union import pandas as pd +from hbutils.design import SingletonMark from hbutils.string import titleize from hfutils.repository import hf_hub_repo_url from huggingface_hub import hf_hub_download @@ -64,7 +65,7 @@ def _check_gradio_env(): f'Please install it with `pip install dghs-imgutils[demo]`.') -FMT_UNSET = object() +FMT_UNSET = SingletonMark('FMT_UNSET') class MultiLabelTIMMModel: diff --git a/imgutils/tagging/__init__.py b/imgutils/tagging/__init__.py index 289dc6775e2..f4124f1c41e 100644 --- a/imgutils/tagging/__init__.py +++ b/imgutils/tagging/__init__.py @@ -18,4 +18,5 @@ from .mldanbooru import get_mldanbooru_tags from .order import sort_tags from .overlap import drop_overlap_tags +from .pixai import get_pixai_tags from .wd14 import get_wd14_tags, convert_wd14_emb_to_prediction, denormalize_wd14_emb diff --git a/imgutils/tagging/pixai.py b/imgutils/tagging/pixai.py new file mode 100644 index 00000000000..2063bc3eb3d --- /dev/null +++ b/imgutils/tagging/pixai.py @@ -0,0 +1,363 @@ +""" +Overview: + This module provides utilities for image tagging using PixAI taggers, which are specialized models + for analyzing anime-style images and extracting relevant tags. The module supports loading ONNX + models from Hugging Face Hub and processing images to generate categorized tags with confidence scores. + + The models are originally developed by the PixAI team and available at + `pixai-labs `_ on Hugging Face. This module uses ONNX-converted + versions of these models for efficient inference, available at + `deepghs `_ repositories. + + In addition to standard tagging, the models can identify anime character IP (Intellectual Property) + associations. For example, if a character like "misaka_mikoto" is detected, the system can map + this to the "toaru_kagaku_no_railgun" (A Certain Scientific Railgun) IP. All IP names follow + Danbooru-style tag conventions for consistency with existing anime tagging systems. + + Example:: + >>> from imgutils.tagging.pixai import get_pixai_tags + >>> # Get tags with default thresholds + >>> result = get_pixai_tags('path/to/anime_image.jpg', model_name='v0.9') + >>> general_tags, character_tags = result + >>> print("General tags:", general_tags) + >>> print("Character tags:", character_tags) + + >>> # Get all tags in a single dictionary + >>> all_tags = get_pixai_tags('path/to/image.jpg', fmt='tag') + >>> print("All tags:", all_tags) + + >>> # Get IP information for detected characters + >>> ips = get_pixai_tags('path/to/image.jpg', fmt='ips') + >>> print("Detected IPs:", ips) +""" + +import json +from collections import defaultdict +from typing import Union, Dict, Any, Tuple, List + +import pandas as pd +from hbutils.design import SingletonMark +from huggingface_hub import hf_hub_download +from huggingface_hub.errors import EntryNotFoundError + +from imgutils.data import ImageTyping, load_image +from imgutils.preprocess import create_pillow_transforms +from imgutils.utils import open_onnx_model, ts_lru_cache, vreplace + +FMT_UNSET = SingletonMark('FMT_UNSET') + + +def _get_repo_id(model_name: str) -> str: + """ + Get the repository ID for the specified model name. + + :param model_name: Name of the model (e.g., 'v0.9') or full repository path + :type model_name: str + + :return: Full repository ID for Hugging Face Hub + :rtype: str + + Example:: + >>> _get_repo_id('v0.9') + 'deepghs/pixai-tagger-v0.9-onnx' + >>> _get_repo_id('custom/model-repo') + 'custom/model-repo' + """ + if '/' in model_name: + return model_name + else: + return f'deepghs/pixai-tagger-{model_name}-onnx' + + +@ts_lru_cache() +def _open_onnx_model(model_name: str): + """ + Load the ONNX model from Hugging Face Hub with caching. + + This function downloads and loads the ONNX model file for the specified PixAI tagger. + Results are cached to avoid repeated downloads and model loading. + + :param model_name: Name of the model to load + :type model_name: str + + :return: The loaded ONNX model session + :rtype: onnxruntime.InferenceSession + """ + return open_onnx_model(hf_hub_download( + repo_id=_get_repo_id(model_name), + repo_type='model', + filename='model.onnx', + )) + + +@ts_lru_cache() +def _open_tags(model_name: str) -> Tuple[pd.DataFrame, Dict[str, List[str]]]: + """ + Load the tag metadata from Hugging Face Hub with caching. + + This function downloads and loads the CSV file containing tag names, categories, + and IP (Intellectual Property) associations for the specified model. The DataFrame + contains columns for tag names, categories, and other metadata including character + IP mappings when available. + + :param model_name: Name of the model + :type model_name: str + + :return: Tuple containing (DataFrame with tag information, dictionary mapping character tags to their IPs) + :rtype: Tuple[pd.DataFrame, Dict[str, List[str]]] + """ + df_tags = pd.read_csv(hf_hub_download( + repo_id=_get_repo_id(model_name), + repo_type='model', + filename='selected_tags.csv', + )) + d_ips = {} + if 'ips' in df_tags: + df_tags['ips'] = df_tags['ips'].map(json.loads) + for name, ips in zip(df_tags['name'], df_tags['ips']): + if ips: + d_ips[name] = ips + return df_tags, d_ips + + +@ts_lru_cache() +def _open_preprocess(model_name: str): + """ + Load the preprocessing pipeline configuration from Hugging Face Hub with caching. + + This function downloads the preprocessing configuration and creates a PIL transforms + pipeline for image preprocessing before model inference. + + :param model_name: Name of the model + :type model_name: str + + :return: Preprocessing transform pipeline + """ + with open(hf_hub_download( + repo_id=_get_repo_id(model_name), + repo_type='model', + filename='preprocess.json' + ), 'r') as f: + data_ = json.load(f) + return create_pillow_transforms(data_['stages']) + + +@ts_lru_cache() +def _open_default_category_thresholds(model_name: str) -> Tuple[Dict[int, float], Dict[int, str]]: + """ + Load default category thresholds and names from the Hugging Face Hub with caching. + + This function attempts to load predefined threshold values for each category from + a CSV file. If the file doesn't exist, empty dictionaries are returned. + + :param model_name: Name of the model + :type model_name: str + + :return: Tuple containing (category_thresholds, category_names) dictionaries + :rtype: tuple[Dict[int, float], Dict[int, str]] + + Example:: + >>> thresholds, names = _open_default_category_thresholds('v0.9') + >>> print(thresholds) # {0: 0.35, 1: 0.4, ...} + >>> print(names) # {0: 'general', 1: 'character', ...} + """ + _default_category_thresholds: Dict[int, float] = {} + _category_names: Dict[int, str] = {} + try: + df_category_thresholds = pd.read_csv(hf_hub_download( + repo_id=_get_repo_id(model_name), + repo_type='model', + filename='thresholds.csv' + ), keep_default_na=False) + except (EntryNotFoundError,): + pass + else: + for item in df_category_thresholds.to_dict('records'): + if item['category'] not in _default_category_thresholds: + _default_category_thresholds[item['category']] = item['threshold'] + _category_names[item['category']] = item['name'] + + return _default_category_thresholds, _category_names + + +def _raw_predict(image: ImageTyping, model_name: str): + """ + Make a raw prediction with the PixAI tagger model. + + This function preprocesses the input image and runs inference using the specified + ONNX model. It returns the raw model outputs without any post-processing or + threshold application. + + :param image: The input image to analyze + :type image: ImageTyping + :param model_name: Name of the model to use for prediction + :type model_name: str + + :return: Dictionary containing raw model outputs with keys like 'prediction', 'embedding', 'logits' + :rtype: dict + + Example:: + >>> raw_output = _raw_predict('anime_image.jpg', 'v0.9') + >>> print(raw_output.keys()) # dict_keys(['prediction', 'embedding', 'logits']) + """ + image = load_image(image, force_background='white', mode='RGB') + model = _open_onnx_model(model_name=model_name) + trans = _open_preprocess(model_name=model_name) + input_ = trans(image)[None, ...] + output_names = [output.name for output in model.get_outputs()] + output_values = model.run(output_names, {'input': input_}) + return {name: value[0] for name, value in zip(output_names, output_values)} + + +def get_pixai_tags(image: ImageTyping, model_name: str = 'v0.9', + thresholds: Union[float, Dict[Any, float]] = None, fmt=FMT_UNSET): + """ + Extract tags from an image using PixAI tagger models. + + This function processes an image through a PixAI tagger model and applies confidence + thresholds to determine which tags to include in the results. The output format can + be customized to return specific categories or all tags together. + + :param image: The input image to analyze (file path, PIL Image, numpy array, etc.) + :type image: ImageTyping + :param model_name: Name or repository ID of the PixAI tagger model to use + :type model_name: str + :param thresholds: Confidence threshold values. Can be a single float applied to all + categories, or a dictionary mapping category IDs/names to specific thresholds + :type thresholds: Union[float, Dict[Any, float]], optional + :param fmt: Output format specification. If FMT_UNSET, returns all available categories. + Can be a tuple of category names to include in output + :type fmt: Any + + :return: Formatted prediction results. Default returns tuple of (general_tags, character_tags, ...) + based on available categories. Can return custom format based on fmt parameter + :rtype: Any + + .. note:: + The fmt parameter can include the following keys: + + - Category names (e.g., 'general', 'character'): dictionaries containing category-specific + tags and their confidence scores + - ``tag``: a dictionary containing all tags across categories and their confidences + - ``embedding``: a 1-dimensional embedding vector of the image, recommended for similarity + search after L2 normalization + - ``logits``: raw 1-dimensional logits output from the model + - ``prediction``: 1-dimensional prediction probabilities from the model + - ``ips_mapping``: a dictionary mapping detected character tags to their associated + IP (Intellectual Property) names in Danbooru-style format + - ``ips_count``: a dictionary containing IP names and their occurrence counts based + on detected characters + - ``ips``: a list of IP names sorted by occurrence count (descending) and name + (ascending), representing the most likely anime/game series in the image + + Default category thresholds are used if not specified. These vary by model and category + but typically range from 0.35 to 0.5. + + You can extract embedding of the given image with the following code + + >>> from imgutils.tagging import get_pixai_tags + >>> + >>> embedding = get_pixai_tags('skadi.jpg', fmt='embedding') + >>> embedding.shape + (1024, ) + + This embedding is valuable for constructing indices that enable rapid querying of images based on + visual features within large-scale datasets. + + Example:: + >>> from imgutils.tagging.pixai import get_pixai_tags + >>> + >>> # Get tags with default format (all categories) + >>> general_tags, character_tags = get_pixai_tags('anime_image.jpg', model_name='v0.9') + >>> print("General tags:", general_tags) + >>> print("Character tags:", character_tags) + >>> + >>> # Get all tags in a single dictionary + >>> all_tags = get_pixai_tags('image.jpg', fmt='tag') + >>> print("All tags:", all_tags) + >>> + >>> # Use custom thresholds + >>> result = get_pixai_tags('image.jpg', thresholds={'general': 0.3, 'character': 0.5}) + >>> + >>> # Get embedding for similarity search + >>> embedding = get_pixai_tags('image.jpg', fmt='embedding') + >>> # Normalize for cosine similarity + >>> import numpy as np + >>> normalized_embedding = embedding / np.linalg.norm(embedding) + >>> + >>> # Get IP information for character identification + >>> ips_mapping = get_pixai_tags('image.jpg', fmt='ips_mapping') + >>> print("Character to IP mapping:", ips_mapping) + >>> # Example output: {'misaka_mikoto': ['toaru_kagaku_no_railgun'], 'hu_tao_(genshin_impact)': ['genshin_impact']} + >>> + >>> # Get most likely anime/game series + >>> top_ips = get_pixai_tags('image.jpg', fmt='ips') + >>> print("Most likely series:", top_ips) + >>> # Example output: ['genshin_impact', 'toaru_kagaku_no_railgun'] + + Here are some images for example + + .. image:: tagging_demo.plot.py.svg + :align: center + + >>> general, character = get_pixai_tags('skadi.jpg') + >>> general + {'patreon_username': 0.9988852739334106, 'baseball_bat': 0.9977256059646606, 'holding_baseball_bat': 0.9858889579772949, 'navel': 0.9830228090286255, 'crop_top': 0.9666315317153931, 'sportswear': 0.9664723873138428, '1girl': 0.9572311639785767, 'long_hair': 0.9550737738609314, 'outdoors': 0.9501817226409912, 'solo': 0.9466996788978577, 'day': 0.9394471049308777, 'breasts': 0.938787579536438, 'web_address': 0.9387772679328918, 'stomach': 0.935083270072937, 'red_eyes': 0.9326196908950806, 'shorts': 0.9305683374404907, 'motion_blur': 0.9278550148010254, 'playing_sports': 0.9263769388198853, 'blue_sky': 0.9213213920593262, 'midriff': 0.9191423654556274, 'large_breasts': 0.9174768924713135, 'artist_name': 0.9089528322219849, 'sky': 0.9054281711578369, 'baseball': 0.904181957244873, 'gloves': 0.9033604860305786, 'thighs': 0.893738865852356, 'black_shorts': 0.8926981687545776, 'volleyball': 0.8198539614677429, 'very_long_hair': 0.7967187166213989, 'short_shorts': 0.7873305082321167, 'black_gloves': 0.7765249013900757, 'white_hair': 0.770541787147522, 'baseball_mitt': 0.7684446573257446, 'thigh_gap': 0.73811936378479, 'sweat': 0.7263807654380798, 'cowboy_shot': 0.7235408425331116, 'short_sleeves': 0.7062878012657166, 'parted_lips': 0.7025120258331299, 'patreon_logo': 0.6970672607421875, 'cloud': 0.6967148780822754, 'looking_at_viewer': 0.6898926496505737, 'holding': 0.6879030466079712, 'swinging': 0.6736525893211365, 'ass_visible_through_thighs': 0.6636734008789062, 'elbow_pads': 0.6630151867866516, 'shirt': 0.6250661611557007, 'hair_between_eyes': 0.6075361967086792, 'standing': 0.5285079479217529, 'black_shirt': 0.5173394680023193, 'linea_alba': 0.513701319694519, 'baseball_uniform': 0.48175835609436035, 'crop_top_overhang': 0.4682246744632721, 'ball': 0.43616628646850586, 'blurry': 0.4201475977897644, 'baseball_stadium': 0.41493287682533264, 'grey_hair': 0.39384859800338745, 'watermark': 0.3919041156768799, 'black_sports_bra': 0.3877854645252228, 'fanbox_username': 0.3790855407714844, 'narrow_waist': 0.36392998695373535} + >>> character + {'skadi_(arknights)': 0.8926791548728943} + >>> + >>> general, character = get_pixai_tags('hutao.jpg') + >>> general + {'bag': 0.9833353757858276, 'backpack': 0.9766197204589844, 'flower-shaped_pupils': 0.962916910648346, 'tongue_out': 0.960152804851532, 'tongue': 0.9526823163032532, 'ghost': 0.9514724016189575, 'plaid_skirt': 0.9499615430831909, '1girl': 0.9378864765167236, 'skirt': 0.9353114366531372, 'bag_charm': 0.9314961433410645, 'symbol-shaped_pupils': 0.9252510070800781, 'charm_(object)': 0.9249529242515564, 'twintails': 0.9239017367362976, 'flower': 0.9175764322280884, 'outdoors': 0.9175151586532593, 'hair_ornament': 0.9161680936813354, 'plaid_clothes': 0.9144806861877441, 'long_hair': 0.8768749833106995, 'pleated_skirt': 0.8597153425216675, 'school_uniform': 0.8573414087295532, 'looking_at_viewer': 0.8392735719680786, ':p': 0.8193913698196411, 'hair_between_eyes': 0.8070638179779053, 'hair_flower': 0.8054562211036682, 'nail_polish': 0.8011300563812256, 'building': 0.7961824536323547, 'jacket': 0.7647742629051208, 'brown_hair': 0.7541910409927368, 'solo': 0.7539198398590088, 'long_sleeves': 0.7471930980682373, 'ahoge': 0.7171378135681152, 'hair_ribbon': 0.6994943618774414, 'red_eyes': 0.6819245219230652, 'bowtie': 0.6639955043792725, 'sidelocks': 0.6275356411933899, 'bush': 0.6164096593856812, 'gate': 0.612525224685669, 'smile': 0.6077383160591125, 'shirt': 0.6042965650558472, 'contemporary': 0.5968752503395081, 'brick_floor': 0.5933602452278137, 'cardigan': 0.5810519456863403, 'gradient_hair': 0.5570307970046997, 'diagonal-striped_bowtie': 0.5565160512924194, 'alternate_costume': 0.5535630583763123, 'school_bag': 0.5535626411437988, 'black_hair': 0.5434530973434448, 'ribbon': 0.5332301259040833, 'hairclip': 0.523446261882782, 'day': 0.5164296627044678, 'street': 0.49987316131591797, 'bow': 0.4941294193267822, 'plum_blossoms': 0.4940766990184784, 'collared_shirt': 0.49013280868530273, 'standing': 0.4820355772972107, 'blue_cardigan': 0.47836002707481384, 'cowboy_shot': 0.4782080054283142, 'pocket': 0.477585107088089, 'pavement': 0.4712265729904175, 'multicolored_hair': 0.4610708951950073, 'blue_jacket': 0.45271334052085876, 'blush': 0.45005902647972107, 'sleeves_past_wrists': 0.440649151802063, 'black_nails': 0.4402858316898346, 'black_bag': 0.4206739366054535, 'miniskirt': 0.4187837243080139, 'red_bow': 0.414681613445282, 'very_long_hair': 0.4129619002342224, 'diagonal-striped_clothes': 0.4112803339958191, 'blazer': 0.40750616788864136, 'striped_bowtie': 0.40123170614242554, 'sunlight': 0.4008329212665558, 'grey_skirt': 0.3930213749408722, 'road': 0.3819067180156708, 'black_ribbon': 0.3776353895664215, 'thighs': 0.3722286820411682, 'hug': 0.37215015292167664, 'brick_wall': 0.3717171549797058, 'white_shirt': 0.3694952428340912, 'open_clothes': 0.36442798376083374, 'open_jacket': 0.3525886535644531, ':d': 0.3343055844306946, 'multicolored_nails': 0.32190075516700745, 'red_bowtie': 0.3157669007778168, 'star-shaped_pupils': 0.309164822101593, 'open_mouth': 0.30890953540802, 'beads': 0.3084579110145569, 'stone_stairs': 0.30559882521629333, 'randoseru': 0.30517613887786865} + >>> character + {'hu_tao_(genshin_impact)': 0.9997367858886719, 'boo_tao_(genshin_impact)': 0.999537467956543} + """ + df_tags, d_ips = _open_tags(model_name=model_name) + values = _raw_predict(image, model_name=model_name) + prediction = values['prediction'] + tags = {} + + default_category_thresholds, category_names = _open_default_category_thresholds(model_name=model_name) + if fmt is FMT_UNSET: + fmt = tuple(category_names[category] for category in sorted(set(df_tags['category'].tolist()))) + + for category in sorted(set(df_tags['category'].tolist())): + mask = df_tags['category'] == category + tag_names = df_tags['name'][mask] + category_pred = prediction[mask] + + if isinstance(thresholds, float): + category_threshold = thresholds + elif isinstance(thresholds, dict) and \ + (category in thresholds or category_names[category] in thresholds): + if category in thresholds: + category_threshold = thresholds[category] + elif category_names[category] in thresholds: + category_threshold = thresholds[category_names[category]] + else: + assert False, 'Should not reach this line' # pragma: no cover + else: + if category in default_category_thresholds: + category_threshold = default_category_thresholds[category] + else: + category_threshold = 0.4 + + mask = category_pred >= category_threshold + tag_names = tag_names[mask].tolist() + category_pred = category_pred[mask].tolist() + cate_tags = dict(sorted(zip(tag_names, category_pred), key=lambda x: (-x[1], x[0]))) + values[category_names[category]] = cate_tags + tags.update(cate_tags) + + values['tag'] = tags + if 'ips' in df_tags.columns: + ips_mapping, ips_counts = {}, defaultdict(lambda: 0) + for tag, _ in tags.items(): + if tag in d_ips: + ips_mapping[tag] = d_ips[tag] + for ip_name in d_ips[tag]: + ips_counts[ip_name] += 1 + values['ips_mapping'] = ips_mapping + values['ips_count'] = dict(ips_counts) + values['ips'] = [x for x, _ in sorted(ips_counts.items(), key=lambda x: (-x[1], x[0]))] + return vreplace(fmt, values) diff --git a/requirements-zoo.txt b/requirements-zoo.txt index 515110506c7..48dd4ff5acd 100644 --- a/requirements-zoo.txt +++ b/requirements-zoo.txt @@ -13,7 +13,7 @@ tensorboard einops thop accelerate -timm~=0.6.13 +timm>=1 ftfy regex git+https://github.com/openai/CLIP.git @@ -27,4 +27,6 @@ tabulate git+https://github.com/deepghs/waifuc.git@main#egg=waifuc pyquery httpx -onnxslim==0.1.32 \ No newline at end of file +onnxslim==0.1.32 +calflops +transformers \ No newline at end of file diff --git a/test/conftest.py b/test/conftest.py index d4f615a88d0..1b93873f6e3 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,5 +1,8 @@ +import os + import pytest from hbutils.testing import TextAligner +from hfutils.cache import delete_cache try: import torch @@ -15,3 +18,12 @@ def _try_import_torch(): @pytest.fixture() def text_aligner(): return TextAligner().multiple_lines() + + +@pytest.fixture(autouse=True, scope='module') +def clean_hf_cache(): + try: + yield + finally: + if os.environ.get('CI'): + delete_cache() diff --git a/test/generic/test_classify.py b/test/generic/test_classify.py index ed789b940f6..06ccad9d0cc 100644 --- a/test/generic/test_classify.py +++ b/test/generic/test_classify.py @@ -1,8 +1,10 @@ +from unittest import skipUnless from unittest.mock import patch import numpy as np import pytest from PIL import Image +from hbutils.testing import vpip from huggingface_hub.utils import reset_sessions from imgutils.generic import classify_predict_score @@ -232,6 +234,7 @@ def test_classify_predict_fmt_complex(self): assert np.linalg.norm(results['embedding']) == pytest.approx(np.linalg.norm(expected_embedding)) @patch("huggingface_hub.constants.HF_HUB_OFFLINE", True) + @skipUnless(vpip('huggingface_hub') < '0.34', 'Has problem on huggingface 0.34+') def test_classify_predict_score_top5_offline_mode(self, clean_session): image = Image.open(get_testfile('png_640.png')) scores = classify_predict_score( diff --git a/test/generic/test_yolo.py b/test/generic/test_yolo.py index d472af2fb16..ba9bf341635 100644 --- a/test/generic/test_yolo.py +++ b/test/generic/test_yolo.py @@ -1,6 +1,8 @@ +from unittest import skipUnless from unittest.mock import patch import pytest +from hbutils.testing import vpip from huggingface_hub import configure_http_backend from huggingface_hub.utils import reset_sessions @@ -52,6 +54,7 @@ def test_detect_faces_none(self): ) == [] @patch("huggingface_hub.constants.HF_HUB_OFFLINE", True) + @skipUnless(vpip('huggingface_hub') < '0.34', 'Has problem on huggingface 0.34+') def test_detect_faces_with_offline_mode(self, clean_session): configure_http_backend() detection = yolo_predict( diff --git a/test/tagging/test_pixai.py b/test/tagging/test_pixai.py new file mode 100644 index 00000000000..afb32d20391 --- /dev/null +++ b/test/tagging/test_pixai.py @@ -0,0 +1,173 @@ +import numpy as np +import pytest + +from imgutils.tagging import get_pixai_tags +from imgutils.tagging.pixai import _open_tags, _open_preprocess, _open_onnx_model, _open_default_category_thresholds +from test.testings import get_testfile + + +@pytest.fixture(autouse=True, scope='module') +def _release_model_after_run(): + try: + yield + finally: + _open_tags.cache_clear() + _open_preprocess.cache_clear() + _open_onnx_model.cache_clear() + _open_default_category_thresholds.cache_clear() + + +@pytest.mark.unittest +class TestTaggingPixAI: + def test_get_pixai_tags(self): + tags, chars = get_pixai_tags(get_testfile('6124220.jpg')) + + assert tags['cat_girl'] >= 0.8 + assert not chars + assert isinstance(tags['cat_girl'], float) + + tags, chars = get_pixai_tags(get_testfile('6125785.jpg')) + assert tags['1girl'] >= 0.90 + assert chars['hu_tao_(genshin_impact)'] >= 0.95 + assert isinstance(tags['1girl'], float) + assert isinstance(chars['hu_tao_(genshin_impact)'], float) + + def test_pixai_tags_sample(self): + tags, chars = get_pixai_tags(get_testfile('6125785.jpg')) + + assert tags == pytest.approx({ + 'symbol-shaped_pupils': 0.9938011169433594, 'ghost': 0.9909012317657471, + 'flower-shaped_pupils': 0.9810856580734253, 'porkpie_hat': 0.9693692922592163, + 'black_nails': 0.9692082405090332, 'ghost_pose': 0.9579154253005981, 'ring': 0.9509532451629639, + '1girl': 0.9435760974884033, 'hat_flower': 0.9315004348754883, 'hat': 0.9308193922042847, + 'flower': 0.9203469753265381, 'plum_blossoms': 0.8814681768417358, 'red_shirt': 0.8734415769577026, + 'jewelry': 0.8436187505722046, 'claw_pose': 0.8382970094680786, 'chinese_clothes': 0.8213527202606201, + 'long_sleeves': 0.8101956844329834, 'long_hair': 0.7956619262695312, 'twintails': 0.7865841388702393, + 'looking_at_viewer': 0.7795377373695374, 'hat_tassel': 0.7775378227233887, + 'multiple_rings': 0.7483937740325928, 'signature': 0.7478364109992981, 'smile': 0.741124153137207, + 'open_mouth': 0.7302751541137695, 'red_flower': 0.705192506313324, 'solo': 0.6761863827705383, + 'hat_ornament': 0.6489882469177246, 'red_eyes': 0.6290297508239746, 'black_hat': 0.6113986372947693, + 'brown_coat': 0.6067467927932739, 'nail_polish': 0.5937596559524536, 'coat': 0.5664999485015869, + 'upper_body': 0.5661808848381042, 'brown_hair': 0.5366549491882324, 'thumb_ring': 0.5287110209465027, + 'star_(symbol)': 0.5028221011161804, 'shirt': 0.4754156470298767, 'star-shaped_pupils': 0.45203927159309387, + 'hair_between_eyes': 0.4512987434864044, 'orange_eyes': 0.4298587739467621, ':d': 0.4219313859939575, + 'tangzhuang': 0.3894977569580078, ':3': 0.3866003453731537, 'blush': 0.3761531412601471, + 'v-shaped_eyebrows': 0.34583818912506104, 'sidelocks': 0.3450319468975067, + 'brown_shirt': 0.31902527809143066, 'blurry': 0.3167526423931122 + }, abs=2e-2) + assert chars == pytest.approx({ + 'hu_tao_(genshin_impact)': 0.9999773502349854, 'boo_tao_(genshin_impact)': 0.9996482133865356 + }, abs=2e-2) + + def test_pixai_rgba(self): + tags, chars = get_pixai_tags(get_testfile('nian.png')) + assert tags == pytest.approx({ + 'red_tube_top': 0.9998065233230591, 'tube_top': 0.9994766712188721, 'red_bandeau': 0.9992805123329163, + 'white_shorts': 0.9976104497909546, 'bandeau': 0.9968563914299011, + 'transparent_background': 0.9929141402244568, 'horns': 0.9903750419616699, 'full_body': 0.9860043525695801, + 'midriff': 0.9834859371185303, 'shorts': 0.9828792810440063, '1girl': 0.982780933380127, + 'strapless': 0.9817174673080444, 'navel': 0.9786539673805237, 'white_footwear': 0.9764545559883118, + 'white_hair': 0.9713669419288635, 'solo': 0.9699287414550781, 'tail': 0.9654039740562439, + 'stomach': 0.9624348282814026, 'standing': 0.9548828601837158, 'pointy_ears': 0.9536501169204712, + 'open_coat': 0.9320983290672302, 'looking_at_viewer': 0.9311813116073608, + 'open_clothes': 0.9311613440513611, 'long_hair': 0.9210299253463745, 'hand_on_own_hip': 0.9110379219055176, + 'holding': 0.9101808071136475, 'short_shorts': 0.9012227058410645, 'tongue': 0.8910117745399475, + 'tongue_out': 0.8871811628341675, 'dragon_horns': 0.8710488080978394, 'weapon': 0.8667179346084595, + 'coat': 0.8642098307609558, 'belt': 0.8602883219718933, 'tachi-e': 0.8591709136962891, + 'crop_top': 0.856552004814148, 'white_coat': 0.8258588314056396, 'dragon_tail': 0.822641909122467, + 'purple_eyes': 0.7998863458633423, 'streaked_hair': 0.796852707862854, 'half_updo': 0.790149450302124, + 'thighs': 0.7805066108703613, 'bare_legs': 0.7744182348251343, ':d': 0.7707932591438293, + 'wide_sleeves': 0.767501711845398, 'smile': 0.765789270401001, 'multicolored_hair': 0.7620936632156372, + 'bead_bracelet': 0.7511173486709595, 'red_hair': 0.7497332692146301, 'drop_shadow': 0.7467425465583801, + 'long_sleeves': 0.7251672744750977, 'red_skin': 0.7151302099227905, 'sidelocks': 0.6788227558135986, + 'socks': 0.6740190982818604, ':p': 0.6698598861694336, 'sword': 0.6676492691040039, + 'jewelry': 0.6632057428359985, 'hand_up': 0.637663722038269, 'boots': 0.6372023224830627, + 'black_socks': 0.626619815826416, 'pixel_art': 0.6166024208068848, 'black_belt': 0.6042207479476929, + 'holding_weapon': 0.6014618873596191, 'breasts': 0.6006667017936707, 'ankle_boots': 0.5968820452690125, + 'white_jacket': 0.5880962014198303, 'jacket': 0.5767049789428711, 'tassel_earrings': 0.5749104022979736, + 'open_mouth': 0.5724694132804871, 'beads': 0.5716468095779419, 'braid': 0.5678756237030029, + 'open_jacket': 0.5464628338813782, 'bare_shoulders': 0.5238108038902283, 'shoes': 0.506024181842804, + 'bracelet': 0.4898264706134796, 'pouch': 0.48912930488586426, 'colored_skin': 0.46792706847190857, + 'originium_arts_(arknights)': 0.4517717957496643, 'dragon_girl': 0.4428599774837494, + 'thigh_strap': 0.4376051723957062, 'earrings': 0.4229695200920105, 'collarbone': 0.41589391231536865, + 'red_eyes': 0.4031139016151428, 'body_markings': 0.4004215598106384, 'medium_breasts': 0.384432315826416, + 'holding_sword': 0.3541863262653351, 'white_pants': 0.34769535064697266, + 'club_(weapon)': 0.33692681789398193, 'eyeshadow': 0.3357437252998352, 'ponytail': 0.3333626389503479, + 'tassel': 0.33294352889060974, 'multiple_weapons': 0.3174567222595215, 'pink_eyes': 0.31579816341400146, + 'flame-tipped_tail': 0.3054245710372925 + }, abs=2e-2) + assert chars == pytest.approx({'nian_(arknights)': 0.9999774694442749}, abs=2e-2) + + def test_pixai_tags_sample_ips(self): + tags, chars, ips, ips_count, ips_mapping = get_pixai_tags( + get_testfile('6125785.jpg'), + fmt=('general', 'character', 'ips', 'ips_count', 'ips_mapping') + ) + + assert tags == pytest.approx({ + 'symbol-shaped_pupils': 0.9938011169433594, 'ghost': 0.9909012317657471, + 'flower-shaped_pupils': 0.9810856580734253, 'porkpie_hat': 0.9693692922592163, + 'black_nails': 0.9692082405090332, 'ghost_pose': 0.9579154253005981, 'ring': 0.9509532451629639, + '1girl': 0.9435760974884033, 'hat_flower': 0.9315004348754883, 'hat': 0.9308193922042847, + 'flower': 0.9203469753265381, 'plum_blossoms': 0.8814681768417358, 'red_shirt': 0.8734415769577026, + 'jewelry': 0.8436187505722046, 'claw_pose': 0.8382970094680786, 'chinese_clothes': 0.8213527202606201, + 'long_sleeves': 0.8101956844329834, 'long_hair': 0.7956619262695312, 'twintails': 0.7865841388702393, + 'looking_at_viewer': 0.7795377373695374, 'hat_tassel': 0.7775378227233887, + 'multiple_rings': 0.7483937740325928, 'signature': 0.7478364109992981, 'smile': 0.741124153137207, + 'open_mouth': 0.7302751541137695, 'red_flower': 0.705192506313324, 'solo': 0.6761863827705383, + 'hat_ornament': 0.6489882469177246, 'red_eyes': 0.6290297508239746, 'black_hat': 0.6113986372947693, + 'brown_coat': 0.6067467927932739, 'nail_polish': 0.5937596559524536, 'coat': 0.5664999485015869, + 'upper_body': 0.5661808848381042, 'brown_hair': 0.5366549491882324, 'thumb_ring': 0.5287110209465027, + 'star_(symbol)': 0.5028221011161804, 'shirt': 0.4754156470298767, 'star-shaped_pupils': 0.45203927159309387, + 'hair_between_eyes': 0.4512987434864044, 'orange_eyes': 0.4298587739467621, ':d': 0.4219313859939575, + 'tangzhuang': 0.3894977569580078, ':3': 0.3866003453731537, 'blush': 0.3761531412601471, + 'v-shaped_eyebrows': 0.34583818912506104, 'sidelocks': 0.3450319468975067, + 'brown_shirt': 0.31902527809143066, 'blurry': 0.3167526423931122 + }, abs=2e-2) + assert chars == pytest.approx({ + 'hu_tao_(genshin_impact)': 0.9999773502349854, 'boo_tao_(genshin_impact)': 0.9996482133865356 + }, abs=2e-2) + assert ips == ['genshin_impact'] + assert ips_count == {'genshin_impact': 2} + assert ips_mapping == { + 'hu_tao_(genshin_impact)': ['genshin_impact'], + 'boo_tao_(genshin_impact)': ['genshin_impact'], + } + + def test_pixai_tags_sample_custom(self): + tags, chars, embedding, logits, prediction = get_pixai_tags( + get_testfile('6125785.jpg'), + fmt=('general', 'character', 'embedding', 'logits', 'prediction'), + model_name='deepghs/pixai-tagger-v0.9-onnx', + ) + + assert tags == pytest.approx({ + 'symbol-shaped_pupils': 0.9938011169433594, 'ghost': 0.9909012317657471, + 'flower-shaped_pupils': 0.9810856580734253, 'porkpie_hat': 0.9693692922592163, + 'black_nails': 0.9692082405090332, 'ghost_pose': 0.9579154253005981, 'ring': 0.9509532451629639, + '1girl': 0.9435760974884033, 'hat_flower': 0.9315004348754883, 'hat': 0.9308193922042847, + 'flower': 0.9203469753265381, 'plum_blossoms': 0.8814681768417358, 'red_shirt': 0.8734415769577026, + 'jewelry': 0.8436187505722046, 'claw_pose': 0.8382970094680786, 'chinese_clothes': 0.8213527202606201, + 'long_sleeves': 0.8101956844329834, 'long_hair': 0.7956619262695312, 'twintails': 0.7865841388702393, + 'looking_at_viewer': 0.7795377373695374, 'hat_tassel': 0.7775378227233887, + 'multiple_rings': 0.7483937740325928, 'signature': 0.7478364109992981, 'smile': 0.741124153137207, + 'open_mouth': 0.7302751541137695, 'red_flower': 0.705192506313324, 'solo': 0.6761863827705383, + 'hat_ornament': 0.6489882469177246, 'red_eyes': 0.6290297508239746, 'black_hat': 0.6113986372947693, + 'brown_coat': 0.6067467927932739, 'nail_polish': 0.5937596559524536, 'coat': 0.5664999485015869, + 'upper_body': 0.5661808848381042, 'brown_hair': 0.5366549491882324, 'thumb_ring': 0.5287110209465027, + 'star_(symbol)': 0.5028221011161804, 'shirt': 0.4754156470298767, 'star-shaped_pupils': 0.45203927159309387, + 'hair_between_eyes': 0.4512987434864044, 'orange_eyes': 0.4298587739467621, ':d': 0.4219313859939575, + 'tangzhuang': 0.3894977569580078, ':3': 0.3866003453731537, 'blush': 0.3761531412601471, + 'v-shaped_eyebrows': 0.34583818912506104, 'sidelocks': 0.3450319468975067, + 'brown_shirt': 0.31902527809143066, 'blurry': 0.3167526423931122 + }, abs=2e-2) + assert chars == pytest.approx({ + 'hu_tao_(genshin_impact)': 0.9999773502349854, 'boo_tao_(genshin_impact)': 0.9996482133865356 + }, abs=2e-2) + + assert isinstance(embedding, np.ndarray) + assert embedding.shape == (1024,) + assert isinstance(logits, np.ndarray) + assert logits.shape == (13461,) + assert isinstance(prediction, np.ndarray) + assert prediction.shape == (13461,) diff --git a/zoo/pixai_tagger/__init__.py b/zoo/pixai_tagger/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/zoo/pixai_tagger/export.py b/zoo/pixai_tagger/export.py new file mode 100644 index 00000000000..e3be359a608 --- /dev/null +++ b/zoo/pixai_tagger/export.py @@ -0,0 +1,291 @@ +import json +import os.path +import re +from pprint import pformat +from textwrap import indent + +import numpy as np +import onnx +import onnxruntime +import pandas as pd +import torch +from ditk import logging +from hbutils.string import titleize, underscore +from hbutils.system import TemporaryDirectory +from hbutils.testing import vpip +from hfutils.operate import get_hf_client, upload_directory_as_directory +from hfutils.repository import hf_hub_repo_url, hf_hub_repo_file_url +from hfutils.utils import hf_normpath +from huggingface_hub import hf_hub_url +from thop import clever_format + +from imgutils.data import load_image +from imgutils.preprocess import parse_torchvision_transforms +from imgutils.tagging import get_pixai_tags +from zoo.pixai_tagger.tags import load_tags +from zoo.utils import onnx_optimize, get_testfile, torch_model_profile_via_calflops +from .min_script import EndpointHandler +from .onnx import get_model + + +def sync(src_repo: str, dst_repo: str, no_optimize: bool = False, show_current_results: bool = False): + hf_client = get_hf_client() + if not hf_client.repo_exists(repo_id=dst_repo, repo_type='model'): + hf_client.create_repo(repo_id=dst_repo, repo_type='model', private=True) + + handler = EndpointHandler(repo_id=src_repo) + meta_info = {} + meta_info['repo_id'] = src_repo + with TemporaryDirectory() as upload_dir: + preprocessor = handler.transform + preprocessor_file = os.path.join(upload_dir, 'preprocess.json') + logging.info(f'Dumping preprocessor:\n{preprocessor}\nto file {preprocessor_file!r}.') + with open(preprocessor_file, 'w') as f: + json.dump({ + 'stages': parse_torchvision_transforms(handler.transform), + }, f, sort_keys=True, ensure_ascii=False, indent=4) + + logging.info('Scanning tags ...') + categories = np.zeros((len(handler.index_to_tag_map),), dtype=np.int32) + idx = np.array(range(len(categories))) + categories[idx < handler.gen_tag_count] = 0 + categories[idx >= handler.gen_tag_count] = 4 + df_src_tags = pd.DataFrame({ + 'name': [v for _, v in sorted(handler.index_to_tag_map.items())], + 'category': categories, + }) + df_tags = load_tags(df_src_tags) + exts = [] + for titem in df_tags.to_dict('records'): + if titem['category'] == 4: + if titem['name'] in handler.character_ip_mapping: + exts.append(handler.character_ip_mapping[titem['name']]) + else: + exts.append([]) + else: + exts.append([]) + df_tags['ips'] = list(map(json.dumps, exts)) + logging.info(f'Tags:\n{df_tags}') + df_tags.to_csv(os.path.join(upload_dir, 'selected_tags.csv'), index=False) + + d_category_names = { + 0: 'general', + 4: 'character', + } + d_category_thresholds = { + 0: handler.default_general_threshold, + 4: handler.default_character_threshold, + } + with open(os.path.join(upload_dir, 'categories.json'), 'w') as f: + json.dump([ + { + "category": cate_id, + "name": cate_name, + } for cate_id, cate_name in sorted(d_category_names.items()) + ], f, sort_keys=True, ensure_ascii=False, indent=4) + df_th = pd.DataFrame([ + { + "category": cate_id, + "name": cate_name, + 'threshold': d_category_thresholds[cate_id] + } for cate_id, cate_name in sorted(d_category_names.items()) + ]) + df_th.to_csv(os.path.join(upload_dir, 'thresholds.csv'), index=False) + + dummy_image = load_image(get_testfile('6125785.jpg'), mode='RGB', force_background='white') + dummy_input = handler.transform(dummy_image).unsqueeze(0).to(handler.device) + meta_info['input_size'] = dummy_input.shape[-1] + flops, params, macs = torch_model_profile_via_calflops(model=handler.model, input_=dummy_input) + meta_info['flops'] = flops + meta_info['params'] = params + meta_info['macs'] = macs + + wrapped_model, (conv_features, conv_logits, conv_preds) = get_model(handler.model, dummy_input) + conv_features = conv_features.detach().cpu() + conv_logits = conv_logits.detach().cpu() + conv_preds = conv_preds.detach().cpu() + meta_info['num_features'] = conv_features.shape[-1] + meta_info['num_classes'] = conv_preds.shape[-1] + new_meta_file = os.path.join(upload_dir, 'meta.json') + logging.info(f'Saving metadata to {new_meta_file!r} ...') + with open(new_meta_file, 'w') as f: + json.dump(meta_info, f, indent=4, sort_keys=True, ensure_ascii=False) + + onnx_filename = os.path.join(upload_dir, 'model.onnx') + with TemporaryDirectory() as td: + temp_model_onnx = os.path.join(td, 'model.onnx') + logging.info(f'Exporting temporary ONNX model to {temp_model_onnx!r} ...') + torch.onnx.export( + wrapped_model, + dummy_input, + temp_model_onnx, + input_names=['input'], + output_names=['embedding', 'logits', 'prediction'], + dynamic_axes={ + 'input': {0: 'batch_size'}, + 'embedding': {0: 'batch_size'}, + 'logits': {0: 'batch_size'}, + 'prediction': {0: 'batch_size'}, + }, + opset_version=14, + do_constant_folding=True, + export_params=True, + verbose=False, + custom_opsets=None, + ) + + model = onnx.load(temp_model_onnx) + if not no_optimize: + logging.info('Optimizing onnx model ...') + model = onnx_optimize(model) + + output_model_dir, _ = os.path.split(onnx_filename) + if output_model_dir: + os.makedirs(output_model_dir, exist_ok=True) + logging.info(f'Complete model saving to {onnx_filename!r} ...') + onnx.save(model, onnx_filename) + + session = onnxruntime.InferenceSession(onnx_filename) + o_embeddings, = session.run(['embedding'], {'input': dummy_input.detach().cpu().numpy()}) + emb_1 = o_embeddings / np.linalg.norm(o_embeddings, axis=-1, keepdims=True) + emb_2 = conv_features.numpy() / np.linalg.norm(conv_features.numpy(), axis=-1, keepdims=True) + emb_sims = (emb_1 * emb_2).sum() + logging.info(f'Similarity of the embeddings is {emb_sims:.5f}.') + assert emb_sims >= 0.98, f'Similarity of the embeddings is {emb_sims:.5f}, ONNX validation failed.' + + with open(os.path.join(upload_dir, 'README.md'), 'w') as f: + print('---', file=f) + print('pipeline_tag: image-classification', file=f) + print('base_model:', file=f) + print(f'- {src_repo}', file=f) + print('language:', file=f) + print('- en', file=f) + print('tags:', file=f) + print('- image', file=f) + print('- dghs-imgutils', file=f) + print('library_name: dghs-imgutils', file=f) + print('license: mit', file=f) + print('---', file=f) + print('', file=f) + + print(f'# ONNX Version for {src_repo}', file=f) + print(f'', file=f) + + print(f'This is the ONNX-exported version of PixAI\'s tagger ' + f'[{src_repo}]({hf_hub_repo_url(repo_id=src_repo, repo_type="model")}).', file=f) + print(f'', file=f) + + s_flops, s_params, s_macs = clever_format([flops, params, macs], "%.1f") + print(f'## Model Details', file=f) + print(f'', file=f) + print(f'- **Model Type:** Multilabel Image classification / feature backbone', file=f) + print(f'- **Model Stats:**', file=f) + print(f' - Params: {s_params}', file=f) + print(f' - FLOPs / MACs: {s_flops} / {s_macs}', file=f) + print(f' - Image size: {dummy_input.shape[-1]} x {dummy_input.shape[-2]}', file=f) + print(f' - Tags Count: {len(df_tags)}', file=f) + for category in sorted(set(df_tags['category'])): + print(f' - {titleize(d_category_names[category])} (#{category}) Tags Count: ' + f'{len(df_tags[df_tags["category"] == category])}', file=f) + print(f'', file=f) + + print(f'## Thresholds', file=f) + print(f'', file=f) + ths = [] + for item in df_th.to_dict('records'): + ths.append({ + 'Category': item['category'], + 'Name': item['name'], + 'Count': len(df_tags[df_tags['category'] == item['category']]), + 'Threshold': item['threshold'], + }) + df_th_shown = pd.DataFrame(ths) + print(df_th_shown.to_markdown(index=False), file=f) + print(f'', file=f) + + print(f'## How to Use', file=f) + print(f'', file=f) + imgutils_version = str(vpip('dghs-imgutils')._actual_version) + sample_input = dummy_image + if min(sample_input.width, sample_input.height) > 640: + r = min(sample_input.width, sample_input.height) / 640 + new_width = int(sample_input.width / r) + new_height = int(sample_input.height / r) + sample_input = sample_input.resize((new_width, new_height)) + sample_input_file = os.path.join(upload_dir, 'sample.webp') + sample_input_relfile = hf_normpath(os.path.relpath(sample_input_file, upload_dir)) + sample_input.save(sample_input_file) + sample_input_url = hf_hub_url(repo_id=dst_repo, repo_type='model', filename=sample_input_relfile) + sample_input_page_url = hf_hub_repo_file_url(repo_id=dst_repo, repo_type='model', path=sample_input_relfile) + + print(f'We provided a sample image for our code samples, ' + f'you can find it [here]({sample_input_page_url}).', file=f) + print(f'', file=f) + + print(f'Install [dghs-imgutils](https://github.com/deepghs/imgutils) with the following command', file=f) + print(f'', file=f) + print(f'```shell', file=f) + print(f'pip install \'dghs-imgutils>={imgutils_version}\' torch huggingface_hub timm pillow pandas', file=f) + print(f'```', file=f) + print(f'', file=f) + + print(f'You can use function `get_pixai_tags` to tag your image. ' + f'For more details of this function, see [documentation of imgutils.tagging.pixai]' + f'(https://dghs-imgutils.deepghs.org/main/api_doc/tagging/pixai.html).', file=f) + print(f'', file=f) + + matching = re.fullmatch(r'^deepghs/pixai-tagger-(?P[\s\S]+)-onnx$', dst_repo) + model_name = matching.group('version') if matching else dst_repo + print(f'```python', file=f) + print(f'from imgutils.tagging import get_pixai_tags', file=f) + print(f'', file=f) + + cate_names = tuple((cate_name for _, cate_name in sorted(d_category_names.items()))) + fmt_names = tuple([*cate_names, 'ips', 'ips_mapping']) + var_names = tuple(map(underscore, fmt_names)) + print(f'{", ".join(var_names)} = get_pixai_tags(', file=f) + print(f' {sample_input_url!r},', file=f) + print(f' model_name={model_name!r},', file=f) + print(f' fmt={fmt_names!r},', file=f) + print(f')', file=f) + print(f'', file=f) + + var_values = get_pixai_tags(sample_input_file, model_name=model_name, fmt=fmt_names) + for varname, fmtname, varvalue in zip(var_names, fmt_names, var_values): + print(f'print({varname})', file=f) + print(indent( + pformat(varvalue, sort_dicts=False), + prefix='# ', + ), file=f) + + print(f'', file=f) + print(f'```', file=f) + print(f'', file=f) + + upload_directory_as_directory( + repo_id=dst_repo, + repo_type='model', + local_directory=upload_dir, + path_in_repo='.', + message=f'Upload ONNX export of model {src_repo!r}', + clear=True, + ) + + # print(df_tags) + + # pprint(handler.index_to_tag_map) + # pprint(handler.gen_tag_count) + # pprint(handler.character_tag_count) + # pprint(handler.character_ip_mapping) + + +if __name__ == '__main__': + logging.try_init_root(logging.INFO) + # matching = re.fullmatch(r'^deepghs/pixai-tagger-(?P[\s\S]+)-onnx$', 'deepghs/pixai-tagger-v0.9-onnx') + # print(matching) + # print(matching.group('version')) + sync( + src_repo='pixai-labs/pixai-tagger-v0.9', + dst_repo='deepghs/pixai-tagger-v0.9-onnx', + show_current_results=True, + ) diff --git a/zoo/pixai_tagger/min_script.py b/zoo/pixai_tagger/min_script.py new file mode 100644 index 00000000000..1c06ace1fcf --- /dev/null +++ b/zoo/pixai_tagger/min_script.py @@ -0,0 +1,229 @@ +import base64 +import io +import json +import logging +import time +from pathlib import Path +from typing import Any + +import requests +import timm +import torch +import torchvision.transforms as transforms +from PIL import Image +from huggingface_hub import hf_hub_download + +from imgutils.preprocess import parse_torchvision_transforms + + +class TaggingHead(torch.nn.Module): + def __init__(self, input_dim, num_classes): + super().__init__() + self.input_dim = input_dim + self.num_classes = num_classes + self.head = torch.nn.Sequential(torch.nn.Linear(input_dim, num_classes)) + self.sigmoid = torch.nn.Sigmoid() + + def forward(self, x): + logits = self.head(x) + probs = self.sigmoid(logits) + return probs + + +def get_tags(tags_file: Path) -> tuple[dict[str, int], int, int]: + with tags_file.open("r", encoding="utf-8") as f: + tag_info = json.load(f) + tag_map = tag_info["tag_map"] + tag_split = tag_info["tag_split"] + gen_tag_count = tag_split["gen_tag_count"] + character_tag_count = tag_split["character_tag_count"] + return tag_map, gen_tag_count, character_tag_count + + +def get_character_ip_mapping(mapping_file: Path): + with mapping_file.open("r", encoding="utf-8") as f: + mapping = json.load(f) + return mapping + + +def get_encoder(): + base_model_repo = "hf_hub:SmilingWolf/wd-eva02-large-tagger-v3" + encoder = timm.create_model(base_model_repo, pretrained=False) + encoder.reset_classifier(0) + return encoder + + +def get_decoder(): + decoder = TaggingHead(1024, 13461) + return decoder + + +def get_model(): + encoder = get_encoder() + decoder = get_decoder() + model = torch.nn.Sequential(encoder, decoder) + return model + + +def load_model(weights_file, device): + model = get_model() + states_dict = torch.load(weights_file, map_location=device, weights_only=True) + model.load_state_dict(states_dict) + model.to(device) + model.eval() + return model + + +def pure_pil_alpha_to_color_v2( + image: Image.Image, color: tuple[int, int, int] = (255, 255, 255) +) -> Image.Image: + """ + Convert a PIL image with an alpha channel to a RGB image. + This is a workaround for the fact that the model expects a RGB image, but the image may have an alpha channel. + This function will convert the image to a RGB image, and fill the alpha channel with the given color. + The alpha channel is the 4th channel of the image. + """ + image.load() # needed for split() + background = Image.new("RGB", image.size, color) + background.paste(image, mask=image.split()[3]) # 3 is the alpha channel + return background + + +def pil_to_rgb(image: Image.Image) -> Image.Image: + if image.mode == "RGBA": + image = pure_pil_alpha_to_color_v2(image) + elif image.mode == "P": + image = pure_pil_alpha_to_color_v2(image.convert("RGBA")) + else: + image = image.convert("RGB") + return image + + +class EndpointHandler: + def __init__(self, repo_id: str = 'pixai-labs/pixai-tagger-v0.9'): + weights_file = Path(hf_hub_download( + repo_id=repo_id, + repo_type='model', + filename="model_v0.9.pth", + )) + tags_file = Path(hf_hub_download( + repo_id=repo_id, + repo_type='model', + filename="tags_v0.9_13k.json", + )) + mapping_file = Path(hf_hub_download( + repo_id=repo_id, + repo_type='model', + filename="char_ip_map.json", + )) + + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.model = load_model(str(weights_file), self.device) + self.transform = transforms.Compose( + [ + transforms.Resize((448, 448)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), + ] + ) + self.fetch_image_timeout = 5.0 + self.default_general_threshold = 0.3 + self.default_character_threshold = 0.85 + + tag_map, self.gen_tag_count, self.character_tag_count = get_tags(tags_file) + + # Invert the tag_map for efficient index-to-tag lookups + self.index_to_tag_map = {v: k for k, v in tag_map.items()} + + self.character_ip_mapping = get_character_ip_mapping(mapping_file) + + def __call__(self, data: dict[str, Any]) -> dict[str, Any]: + inputs = data.pop("inputs", data) + + fetch_start_time = time.time() + if isinstance(inputs, Image.Image): + image = inputs + elif image_url := inputs.pop("url", None): + with requests.get( + image_url, stream=True, timeout=self.fetch_image_timeout + ) as res: + res.raise_for_status() + image = Image.open(res.raw) + elif image_base64_encoded := inputs.pop("image", None): + image = Image.open(io.BytesIO(base64.b64decode(image_base64_encoded))) + else: + raise ValueError(f"No image or url provided: {data}") + # remove alpha channel if it exists + image = pil_to_rgb(image) + fetch_time = time.time() - fetch_start_time + + parameters = data.pop("parameters", {}) + general_threshold = parameters.pop( + "general_threshold", self.default_general_threshold + ) + character_threshold = parameters.pop( + "character_threshold", self.default_character_threshold + ) + + inference_start_time = time.time() + with torch.inference_mode(): + # Preprocess image on CPU, then pin memory for faster async transfer + image_tensor = self.transform(image).unsqueeze(0).pin_memory() + + # Asynchronously move image to GPU + image_tensor = image_tensor.to(self.device, non_blocking=True) + + # Run model on GPU + probs = self.model(image_tensor)[0] # Get probs for the single image + + # Perform thresholding directly on the GPU + general_mask = probs[: self.gen_tag_count] > general_threshold + character_mask = probs[self.gen_tag_count:] > character_threshold + + # Get the indices of positive tags on the GPU + general_indices = general_mask.nonzero(as_tuple=True)[0] + character_indices = ( + character_mask.nonzero(as_tuple=True)[0] + self.gen_tag_count + ) + + # Combine indices and move the small result tensor to the CPU + combined_indices = torch.cat((general_indices, character_indices)).cpu() + + inference_time = time.time() - inference_start_time + + post_process_start_time = time.time() + + cur_gen_tags = [] + cur_char_tags = [] + + # Use the efficient pre-computed map for lookups + for i in combined_indices: + idx = i.item() + tag = self.index_to_tag_map[idx] + if idx < self.gen_tag_count: + cur_gen_tags.append(tag) + else: + cur_char_tags.append(tag) + + ip_tags = [] + for tag in cur_char_tags: + if tag in self.character_ip_mapping: + ip_tags.extend(self.character_ip_mapping[tag]) + ip_tags = sorted(set(ip_tags)) + post_process_time = time.time() - post_process_start_time + + logging.info( + f"Timing - Fetch: {fetch_time:.3f}s, Inference: {inference_time:.3f}s, Post-process: {post_process_time:.3f}s, Total: {fetch_time + inference_time + post_process_time:.3f}s" + ) + + return { + "feature": cur_gen_tags, + "character": cur_char_tags, + "ip": ip_tags, + } + + +if __name__ == '__main__': + handler = EndpointHandler() + print(handler.transform) + print(parse_torchvision_transforms(handler.transform)) diff --git a/zoo/pixai_tagger/onnx.py b/zoo/pixai_tagger/onnx.py new file mode 100644 index 00000000000..1644322ac8f --- /dev/null +++ b/zoo/pixai_tagger/onnx.py @@ -0,0 +1,61 @@ +import logging + +import torch +from torch import nn + + +class ModuleWrapper(nn.Module): + def __init__(self, base_module: nn.Module, classifier: nn.Module, sigmoid: nn.Module): + super().__init__() + self.base_module = base_module + self.classifier = classifier + self.sigmoid = sigmoid + + self._output_features = None + self._output_logits = None + self._register_hook() + + def _register_hook(self): + def hook_fn_embeddings(module, input_tensor, output_tensor): + assert isinstance(input_tensor, tuple) and len(input_tensor) == 1 + input_tensor = input_tensor[0] + self._output_features = input_tensor + + self.classifier.register_forward_hook(hook_fn_embeddings) + + def hook_fn_logits(module, input_tensor, output_tensor): + assert isinstance(input_tensor, tuple) and len(input_tensor) == 1 + input_tensor = input_tensor[0] + self._output_logits = input_tensor + + self.sigmoid.register_forward_hook(hook_fn_logits) + + def forward(self, x: torch.Tensor): + preds = self.base_module(x) + + if self._output_features is None: + raise RuntimeError("Target module did not receive any input during forward pass (features)") + if self._output_logits is None: + raise RuntimeError("Target module did not receive any input during forward pass (logits)") + features, self._output_features = self._output_features, None + logits, self._output_logits = self._output_logits, None + assert all([x == 1 for x in features.shape[2:]]), f'Invalid feature shape: {features.shape!r}' + features = torch.flatten(features, start_dim=1) + + return features, logits, preds + + +def get_model(model: nn.Module, dummy_input: torch.Tensor): + assert isinstance(model, nn.Sequential) + head = model[-1] + wrapped_model = ModuleWrapper(model, head, head.sigmoid) + + logging.info(f'Input size: {dummy_input.shape!r}') + with torch.no_grad(): + dummy_embedding, dummy_logits, dummy_preds = wrapped_model(dummy_input) + logging.info(f'Embedding size: {dummy_embedding.shape!r}') + logging.info(f'Logits size: {dummy_preds.shape!r}') + logging.info(f'Preds size: {dummy_preds.shape!r}') + + return wrapped_model, (dummy_embedding, dummy_logits, dummy_preds) + # print(model[-1]) diff --git a/zoo/pixai_tagger/tags.py b/zoo/pixai_tagger/tags.py new file mode 100644 index 00000000000..3eabe7c1c7a --- /dev/null +++ b/zoo/pixai_tagger/tags.py @@ -0,0 +1,82 @@ +from functools import lru_cache + +import pandas as pd +from ditk import logging +from hfutils.operate import get_hf_client +from tqdm import tqdm +from waifuc.utils import srequest + +from zoo.wd14.tags import _get_tag_by_name, _db_session + +_CATEGORY_MAPS = { + 'general': 0, + 'character': 4, +} + + +@lru_cache() +def _get_rating_count_by_name(tag_name: str): + session = _db_session() + logging.info(f'Getting count for {tag_name!r} ...') + vs = srequest( + session, 'GET', f'https://danbooru.donmai.us/counts/posts.json', + params={'tags': f'rating:{tag_name}'} + ).json() + logging.info(f'Result of {tag_name!r}: {vs!r}') + return vs['counts']['posts'] + + +def load_tags(df_src_tags: pd.DataFrame): + hf_client = get_hf_client() + df_p_tags = pd.read_csv(hf_client.hf_hub_download( + repo_id='deepghs/site_tags', + repo_type='dataset', + filename='danbooru.donmai.us/tags.csv' + )) + logging.info(f'Loaded danbooru tags pool, columns: {df_p_tags.columns!r}') + d_p_tags = {(item['category'], item['name']): item for item in df_p_tags.to_dict('records')} + + rows = [] + src_tags = df_src_tags.to_dict('records') + for i in tqdm(range(len(src_tags)), desc='Scan Tags'): + tag_name = src_tags[i]['name'] + category = src_tags[i]['category'] + if (category, tag_name) in d_p_tags: + tag_id = d_p_tags[(category, tag_name)]['id'] + count = d_p_tags[(category, tag_name)]['post_count'] + elif category < 9: + logging.warning(f'Cannot find tag {tag_name!r}, category: {category!r}.') + tag_info = _get_tag_by_name(tag_name) + if tag_info['name'] != tag_name: + logging.warning(f'Not found matching tags for {tag_name!r}, will be ignored.') + tag_id = -1 + count = -1 + else: + logging.info(f'Tag info found from danbooru - {tag_info!r}.') + tag_id = tag_info['id'] + count = tag_info['post_count'] + if category != tag_info['category']: + logging.warning(f'Category not match for tag {tag_name!r}, ' + f'replace category {category!r} --> {tag_info["category"]!r}') + category = tag_info['category'] + else: + logging.warning(f'Unknown tag {tag_name!r} ...') + tag_id = -1 + count = -1 + + rows.append({ + 'id': i, + 'tag_id': tag_id, + 'name': tag_name, + 'category': category, + 'count': count, + }) + + df = pd.DataFrame(rows) + return df + + +if __name__ == '__main__': + logging.try_init_root(level=logging.INFO) + df = load_tags() + df.to_csv('test_df.csv', index=False) diff --git a/zoo/utils/__init__.py b/zoo/utils/__init__.py index d699c746f2d..40d974e41f0 100644 --- a/zoo/utils/__init__.py +++ b/zoo/utils/__init__.py @@ -2,4 +2,5 @@ from .lr import get_init_lr, get_dynamic_lr_scheduler, LRTyping from .onnx import onnx_quick_export from .optimize import onnx_optimize +from .profile import torch_model_profile_via_thop, torch_model_profile_via_calflops from .testfile import get_testfile diff --git a/zoo/utils/profile.py b/zoo/utils/profile.py new file mode 100644 index 00000000000..169aabff244 --- /dev/null +++ b/zoo/utils/profile.py @@ -0,0 +1,39 @@ +import torch +from ditk import logging +from thop import profile, clever_format + + +def torch_model_profile_via_thop(model, input_): + with torch.no_grad(): + flops, params = profile(model, (input_,)) + + s_flops, s_params = clever_format([flops, params], "%.1f") + logging.info(f'Params: {s_params}, FLOPs: {s_flops}.') + + return flops, params + + +def torch_model_profile_via_calflops(model, input_): + from calflops import calculate_flops + flops, macs, params = calculate_flops( + model=model, + input_shape=tuple(input_.shape), + output_as_string=False, + print_detailed=False, + # output_as_string=True, + # output_precision=4 + ) + s_flops, s_params, s_macs = clever_format([flops, params, macs], "%.1f") + logging.info(f'Params: {s_params}, FLOPs: {s_flops}, MACs: {s_macs}.') + return flops, params, macs + + +if __name__ == '__main__': + logging.try_init_root(level=logging.INFO) + from timm import create_model + + # model = create_model('hf-hub:animetimm/swinv2_base_window8_256.dbv4-full', pretrained=False) + model = create_model('caformer_b36.sail_in22k_ft_in1k_384', pretrained=False) + dummy_input = torch.randn(1, 3, 448, 448) + print(torch_model_profile_via_thop(model, dummy_input)) + print(torch_model_profile_via_calflops(model, dummy_input))