dev(naurgo): upgrade demo

detection/booru_yolo.py

@@ -1,7 +1,8 @@
 from typing import List, Tuple
 
 from imgutils.data import ImageTyping
-from imgutils.detect.booru_yolo import detect_with_booru_yolo, _get_booru_yolo_labels, _DEFAULT_MODEL
+from imgutils.detect.booru_yolo import detect_with_booru_yolo, _DEFAULT_MODEL
+from imgutils.generic.yolo import _open_models_for_repo_id
 
 from .base import DeepGHSObjectDetection
 
@@ -17,7 +18,8 @@ class BooruYOLODetection(DeepGHSObjectDetection):
         return 0.7, 0.25
 
     def _get_labels(self, model_name: str) -> List[str]:
-        return _get_booru_yolo_labels(model_name)
+        _, _, labels = _open_models_for_repo_id(self._repo_id)._open_model(model_name)
+        return labels
 
     def detect(self, image: ImageTyping, model_name: str,
                iou_threshold: float = 0.7, score_threshold: float = 0.25) -> \

detection/censor.py

@@ -2,7 +2,8 @@ import re
 from typing import List, Tuple
 
 from imgutils.data import ImageTyping
-from imgutils.detect.censor import detect_censors, _LABELS
+from imgutils.detect.censor import detect_censors
+from imgutils.generic.yolo import _open_models_for_repo_id
 
 from .base import DeepGHSObjectDetection
 
@@ -23,7 +24,8 @@ class CensorDetection(DeepGHSObjectDetection):
         return 0.7, 0.3
 
     def _get_labels(self, model_name: str) -> List[str]:
-        return _LABELS
+        _, _, labels = _open_models_for_repo_id(self._repo_id)._open_model(model_name)
+        return labels
 
     def detect(self, image: ImageTyping, model_name: str,
                iou_threshold: float = 0.7, score_threshold: float = 0.25) \

detection/eyes.py

@@ -2,7 +2,7 @@ import re
 from typing import List, Tuple
 
 from imgutils.data import ImageTyping
-from imgutils.detect.eye import detect_eyes, _LABELS
+from imgutils.detect.eye import detect_eyes
 
 from .base import DeepGHSObjectDetection
 
@@ -23,7 +23,7 @@ class EyesDetection(DeepGHSObjectDetection):
         return 0.3, 0.3
 
     def _get_labels(self, model_name: str) -> List[str]:
-        return _LABELS
+        return ['eye']
 
     def detect(self, image: ImageTyping, model_name: str,
                iou_threshold: float = 0.7, score_threshold: float = 0.25) \

detection/halfbody.py

@@ -2,7 +2,7 @@ import re
 from typing import List, Tuple
 
 from imgutils.data import ImageTyping
-from imgutils.detect.halfbody import detect_halfbody, _LABELS
+from imgutils.detect.halfbody import detect_halfbody
 
 from .base import DeepGHSObjectDetection
 
@@ -23,7 +23,7 @@ class HalfBodyDetection(DeepGHSObjectDetection):
         return 0.7, 0.5
 
     def _get_labels(self, model_name: str) -> List[str]:
-        return _LABELS
+        return ['halfbody']
 
     def detect(self, image: ImageTyping, model_name: str,
                iou_threshold: float = 0.7, score_threshold: float = 0.25) \

detection/hand.py

@@ -2,7 +2,7 @@ import re
 from typing import List, Tuple
 
 from imgutils.data import ImageTyping
-from imgutils.detect.hand import detect_hands, _LABELS
+from imgutils.detect.hand import detect_hands
 
 from .base import DeepGHSObjectDetection
 
@@ -23,7 +23,7 @@ class HandDetection(DeepGHSObjectDetection):
         return 0.7, 0.35
 
     def _get_labels(self, model_name: str) -> List[str]:
-        return _LABELS
+        return ['hand']
 
     def detect(self, image: ImageTyping, model_name: str,
                iou_threshold: float = 0.7, score_threshold: float = 0.25) \

detection/head.py

@@ -1,37 +1,17 @@
-import os.path
-import re
 from typing import List, Tuple
 
-from hfutils.operate import get_hf_fs
-from hfutils.utils import hf_fs_path, parse_hf_fs_path
 from imgutils.data import ImageTyping
 from imgutils.detect import detect_heads
 
-from .base import ObjectDetection
+from .base import DeepGHSObjectDetection
 
 
-def _parse_model_name(model_name: str):
-    matching = re.fullmatch(r'^head_detect_best_(?P<level>[\s\S]+?)$', model_name)
-    return matching.group('level')
-
-
-class HeadDetection(ObjectDetection):
+class HeadDetection(DeepGHSObjectDetection):
     def __init__(self):
-        self.repo_id = 'deepghs/imgutils-models'
+        DeepGHSObjectDetection.__init__(self, repo_id='deepghs/anime_head_detection')
 
     def _get_default_model(self) -> str:
-        return 'head_detect_best_s'
-
-    def _list_models(self) -> List[str]:
-        hf_fs = get_hf_fs()
-        return [
-            os.path.splitext(os.path.basename(parse_hf_fs_path(path).filename))[0]
-            for path in hf_fs.glob(hf_fs_path(
-                repo_id=self.repo_id,
-                repo_type='model',
-                filename='head_detect/*.onnx',
-            ))
-        ]
+        return 'head_detect_v0.5_s_pruned'
 
     def _get_default_iou_and_score(self, model_name: str) -> Tuple[float, float]:
         return 0.7, 0.3
@@ -42,6 +22,5 @@ class HeadDetection(ObjectDetection):
     def detect(self, image: ImageTyping, model_name: str,
                iou_threshold: float = 0.7, score_threshold: float = 0.25) -> \
             List[Tuple[Tuple[float, float, float, float], str, float]]:
-        level = _parse_model_name(model_name)
-        return detect_heads(image=image, level=level,
+        return detect_heads(image=image, model_name=model_name,
                             iou_threshold=iou_threshold, conf_threshold=score_threshold)

requirements.txt

@@ -8,4 +8,5 @@ pandas
 opencv-python>=4.6.0
 hbutils>=0.9.0
 dghs-imgutils>=0.5.0
+git+https://github.com/deepghs/imgutils.git
 httpx