Update vtoonify_model.py

vtoonify_model.py

@@ -8,7 +8,7 @@ from util import load_psp_standalone, get_video_crop_parameter, tensor2cv2
 import torch
 import torch.nn as nn
 import numpy as np
-import dlib
+import insightface
 import cv2
 from model.vtoonify import VToonify
 from model.bisenet.model import BiSeNet
@@ -51,13 +51,13 @@ class Model():
             'illustration5-d': ['vtoonify_d_illustration/vtoonify_s086_d_c.pt', 86],
         }
         
-        self.landmarkpredictor = self._create_dlib_landmark_model()
+        self.face_detector = self._create_insightface_detector()
         self.parsingpredictor = self._create_parsing_model()
         self.pspencoder = self._load_encoder()    
         self.transform = transforms.Compose([
             transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5,0.5,0.5]),
-            ])
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
         
         self.vtoonify, self.exstyle = self._load_default_model()
         self.color_transfer = False
@@ -65,10 +65,11 @@ class Model():
         self.video_limit_cpu = 100
         self.video_limit_gpu = 300
         
-    @staticmethod
-    def _create_dlib_landmark_model():
-        return dlib.shape_predictor(huggingface_hub.hf_hub_download(MODEL_REPO,
-                                                    'models/shape_predictor_68_face_landmarks.dat'))
+    def _create_insightface_detector(self):
+        # Initialize InsightFace
+        app = insightface.app.FaceAnalysis()
+        app.prepare(ctx_id=0, det_size=(640, 640))  # ctx_id=-1 for CPU, 0 for GPU
+        return app
     
     def _create_parsing_model(self):
         parsingpredictor = BiSeNet(n_classes=19)
@@ -78,93 +79,118 @@ class Model():
         return parsingpredictor
     
     def _load_encoder(self) -> nn.Module:
-        style_encoder_path = huggingface_hub.hf_hub_download(MODEL_REPO,'models/encoder.pt')
+        style_encoder_path = huggingface_hub.hf_hub_download(MODEL_REPO, 'models/encoder.pt')
         return load_psp_standalone(style_encoder_path, self.device)
     
     def _load_default_model(self) -> tuple[torch.Tensor, str]:
-        vtoonify = VToonify(backbone = 'dualstylegan')
+        vtoonify = VToonify(backbone='dualstylegan')
         vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO,
                                             'models/vtoonify_d_cartoon/vtoonify_s026_d0.5.pt'), 
                                             map_location=lambda storage, loc: storage)['g_ema'])
         vtoonify.to(self.device)
-        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO,'models/vtoonify_d_cartoon/exstyle_code.npy'), allow_pickle=True).item()
+        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/vtoonify_d_cartoon/exstyle_code.npy'), allow_pickle=True).item()
         exstyle = torch.tensor(tmp[list(tmp.keys())[26]]).to(self.device)
         with torch.no_grad():  
             exstyle = vtoonify.zplus2wplus(exstyle)
         return vtoonify, exstyle
     
-    def load_model(self, style_type: str) -> tuple[torch.Tensor, str]:
-        if 'illustration' in style_type:
-            self.color_transfer = True
-        else:
-            self.color_transfer = False
-        if style_type not in self.style_types.keys():
-            return None, 'Oops, wrong Style Type. Please select a valid model.'
-        self.style_name = style_type
-        model_path, ind = self.style_types[style_type]
-        style_path = os.path.join('models',os.path.dirname(model_path),'exstyle_code.npy')
-        self.vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO,'models/'+model_path), 
-                                            map_location=lambda storage, loc: storage)['g_ema'])
-        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO, style_path), allow_pickle=True).item()
-        exstyle = torch.tensor(tmp[list(tmp.keys())[ind]]).to(self.device)
-        with torch.no_grad():  
-            exstyle = self.vtoonify.zplus2wplus(exstyle)
-        return exstyle, 'Model of %s loaded.'%(style_type)
-    
     def detect_and_align(self, frame, top, bottom, left, right, return_para=False):
         message = 'Error: no face detected! Please retry or change the photo.'
