From 34c49b438cc957f5ba41b9f08dbfb9bf75edc22c Mon Sep 17 00:00:00 2001
From: "Pedro E. Gonzalez V" <1pedrog@gmail.com>
Date: Fri, 25 Oct 2024 18:19:33 -0500
Subject: [PATCH] Adding YOLO detectors for specific use cases

---
 asdff/__init__.py | 12 +++++++++++-
 asdff/base.py     |  4 ++--
 asdff/yolo.py     | 28 ++++++++++++++++++++++++++++
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/asdff/__init__.py b/asdff/__init__.py
index e60473a..8096939 100644
--- a/asdff/__init__.py
+++ b/asdff/__init__.py
@@ -1,10 +1,20 @@
 from .__version__ import __version__
 from .sd import AdCnPipeline, AdPipeline
-from .yolo import yolo_detector
+from .yolo import (
+    yolo_detector,
+    yolo_face_detector,
+    yolo_hand_detector,
+    yolo_person_detector,
+    yolo_clothes_detector
+)
 
 __all__ = [
     "AdPipeline",
     "AdCnPipeline",
     "yolo_detector",
+    "yolo_face_detector",
+    "yolo_hand_detector",
+    "yolo_person_detector",
+    "yolo_clothes_detector",
     "__version__",
 ]
diff --git a/asdff/base.py b/asdff/base.py
index 19e8fb7..d939691 100644
--- a/asdff/base.py
+++ b/asdff/base.py
@@ -14,7 +14,7 @@
     mask_dilate,
     mask_gaussian_blur,
 )
-from asdff.yolo import yolo_detector
+from asdff.yolo import yolo_face_detector
 
 logger = logging.get_logger("diffusers")
 
@@ -122,7 +122,7 @@ def __call__(  # noqa: C901
 
     @property
     def default_detector(self) -> Callable[..., list[Image.Image] | None]:
-        return yolo_detector
+        return yolo_face_detector
 
     def _get_txt2img_args(
         self, common: Mapping[str, Any], txt2img_only: Mapping[str, Any]
diff --git a/asdff/yolo.py b/asdff/yolo.py
index f288e10..8c07a04 100644
--- a/asdff/yolo.py
+++ b/asdff/yolo.py
@@ -78,3 +78,31 @@ def yolo_detector(
         masks = mask_to_pil(pred[0].masks.data, image.size)
 
     return masks
+
+def yolo_face_detector(
+        image: Image.Image,
+        confidence: float = 0.3
+    ) -> list[Image.Image] | None:
+    model_path = hf_hub_download("Bingsu/adetailer", "face_yolov8n.pt")
+    return yolo_detector(image, model_path=model_path, confidence=confidence)
+
+def yolo_hand_detector(
+        image: Image.Image,
+        confidence: float = 0.3
+    ) -> list[Image.Image] | None:
+    model_path = hf_hub_download("Bingsu/adetailer", "hand_yolov9c.pt")
+    return yolo_detector(image, model_path=model_path, confidence=confidence)
+
+def yolo_person_detector(
+        image: Image.Image,
+        confidence: float = 0.3
+    ) -> list[Image.Image] | None:
+    model_path = hf_hub_download("Bingsu/adetailer", "person_yolov8s-seg.pt")
+    return yolo_detector(image, model_path=model_path, confidence=confidence)
+
+def yolo_clothes_detector(
+        image: Image.Image,
+        confidence: float = 0.3
+    ) -> list[Image.Image] | None:
+    model_path = hf_hub_download("Bingsu/adetailer", "deepfashion2_yolov8s-seg.pt")
+    return yolo_detector(image, model_path=model_path, confidence=confidence)