Merge pull request #64 from deepghs/dev/safe

dev(narugo): add safe check model

docs/source/api_doc/validate/index.rst

@@ -18,6 +18,7 @@ imgutils.validate
     portrait
     rating
     real
+    safe
     style_age
     teen
     truncate

docs/source/api_doc/validate/safe.rst

@@ -0,0 +1,21 @@
+imgutils.validate.safe
+=============================================
+
+.. currentmodule:: imgutils.validate.safe
+
+.. automodule:: imgutils.validate.safe
+
+
+safe_check_score
+-----------------------------
+
+.. autofunction:: safe_check_score
+
+
+
+safe_check
+-----------------------------
+
+.. autofunction:: safe_check
+
+

docs/source/api_doc/validate/safe_benchmark.plot.py

@@ -0,0 +1,46 @@
+import os
+import random
+
+from huggingface_hub import HfFileSystem
+from natsort import natsorted
+
+from benchmark import BaseBenchmark, create_plot_cli
+from imgutils.validate import safe_check
+
+hf_fs = HfFileSystem()
+
+REPOSITORY = 'mf666/shit-checker'
+MODELS = natsorted([
+    os.path.splitext(os.path.relpath(file, REPOSITORY))[0]
+    for file in hf_fs.glob(f'{REPOSITORY}/*.onnx')
+])
+
+
+class SafeCheckBenchmark(BaseBenchmark):
+    def __init__(self, model):
+        BaseBenchmark.__init__(self)
+        self.model = model
+
+    def load(self):
+        from imgutils.validate.safe import _open_model
+        _ = _open_model(self.model)
+
+    def unload(self):
+        from imgutils.validate.safe import _open_model
+        _open_model.cache_clear()
+
+    def run(self):
+        image_file = random.choice(self.all_images)
+        _ = safe_check(image_file, self.model)
+
+
+if __name__ == '__main__':
+    create_plot_cli(
+        [
+            (name, SafeCheckBenchmark(name))
+            for name in MODELS
+        ],
+        title='Benchmark for Safe Check Models',
+        run_times=10,
+        try_times=20,
+    )()