-        paras = get_video_crop_parameter(frame, self.landmarkpredictor, [left, right, top, bottom])
-        instyle = None
-        h, w, scale = 0, 0, 0
-        if paras is not None:
-            h,w,top,bottom,left,right,scale = paras
+        # Use InsightFace for face detection
+        faces = self.face_detector.get(frame)
+        if len(faces) > 0:
+            face = faces[0]
+            bbox = face.bbox.astype(int)
+            x, y, w, h = bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]
+            top, bottom, left, right = y, y + h, x, x + w
+            scale = 1.0  # Adjust scale as needed
+            h, w = frame.shape[:2]
             H, W = int(bottom-top), int(right-left)
             # for HR image, we apply gaussian blur to it to avoid over-sharp stylization results
-            kernel_1d = np.array([[0.125],[0.375],[0.375],[0.125]])
+            kernel_1d = np.array([[0.125], [0.375], [0.375], [0.125]])
             if scale <= 0.75:
                 frame = cv2.sepFilter2D(frame, -1, kernel_1d, kernel_1d)
             if scale <= 0.375:
                 frame = cv2.sepFilter2D(frame, -1, kernel_1d, kernel_1d)
             frame = cv2.resize(frame, (w, h))[top:bottom, left:right]
             with torch.no_grad():
-                I = align_face(frame, self.landmarkpredictor)
+                I = align_face(frame, self.face_detector)
                 if I is not None:
                     I = self.transform(I).unsqueeze(dim=0).to(self.device)
                     instyle = self.pspencoder(I)
                     instyle = self.vtoonify.zplus2wplus(instyle)
-                    message = 'Successfully rescale the frame to (%d, %d)'%(bottom-top, right-left)
+                    message = 'Successfully rescale the frame to (%d, %d)' % (bottom-top, right-left)
                 else:
-                    frame = np.zeros((256,256,3), np.uint8)
+                    frame = np.zeros((256, 256, 3), np.uint8)
         else:
-            frame = np.zeros((256,256,3), np.uint8)
+            frame = np.zeros((256, 256, 3), np.uint8)
         if return_para:
             return frame, instyle, message, w, h, top, bottom, left, right, scale
         return frame, instyle, message
     
-    #@torch.inference_mode()
-    def detect_and_align_image(self, frame_rgb: np.ndarray, top: int, bottom: int, left: int, right: int) -> tuple[np.ndarray, torch.Tensor, str]:
-        if frame_rgb is None:
-           return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load the image.'
+    # Other methods remain unchanged
+    def _create_parsing_model(self):
+        parsingpredictor = BiSeNet(n_classes=19)
+        parsingpredictor.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/faceparsing.pth'),
+                                                    map_location=lambda storage, loc: storage))
+        parsingpredictor.to(self.device).eval()
+        return parsingpredictor
+    
+    def _load_encoder(self) -> nn.Module:
+        style_encoder_path = huggingface_hub.hf_hub_download(MODEL_REPO, 'models/encoder.pt')
+        return load_psp_standalone(style_encoder_path, self.device)
+    
+    def _load_default_model(self) -> tuple:
+        vtoonify = VToonify(backbone='dualstylegan')
+        vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO,
+                                            'models/vtoonify_d_cartoon/vtoonify_s026_d0.5.pt'), 
+                                            map_location=lambda storage, loc: storage)['g_ema'])
+        vtoonify.to(self.device)
+        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/vtoonify_d_cartoon/exstyle_code.npy'), allow_pickle=True).item()
+        exstyle = torch.tensor(tmp[list(tmp.keys())[26]]).to(self.device)
+        with torch.no_grad():  
+            exstyle = vtoonify.zplus2wplus(exstyle)
+        return vtoonify, exstyle
+    
+    def load_model(self, style_type: str) -> tuple:
+        if 'illustration' in style_type:
+            self.color_transfer = True
+        else:
+            self.color_transfer = False
+        if style_type not in self.style_types.keys():
+            return None, 'Oops, wrong Style Type. Please select a valid model.'
+        self.style_name = style_type
+        model_path, ind = self.style_types[style_type]
+        style_path = os.path.join('models', os.path.dirname(model_path), 'exstyle_code.npy')
+        self.vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/' + model_path), 
+                                            map_location=lambda storage, loc: storage)['g_ema'])
+        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO, style_path), allow_pickle=True).item()
+        exstyle = torch.tensor(tmp[list(tmp.keys())[ind]]).to(self.device)
+        with torch.no_grad():  
+            exstyle = self.vtoonify.zplus2wplus(exstyle)
+        return exstyle, 'Model of %s loaded.' % (style_type)
 
