dev(narugo): add 2 new models

app.py

@@ -3,7 +3,7 @@ import os
 import gradio as gr
 
 from detection import EyesDetection, FaceDetection, HeadDetection, PersonDetection, HandDetection, CensorDetection, \
-    HalfBodyDetection
+    HalfBodyDetection, NudeNetDetection, BooruYOLODetection
 
 _GLOBAL_CSS = """
 .limit-height {
@@ -36,5 +36,9 @@ if __name__ == '__main__':
                     HandDetection().make_ui()
                 with gr.Tab('Censor Point Detection'):
                     CensorDetection().make_ui()
+                with gr.Tab('NudeNet'):
+                    NudeNetDetection().make_ui()
+                with gr.Tab('BooruYOLO'):
+                    BooruYOLODetection().make_ui()
 
     demo.queue(os.cpu_count()).launch()

detection/__init__.py

@@ -1,8 +1,10 @@
 from .base import ObjectDetection, DeepGHSObjectDetection
+from .booru_yolo import BooruYOLODetection
 from .censor import CensorDetection
 from .eyes import EyesDetection
 from .face import FaceDetection
 from .halfbody import HalfBodyDetection
 from .hand import HandDetection
 from .head import HeadDetection
+from .nudenet import NudeNetDetection
 from .person import PersonDetection

detection/base.py

@@ -9,6 +9,14 @@ from hfutils.utils import hf_fs_path, parse_hf_fs_path
 from imgutils.data import ImageTyping
 
 
+def _v_fix(v):
+    return int(round(v))
+
+
+def _bbox_fix(bbox):
+    return tuple(map(_v_fix, bbox))
+
+
 class ObjectDetection:
     @lru_cache()
     def get_default_model(self) -> str:
@@ -51,7 +59,7 @@ class ObjectDetection:
         _color_map = dict(zip(labels, _colors))
         return gr.AnnotatedImage(
             value=(image, [
-                (bbox, label) for bbox, label, _ in
+                (_bbox_fix(bbox), label) for bbox, label, _ in
                 self.detect(image, model_name, iou_threshold, score_threshold)
             ]),
             color_map=_color_map,

detection/booru_yolo.py

@@ -0,0 +1,26 @@
+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 .base import DeepGHSObjectDetection
+
+
+class BooruYOLODetection(DeepGHSObjectDetection):
+    def __init__(self):
+        DeepGHSObjectDetection.__init__(self, repo_id='deepghs/booru_yolo')
+
+    def _get_default_model(self) -> str:
+        return _DEFAULT_MODEL
+
+    def _get_default_iou_and_score(self, model_name: str) -> Tuple[float, float]:
+        return 0.7, 0.25
+
+    def _get_labels(self, model_name: str) -> List[str]:
+        return _get_booru_yolo_labels(model_name)
+
+    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]]:
+        return detect_with_booru_yolo(image=image, model_name=model_name,
+                                      iou_threshold=iou_threshold, conf_threshold=score_threshold)

detection/nudenet.py

@@ -0,0 +1,25 @@
+from typing import List, Tuple
+
+from imgutils.data import ImageTyping
+from imgutils.detect.nudenet import detect_with_nudenet, _LABELS
+
+from .base import ObjectDetection
+
+
+class NudeNetDetection(ObjectDetection):
+    def _get_default_model(self) -> str:
+        return 'Default'
+
+    def _list_models(self) -> List[str]:
+        return ['Default']
+
+    def _get_default_iou_and_score(self, model_name: str) -> Tuple[float, float]:
+        return 0.45, 0.25
+
+    def _get_labels(self, model_name: str) -> List[str]:
+        return _LABELS
+
+    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]]:
+        return detect_with_nudenet(image=image, iou_threshold=iou_threshold, score_threshold=score_threshold)