Update vtoonify_model.py

vtoonify_model.py

@@ -22,8 +22,6 @@ import logging
 from PIL import Image
 
 
-
-
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 
@@ -98,6 +96,23 @@ class Model():
             exstyle = self.vtoonify.zplus2wplus(exstyle)
         return exstyle, 'Model of %s loaded.' % (style_type)
 
+    def convert_106_to_68(self, landmarks_106):
+        # Mapping from 106 landmarks to 68 landmarks
+        landmark106to68 = [
+            1, 10, 12, 14, 16, 3, 5, 7, 0, 23, 21, 19, 32, 30, 28, 26, 17,  # Face outline
+            43, 48, 49, 51, 50,  # Left eyebrow
+            102, 103, 104, 105, 101,  # Right eyebrow
+            72, 73, 74, 86, 78, 79, 80, 85, 84,  # Nose
+            35, 41, 42, 39, 37, 36,  # Left eye
+            89, 95, 96, 93, 91, 90,  # Right eye
+            52, 64, 63, 71, 67, 68, 61, 58, 59, 53, 56, 55, 65, 66, 62, 70, 69, 57, 60, 54  # Mouth
+        ]
+        
+        # Convert 106 landmarks to 68 landmarks
+        landmarks_68 = [landmarks_106[index] for index in landmark106to68]
+        
+        return landmarks_68
+
     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
@@ -107,10 +122,11 @@ class Model():
         if len(faces) > 0:
             logging.info(f"Detected {len(faces)} face(s).")
             face = faces[0]
-            landmarks = face.landmark_2d_106
+            landmarks_106 = face.landmark_2d_106
+            landmarks_68 = self.convert_106_to_68(landmarks_106)
 
             # Align face based on mapped landmarks
-            aligned_face = self.align_face(frame, landmarks)
+            aligned_face = self.align_face(frame, landmarks_68)
             if aligned_face is not None:
                 logging.info(f"Aligned face shape: {aligned_face.shape}")
                 with torch.no_grad():
@@ -130,8 +146,7 @@ class Model():
         return frame, instyle, message
 
     def align_face(self, image, landmarks):
-        # Map InsightFace landmarks to dlib's 68-point model
-        # Example: use specific indices for eyes and mouth
+        # Example alignment logic using 68 landmarks
         eye_left = np.mean(landmarks[36:42], axis=0)
         eye_right = np.mean(landmarks[42:48], axis=0)
         mouth_left = landmarks[48]
@@ -161,22 +176,16 @@ class Model():
             img = img.resize((output_size, output_size), Image.ANTIALIAS)
 
         return np.array(img)
-    def detect_and_align_image(self, image_data: bytes, top: int, bottom: int, left: int, right: int) -> tuple:
-        if image_data is None:
+
+    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.'
-    
-    # Convert bytes to a NumPy array
-        nparr = np.frombuffer(image_data, np.uint8)
-    
-    # Decode the image data
-        frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-    
+        frame = cv2.imread(image)
         if frame is None:
-            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to decode the image.'
-    
+            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.'
@@ -207,4 +216,6 @@ 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)
+
+