Merge pull request #107 from deepghs/dev/yolo

dev(narugo): refactor yolo system design, add generic yolo modules

docs/source/api_doc/detect/booru_yolo_detect_benchmark.plot.py

@@ -1,17 +1,14 @@
-import os
 import random
 
 from hfutils.operate import get_hf_fs
 
 from benchmark import BaseBenchmark, create_plot_cli
-from imgutils.detect import detect_with_booru_yolo
+from imgutils.detect.booru_yolo import detect_with_booru_yolo, _REPO_ID
+from imgutils.generic.yolo import _open_models_for_repo_id
 
 repository = 'deepghs/booru_yolo'
 hf_fs = get_hf_fs()
-_MODELS = [
-    os.path.basename(os.path.dirname(file))
-    for file in hf_fs.glob(f'{repository}/*/model.onnx')
-]
+_MODELS = _open_models_for_repo_id(_REPO_ID).model_names
 
 
 class BooruYOLODetectBenchmark(BaseBenchmark):
@@ -20,14 +17,12 @@ def __init__(self, model_name: str):
         self.model_name = model_name
 
     def load(self):
-        from imgutils.detect.booru_yolo import _open_booru_yolo_model, _get_booru_yolo_labels
-        _ = _open_booru_yolo_model(self.model_name)
-        _ = _get_booru_yolo_labels(self.model_name)
+        from imgutils.generic.yolo import _open_models_for_repo_id
+        _open_models_for_repo_id(_REPO_ID)._open_model(self.model_name)
 
     def unload(self):
-        from imgutils.detect.booru_yolo import _open_booru_yolo_model, _get_booru_yolo_labels
-        _open_booru_yolo_model.cache_clear()
-        _get_booru_yolo_labels.cache_clear()
+        from imgutils.generic.yolo import _open_models_for_repo_id
+        _open_models_for_repo_id.cache_clear()
 
     def run(self):
         image_file = random.choice(self.all_images)

docs/source/api_doc/detect/censor_detect_benchmark.plot.py

@@ -1,37 +1,35 @@
 import random
 
 from benchmark import BaseBenchmark, create_plot_cli
-from imgutils.detect import detect_censors
+from imgutils.detect.censor import detect_censors, _REPO_ID
+from imgutils.generic.yolo import _open_models_for_repo_id
+
+_MODELS = _open_models_for_repo_id(_REPO_ID).model_names
 
 
 class CensorDetectBenchmark(BaseBenchmark):
-    def __init__(self, level, version):
+    def __init__(self, model_name: str):
         BaseBenchmark.__init__(self)
-        self.level = level
-        self.version = version
+        self.model_name = model_name
 
     def load(self):
-        from imgutils.detect.censor import _open_censor_detect_model
-        _ = _open_censor_detect_model(level=self.level, version=self.version)
+        from imgutils.generic.yolo import _open_models_for_repo_id
+        _ = _open_models_for_repo_id(_REPO_ID)._open_model(self.model_name)
 
     def unload(self):
-        from imgutils.detect.censor import _open_censor_detect_model
-        _open_censor_detect_model.cache_clear()
+        from imgutils.generic.yolo import _open_models_for_repo_id
+        _open_models_for_repo_id.cache_clear()
 
     def run(self):
         image_file = random.choice(self.all_images)
-        _ = detect_censors(image_file, level=self.level, version=self.version)
+        _ = detect_censors(image_file, model_name=self.model_name)
 
 
 if __name__ == '__main__':
     create_plot_cli(
         [
-            ('censor v1.0 (yolov8s)', CensorDetectBenchmark('s', 'v1.0')),
-            ('censor v1.0 (yolov8n)', CensorDetectBenchmark('n', 'v1.0')),
-            ('censor v0.10 (yolov8s)', CensorDetectBenchmark('s', 'v0.10')),
-            ('censor v0.9 (yolov8s)', CensorDetectBenchmark('s', 'v0.9')),
-            # ('censor v0.8 (yolov8s)', CensorDetectBenchmark('s', 'v0.8')),
-            # ('censor v0.7 (yolov8s)', CensorDetectBenchmark('s', 'v0.7')),
+            (model_name, CensorDetectBenchmark(model_name))
+            for model_name in _MODELS
         ],
         title='Benchmark for Anime Censor Detections',
         run_times=10,