From 3079400d2ff9a4331f167ba1e76c1be8158ca0c7 Mon Sep 17 00:00:00 2001
From: Bingsu <ks2515@naver.com>
Date: Wed, 26 Apr 2023 15:38:43 +0900
Subject: [PATCH] feat: mediapipe predict

---
 adetailer/common.py      |  1 -
 adetailer/mediapipe.py   | 48 ++++++++++++++++++++++++++++++++++++++++
 adetailer/ultralytics.py |  5 ++++-
 3 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 adetailer/mediapipe.py

diff --git a/adetailer/common.py b/adetailer/common.py
index 44bf5d6..cac36fb 100644
--- a/adetailer/common.py
+++ b/adetailer/common.py
@@ -3,7 +3,6 @@ from __future__ import annotations
 from collections import OrderedDict
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional
 
 from huggingface_hub import hf_hub_download
 from PIL import Image, ImageDraw
diff --git a/adetailer/mediapipe.py b/adetailer/mediapipe.py
new file mode 100644
index 0000000..a003534
--- /dev/null
+++ b/adetailer/mediapipe.py
@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import numpy as np
+import mediapipe as mp
+from PIL import Image
+
+from adetailer import PredictOutput
+from adetailer.common import create_mask_from_bbox
+
+
+def mediapipe_predict(
+    model_type: int, image: Image.Image, confidence: float = 0.25
+) -> PredictOutput:
+    img_width, img_height = image.size
+
+    mp_face_detection = mp.solutions.face_detection
+    draw_util = mp.solutions.drawing_utils
+
+    img_array = np.array(image)
+
+    with mp_face_detection.FaceDetection(
+        model_selection=model_type, min_detection_confidence=confidence
+    ) as face_detector:
+        pred = face_detector.process(img_array)
+
+    if pred.detections is None:
+        return PredictOutput()
+
+    example_array = img_array.copy()
+
+    bboxes = []
+    for detection in pred.detections:
+        draw_util.draw_detection(example_array, detection)
+
+        bbox = detection.location_data.relative_bounding_box
+        x1 = bbox.xmin * img_width
+        y1 = bbox.ymin * img_height
+        w = bbox.width * img_width
+        h = bbox.height * img_height
+        x2 = x1 + w
+        y2 = y1 + h
+
+        bboxes.append([x1, y1, x2, y2])
+
+    masks = create_mask_from_bbox(image, bboxes)
+    example = Image.fromarray(example_array)
+
+    return PredictOutput(bboxes=bboxes, masks=masks, example=example)
diff --git a/adetailer/ultralytics.py b/adetailer/ultralytics.py
index c55bfb9..ef65079 100644
--- a/adetailer/ultralytics.py
+++ b/adetailer/ultralytics.py
@@ -18,7 +18,10 @@ def ultralytics_predict(
     model = YOLO(model_path)
     pred = model(image, conf=confidence, hide_labels=True)
 
-    bboxes = pred[0].xyxy.cpu().numpy()
+    bboxes = pred[0].xyxy.cpu().numpy().tolist()
+    if len(bboxes) == 0:
+        return PredictOutput()
+
     masks = create_mask_from_bbox(image, bboxes)
     example = pred[0].plot()
     example = cv2.cvtColor(example, cv2.COLOR_BGR2RGB)