From f9fc0bbff1d4126e34904a930597bee010851d99 Mon Sep 17 00:00:00 2001
From: Tran Xen <137925069+glucauze@users.noreply.github.com>
Date: Tue, 15 Aug 2023 14:58:10 +0200
Subject: [PATCH] fix install wip

---
 requirements-gpu.txt                  | 2 ++
 requirements.txt                      | 1 +
 scripts/faceswaplab_utils/imgutils.py | 7 +++++--
 scripts/faceswaplab_utils/typing.py   | 4 ++--
 4 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/requirements-gpu.txt b/requirements-gpu.txt
index a26f6c2..e914f8d 100644
--- a/requirements-gpu.txt
+++ b/requirements-gpu.txt
@@ -3,6 +3,8 @@ dill
 ifnude
 insightface==0.7.3
 onnx>=1.14.0
+protobuf>=3.20.2
+opencv-python
 pandas
 pydantic
 safetensors
diff --git a/requirements.txt b/requirements.txt
index f6581d5..cf934a2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+protobuf>=3.20.2
 cython
 dill
 ifnude
diff --git a/scripts/faceswaplab_utils/imgutils.py b/scripts/faceswaplab_utils/imgutils.py
index e8d67fb..3b55b6c 100644
--- a/scripts/faceswaplab_utils/imgutils.py
+++ b/scripts/faceswaplab_utils/imgutils.py
@@ -29,12 +29,15 @@ def check_against_nsfw(img: PILImage) -> bool:
     chunks: List[Dict[str, Union[int, float]]] = detect(img)
 
     for chunk in chunks:
+        logger.debug(
+            f"chunck score {chunk['score']}, threshold : {NSFW_SCORE_THRESHOLD}"
+        )
         shapes.append(chunk["score"] > NSFW_SCORE_THRESHOLD)
 
     return any(shapes)
 
 
-def pil_to_cv2(pil_img: PILImage) -> CV2ImgU8:  # type: ignore
+def pil_to_cv2(pil_img: PILImage) -> CV2ImgU8:
     """
     Convert a PIL Image into an OpenCV image (cv2).
 
@@ -44,7 +47,7 @@ def pil_to_cv2(pil_img: PILImage) -> CV2ImgU8:  # type: ignore
     Returns:
         CV2ImgU8: The input image converted to OpenCV format (BGR).
     """
-    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR).astype("uint8")
 
 
 def cv2_to_pil(cv2_img: CV2ImgU8) -> PILImage:  # type: ignore
diff --git a/scripts/faceswaplab_utils/typing.py b/scripts/faceswaplab_utils/typing.py
index d102d63..fb14503 100644
--- a/scripts/faceswaplab_utils/typing.py
+++ b/scripts/faceswaplab_utils/typing.py
@@ -1,10 +1,10 @@
 from typing import Tuple
 from numpy import uint8
-from numpy.typing import NDArray
 from insightface.app.common import Face as IFace
 from PIL import Image
+import numpy as np
 
 PILImage = Image.Image
-CV2ImgU8 = NDArray[uint8]
+CV2ImgU8 = np.ndarray[int, np.dtype[uint8]]
 Face = IFace
 BoxCoords = Tuple[int, int, int, int]