Merge branch 'dev'

Bing-su · Nov 13, 2024 · 358d170 · 358d170
2 parents 8ddf919 + 35c44a5
commit 358d170
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,10 @@
 # Changelog
 
+## 2024-11-13
+
+- v24.11.1
+- `mediapipe_face_mesh`, `mediapipe_face_mesh_eyes_only` 모델에 confidences가 없어 발생하는 에러 수정
+
 ## 2024-11-10
 
 - v24.11.0

diff --git a/adetailer/__version__.py b/adetailer/__version__.py
@@ -1 +1 @@
-__version__ = "24.11.0"
+__version__ = "24.11.1"
diff --git a/adetailer/mediapipe.py b/adetailer/mediapipe.py
@@ -97,6 +97,7 @@ def mediapipe_face_mesh(
 
         preview = arr.copy()
         masks = []
+        confidences = []
 
         for landmarks in pred.multi_face_landmarks:
             draw_util.draw_landmarks(
@@ -116,10 +117,13 @@ def mediapipe_face_mesh(
             draw = ImageDraw.Draw(mask)
             draw.polygon(outline, fill="white")
             masks.append(mask)
+            confidences.append(1.0)  # Confidence is unknown
 
         bboxes = create_bbox_from_mask(masks, image.size)
         preview = Image.fromarray(preview)
-        return PredictOutput(bboxes=bboxes, masks=masks, preview=preview)
+        return PredictOutput(
+            bboxes=bboxes, masks=masks, confidences=confidences, preview=preview
+        )
 
 
 def mediapipe_face_mesh_eyes_only(
@@ -145,6 +149,8 @@ def mediapipe_face_mesh_eyes_only(
 
         preview = image.copy()
         masks = []
+        confidences = []
+
         for landmarks in pred.multi_face_landmarks:
             points = np.array(
                 [[land.x * w, land.y * h] for land in landmarks.landmark], dtype=int
@@ -159,10 +165,13 @@ def mediapipe_face_mesh_eyes_only(
             for outline in (left_outline, right_outline):
                 draw.polygon(outline, fill="white")
             masks.append(mask)
+            confidences.append(1.0)  # Confidence is unknown
 
         bboxes = create_bbox_from_mask(masks, image.size)
         preview = draw_preview(preview, bboxes, masks)
-        return PredictOutput(bboxes=bboxes, masks=masks, preview=preview)
+        return PredictOutput(
+            bboxes=bboxes, masks=masks, confidences=confidences, preview=preview
+        )
 
 
 def draw_preview(

diff --git a/tests/test_mediapipe.py b/tests/test_mediapipe.py
@@ -16,3 +16,6 @@
 def test_mediapipe(sample_image2: Image.Image, model_name: str):
     result = mediapipe_predict(model_name, sample_image2)
     assert result.preview is not None
+    assert len(result.bboxes) > 0
+    assert len(result.masks) > 0
+    assert len(result.confidences) > 0
diff --git a/tests/test_ultralytics.py b/tests/test_ultralytics.py
@@ -48,3 +48,6 @@ def test_yolo_world(sample_image2: Image.Image, klass: str):
     model_path = hf_hub_download("Bingsu/yolo-world-mirror", "yolov8x-worldv2.pt")
     result = ultralytics_predict(model_path, sample_image2, classes=klass)
     assert result.preview is not None
+    assert len(result.bboxes) > 0
+    assert len(result.masks) > 0
+    assert len(result.confidences) > 0