-    # Convert RGB to BGR
+    def detect_and_align_image(self, frame_rgb: np.ndarray, top: int, bottom: int, left: int, right: int) -> tuple:
+        if frame_rgb is None:
+            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load the image.'
+        
+        # Convert RGB to BGR
         frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
-
         return self.detect_and_align(frame_bgr, top, bottom, left, right)
 
-    def detect_and_align_video(self, video: str, top: int, bottom: int, left: int, right: int
-                              ) -> tuple[np.ndarray, torch.Tensor, str]:
+    def detect_and_align_video(self, video: str, top: int, bottom: int, left: int, right: int) -> tuple:
         if video is None:
-            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load empty file.'
+            return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load empty file.'
         video_cap = cv2.VideoCapture(video)
         if video_cap.get(7) == 0:
             video_cap.release()
-            return np.zeros((256,256,3), np.uint8), torch.zeros(1,18,512).to(self.device), 'Error: fail to load the video.'
+            return np.zeros((256, 256, 3), np.uint8), torch.zeros(1, 18, 512).to(self.device), 'Error: fail to load the video.'
         success, frame = video_cap.read()
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         video_cap.release()
         return self.detect_and_align(frame, top, bottom, left, right)
-    
-    def detect_and_align_full_video(self, video: str, top: int, bottom: int, left: int, right: int) -> tuple[str, torch.Tensor, str]:
+
+    def detect_and_align_full_video(self, video: str, top: int, bottom: int, left: int, right: int) -> tuple:
         message = 'Error: no face detected! Please retry or change the video.'
         instyle = None
         if video is None:
@@ -172,7 +198,7 @@ class Model():
         video_cap = cv2.VideoCapture(video)
         if video_cap.get(7) == 0:
             video_cap.release()
-            return 'default.mp4', instyle, 'Error: fail to load the video.'    
+            return 'default.mp4', instyle, 'Error: fail to load the video.'
         num = min(self.video_limit_gpu, int(video_cap.get(7)))
         if self.device == 'cpu':
             num = min(self.video_limit_cpu, num)
@@ -180,14 +206,14 @@ class Model():
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         frame, instyle, message, w, h, top, bottom, left, right, scale = self.detect_and_align(frame, top, bottom, left, right, True)
         if instyle is None:
-            return 'default.mp4', instyle, message    
+            return 'default.mp4', instyle, message
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         videoWriter = cv2.VideoWriter('input.mp4', fourcc, video_cap.get(5), (int(right-left), int(bottom-top)))
         videoWriter.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-        kernel_1d = np.array([[0.125],[0.375],[0.375],[0.125]])
+        kernel_1d = np.array([[0.125], [0.375], [0.375], [0.125]])
         for i in range(num-1):
             success, frame = video_cap.read()
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)            
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             if scale <= 0.75:
                 frame = cv2.sepFilter2D(frame, -1, kernel_1d, kernel_1d)
             if scale <= 0.375:
@@ -198,42 +224,40 @@ class Model():
         videoWriter.release()
         video_cap.release()
 
-        return 'input.mp4', instyle, 'Successfully rescale the video to (%d, %d)'%(bottom-top, right-left)
-    
-    def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple[np.ndarray, str]:
-        #print(style_type + ' ' + self.style_name)
+        return 'input.mp4', instyle, 'Successfully rescale the video to (%d, %d)' % (bottom-top, right-left)
+
+    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), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
+            return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
         if self.style_name != style_type:
