Update vtoonify_model.py

vtoonify_model.py

@@ -22,6 +22,8 @@ import logging
 from PIL import Image
 
 
+
+
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 
@@ -113,6 +115,13 @@ class Model():
         
         return landmarks_68
 
+    def detect_and_align_image(self, frame_bgr: np.ndarray, top: int, bottom: int, left: int, right: int) -> tuple:
+        if frame_bgr is None:
+            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load the image.'
+        
+        # Proceed with detection and alignment
+        return self.detect_and_align(frame_bgr, top, bottom, left, right)
+
     def detect_and_align(self, frame, top, bottom, left, right, return_para=False):
         message = 'Error: no face detected! Please retry or change the photo.'
         instyle = None
@@ -177,15 +186,6 @@ class Model():
 
         return np.array(img)
 
-    def detect_and_align_image(self, image: str, top: int, bottom: int, left: int, right: int) -> tuple:
-        if image is None:
-            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load empty file.'
-        frame = cv2.imread(image)
-        if frame is None:
-            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load the image.'       
-        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-        return self.detect_and_align(frame, top, bottom, left, right)
-
     def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple:
         if instyle is None or aligned_face is None:
             return np.zeros((256, 256, 3), np.uint8), 'Oops, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
@@ -216,6 +216,4 @@ class Model():
         """Convert a tensor image to OpenCV format."""
         tmp = ((img.cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8).copy()
         logging.debug(f"Converted image shape: {tmp.shape}, strides: {tmp.strides}")
-        return cv2.cvtColor(tmp, cv2.COLOR_RGB2BGR)
-
-
+        return cv2.cvtColor(tmp, cv2.COLOR_RGB2BGR)