Merge pull request #31 from deepghs/dev/nsfw

dev(narugo): add nsfw prediction from nsfwjs library in javascript
deepghs · Jun 26, 2023 · 3347abe · 3347abe
2 parents 5b318d1 + c671c37
commit 3347abe
Show file tree

Hide file tree

Showing 69 changed files with 4,206 additions and 0 deletions.
diff --git a/docs/source/api_doc/validate/index.rst b/docs/source/api_doc/validate/index.rst
@@ -13,5 +13,6 @@ imgutils.validate
     classify
     color
     monochrome
+    nsfw
     rating
     truncate
diff --git a/docs/source/api_doc/validate/nsfw.plot.py b/docs/source/api_doc/validate/nsfw.plot.py
@@ -0,0 +1,16 @@
+import glob
+import os.path
+
+from natsort import natsorted
+
+from plot import image_plot
+
+if __name__ == '__main__':
+    image_plot(
+        *natsorted(glob.glob(os.path.join('nsfw', 'drawings', '*.jpg'))),
+        *natsorted(glob.glob(os.path.join('nsfw', 'hentai', '*.jpg'))),
+        *natsorted(glob.glob(os.path.join('nsfw', 'neutral', '*.jpg'))),
+        *natsorted(glob.glob(os.path.join('nsfw', 'porn', '*.jpg'))),
+        *natsorted(glob.glob(os.path.join('nsfw', 'sexy', '*.jpg'))),
+        columns=4, figsize=(10, 18.5),
+    )
diff --git a/docs/source/api_doc/validate/nsfw.plot.py.svg b/docs/source/api_doc/validate/nsfw.plot.py.svg
diff --git a/docs/source/api_doc/validate/nsfw.rst b/docs/source/api_doc/validate/nsfw.rst
@@ -0,0 +1,22 @@
+imgutils.validate.nsfw
+=============================================
+
+.. currentmodule:: imgutils.validate.nsfw
+
+.. automodule:: imgutils.validate.nsfw
+
+
+nsfw_pred_score
+----------------------------------------
+
+.. autofunction:: nsfw_pred_score
+
+
+
+nsfw_pred
+----------------------------------------
+
+.. autofunction:: nsfw_pred
+
+
+
diff --git a/docs/source/api_doc/validate/nsfw/drawings/1.jpg b/docs/source/api_doc/validate/nsfw/drawings/1.jpg
diff --git a/docs/source/api_doc/validate/nsfw/drawings/2.jpg b/docs/source/api_doc/validate/nsfw/drawings/2.jpg
diff --git a/docs/source/api_doc/validate/nsfw/drawings/3.jpg b/docs/source/api_doc/validate/nsfw/drawings/3.jpg
diff --git a/docs/source/api_doc/validate/nsfw/drawings/4.jpg b/docs/source/api_doc/validate/nsfw/drawings/4.jpg
diff --git a/docs/source/api_doc/validate/nsfw/hentai/5.jpg b/docs/source/api_doc/validate/nsfw/hentai/5.jpg
diff --git a/docs/source/api_doc/validate/nsfw/hentai/6.jpg b/docs/source/api_doc/validate/nsfw/hentai/6.jpg
diff --git a/docs/source/api_doc/validate/nsfw/hentai/7.jpg b/docs/source/api_doc/validate/nsfw/hentai/7.jpg
diff --git a/docs/source/api_doc/validate/nsfw/hentai/8.jpg b/docs/source/api_doc/validate/nsfw/hentai/8.jpg
diff --git a/docs/source/api_doc/validate/nsfw/neutral/10.jpg b/docs/source/api_doc/validate/nsfw/neutral/10.jpg
diff --git a/docs/source/api_doc/validate/nsfw/neutral/11.jpg b/docs/source/api_doc/validate/nsfw/neutral/11.jpg
diff --git a/docs/source/api_doc/validate/nsfw/neutral/12.jpg b/docs/source/api_doc/validate/nsfw/neutral/12.jpg
diff --git a/docs/source/api_doc/validate/nsfw/neutral/9.jpg b/docs/source/api_doc/validate/nsfw/neutral/9.jpg
diff --git a/docs/source/api_doc/validate/nsfw/porn/13.jpg b/docs/source/api_doc/validate/nsfw/porn/13.jpg
diff --git a/docs/source/api_doc/validate/nsfw/porn/14.jpg b/docs/source/api_doc/validate/nsfw/porn/14.jpg
diff --git a/docs/source/api_doc/validate/nsfw/porn/15.jpg b/docs/source/api_doc/validate/nsfw/porn/15.jpg
diff --git a/docs/source/api_doc/validate/nsfw/porn/16.jpg b/docs/source/api_doc/validate/nsfw/porn/16.jpg
diff --git a/docs/source/api_doc/validate/nsfw/sexy/17.jpg b/docs/source/api_doc/validate/nsfw/sexy/17.jpg
diff --git a/docs/source/api_doc/validate/nsfw/sexy/18.jpg b/docs/source/api_doc/validate/nsfw/sexy/18.jpg
diff --git a/docs/source/api_doc/validate/nsfw/sexy/19.jpg b/docs/source/api_doc/validate/nsfw/sexy/19.jpg
diff --git a/docs/source/api_doc/validate/nsfw/sexy/20.jpg b/docs/source/api_doc/validate/nsfw/sexy/20.jpg
diff --git a/docs/source/api_doc/validate/nsfw_benchmark.plot.py b/docs/source/api_doc/validate/nsfw_benchmark.plot.py
@@ -0,0 +1,35 @@
+import random
+
+from benchmark import BaseBenchmark, create_plot_cli
+from imgutils.validate import nsfw_pred_score
+from imgutils.validate.nsfw import _MODEL_NAMES
+
+
+class AnimeNSFWBenchmark(BaseBenchmark):
+    def __init__(self, model):
+        BaseBenchmark.__init__(self)
+        self.model = model
+
+    def load(self):
+        from imgutils.validate.nsfw import _open_nsfw_model
+        _ = _open_nsfw_model(self.model)
+
+    def unload(self):
+        from imgutils.validate.nsfw import _open_nsfw_model
+        _open_nsfw_model.cache_clear()
+
+    def run(self):
+        image_file = random.choice(self.all_images)
+        _ = nsfw_pred_score(image_file, self.model)
+
+
+if __name__ == '__main__':
+    create_plot_cli(
+        [
+            (name, AnimeNSFWBenchmark(name))
+            for name in _MODEL_NAMES
+        ],
+        title='Benchmark for NSFW Models',
+        run_times=10,
+        try_times=20,
+    )()