-            exstyle, _  = self.load_model(style_type)
+            exstyle, _ = self.load_model(style_type)
         if exstyle is None:
-            return np.zeros((256,256,3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
+            return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
         with torch.no_grad():
             if self.color_transfer:
                 s_w = exstyle
             else:
                 s_w = instyle.clone()
-                s_w[:,:7] = exstyle[:,:7]
+                s_w[:, :7] = exstyle[:, :7]
 
             x = self.transform(aligned_face).unsqueeze(dim=0).to(self.device)
             x_p = F.interpolate(self.parsingpredictor(2*(F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0], 
                                 scale_factor=0.5, recompute_scale_factor=False).detach()
             inputs = torch.cat((x, x_p/16.), dim=1)
-            y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s = style_degree)        
+            y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s=style_degree)        
             y_tilde = torch.clamp(y_tilde, -1, 1)
-        print('*** Toonify %dx%d image with style of %s'%(y_tilde.shape[2], y_tilde.shape[3], style_type))
-        return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s'%(self.style_name)
-    
-    def video_tooniy(self, aligned_video: str, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple[str, str]:
-        #print(style_type + ' ' + self.style_name)
+        print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
+        return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s' % (self.style_name)
+
+    def video_toonify(self, aligned_video: str, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple:
         if aligned_video is None:
-            return 'default.mp4', 'Opps, something wrong with the input. Please go to Step 2 and Rescale Video again.'         
+            return 'default.mp4', 'Opps, something wrong with the input. Please go to Step 2 and Rescale Video again.'
         video_cap = cv2.VideoCapture(aligned_video)
         if instyle is None or aligned_video is None or video_cap.get(7) == 0:
             video_cap.release()
             return 'default.mp4', 'Opps, something wrong with the input. Please go to Step 2 and Rescale Video again.'
         if self.style_name != style_type:
-            exstyle, _  = self.load_model(style_type)
+            exstyle, _ = self.load_model(style_type)
         if exstyle is None:
             return 'default.mp4', 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
         num = min(self.video_limit_gpu, int(video_cap.get(7)))
@@ -247,18 +271,18 @@ class Model():
         batch_frames = []
         if video_cap.get(3) != 0:
             if self.device == 'cpu':
-                batch_size = max(1, int(4 * 256* 256/ video_cap.get(3) / video_cap.get(4)))
+                batch_size = max(1, int(4 * 256 * 256 / video_cap.get(3) / video_cap.get(4)))
             else:
-                batch_size = min(max(1, int(4 * 400 * 360/ video_cap.get(3) / video_cap.get(4))), 4)
+                batch_size = min(max(1, int(4 * 400 * 360 / video_cap.get(3) / video_cap.get(4))), 4)
         else:
             batch_size = 1
-        print('*** Toonify using batch size of %d on %dx%d video of %d frames with style of %s'%(batch_size, int(video_cap.get(3)*4), int(video_cap.get(4)*4), num, style_type))
+        print('*** Toonify using batch size of %d on %dx%d video of %d frames with style of %s' % (batch_size, int(video_cap.get(3)*4), int(video_cap.get(4)*4), num, style_type))
         with torch.no_grad():
             if self.color_transfer:
                 s_w = exstyle
             else:
                 s_w = instyle.clone()
-                s_w[:,:7] = exstyle[:,:7]
+                s_w[:, :7] = exstyle[:, :7]
             for i in range(num):
                 success, frame = video_cap.read()
                 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
@@ -278,6 +302,10 @@ class Model():
 
         videoWriter.release()
         video_cap.release()
-        return 'output.mp4', 'Successfully toonify video of %d frames with style of %s'%(num, self.style_name)
-
+        return 'output.mp4', 'Successfully toonify video of %d frames with style of %s' % (num, self.style_name)
 
+    def tensor2cv2(self, img):
+        """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)