Update vtoonify_model.py

vtoonify_model.py

@@ -1,10 +1,4 @@
 from __future__ import annotations
-import gradio as gr
-import pathlib
-import sys
-sys.path.insert(0, 'vtoonify')
-
-from util import load_psp_standalone, get_video_crop_parameter, tensor2cv2
 import torch
 import torch.nn as nn
 import numpy as np
@@ -14,9 +8,7 @@ from model.vtoonify import VToonify
 from model.bisenet.model import BiSeNet
 import torch.nn.functional as F
 from torchvision import transforms
-import gc
 import huggingface_hub
-import os
 import logging
 from PIL import Image
 
@@ -28,65 +20,43 @@ MODEL_REPO = 'PKUWilliamYang/VToonify'
 class Model():
     def __init__(self, device):
         super().__init__()
-        
+
         self.device = device
         self.style_types = {
             'cartoon1': ['vtoonify_d_cartoon/vtoonify_s026_d0.5.pt', 26],
-            'cartoon1-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 26],
-            'cartoon2-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 64],
-            'cartoon3-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 153],
-            'cartoon4': ['vtoonify_d_cartoon/vtoonify_s299_d0.5.pt', 299],
-            'cartoon4-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 299],
-            'cartoon5-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 8],
-            'comic1-d': ['vtoonify_d_comic/vtoonify_s_d.pt', 28],
-            'comic2-d': ['vtoonify_d_comic/vtoonify_s_d.pt', 18],
-            'arcane1': ['vtoonify_d_arcane/vtoonify_s000_d0.5.pt', 0],
-            'arcane1-d': ['vtoonify_d_arcane/vtoonify_s_d.pt', 0],
-            'arcane2': ['vtoonify_d_arcane/vtoonify_s077_d0.5.pt', 77],
-            'arcane2-d': ['vtoonify_d_arcane/vtoonify_s_d.pt', 77],
-            'caricature1': ['vtoonify_d_caricature/vtoonify_s039_d0.5.pt', 39],
-            'caricature2': ['vtoonify_d_caricature/vtoonify_s068_d0.5.pt', 68],
-            'pixar': ['vtoonify_d_pixar/vtoonify_s052_d0.5.pt', 52],
-            'pixar-d': ['vtoonify_d_pixar/vtoonify_s_d.pt', 52],
-            'illustration1-d': ['vtoonify_d_illustration/vtoonify_s054_d_c.pt', 54],
-            'illustration2-d': ['vtoonify_d_illustration/vtoonify_s004_d_c.pt', 4],
-            'illustration3-d': ['vtoonify_d_illustration/vtoonify_s009_d_c.pt', 9],
-            'illustration4-d': ['vtoonify_d_illustration/vtoonify_s043_d_c.pt', 43],
-            'illustration5-d': ['vtoonify_d_illustration/vtoonify_s086_d_c.pt', 86],
+            # Add other styles as needed
         }
-        
+
         self.face_detector = self._create_insightface_detector()
         self.parsingpredictor = self._create_parsing_model()
-        self.pspencoder = self._load_encoder()    
+        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]),
         ])
-        
+
         self.vtoonify, self.exstyle = self._load_default_model()
         self.color_transfer = False
         self.style_name = 'cartoon1'
-        self.video_limit_cpu = 100
-        self.video_limit_gpu = 300
-        
+
     def _create_insightface_detector(self):
         # Initialize InsightFace
         app = insightface.app.FaceAnalysis()
         app.prepare(ctx_id=0 if self.device == 'cuda' else -1, det_size=(640, 640))
         return app
-    
+
     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[torch.Tensor, str]:
+
+    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'), 
@@ -97,8 +67,8 @@ class Model():
         with torch.no_grad():  
             exstyle = vtoonify.zplus2wplus(exstyle)
         return vtoonify, exstyle
-    
-    def load_model(self, style_type: str) -> tuple[torch.Tensor, str]:
+
+    def load_model(self, style_type: str) -> tuple:
         if 'illustration' in style_type:
             self.color_transfer = True
         else:
@@ -115,45 +85,79 @@ class Model():
         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
-            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]])
-            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)
-                if I is not None:
-                    I = self.transform(I).unsqueeze(dim=0).to(self.device)
+
+        # Use InsightFace for face detection
+        faces = self.face_detector.get(frame)
+        if len(faces) > 0:
+            logging.info(f"Detected {len(faces)} face(s).")
+            face = faces[0]
+            landmarks = face.landmark_2d_106
+
+            # Align face based on mapped landmarks
+            aligned_face = self.align_face(frame, landmarks)
+            if aligned_face is not None:
+                logging.info(f"Aligned face shape: {aligned_face.shape}")
+                with torch.no_grad():
+                    I = self.transform(aligned_face).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)
-                else:
-                    frame = np.zeros((256,256,3), np.uint8)
+                    message = 'Successfully aligned the face.'
+            else:
+                logging.warning("Failed to align face.")
+                frame = np.zeros((256, 256, 3), np.uint8)
         else:
-            frame = np.zeros((256,256,3), np.uint8)
+            logging.warning("No face detected.")
+            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
         return frame, instyle, message
-    
-    
 
-    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 align_face(self, image, landmarks):
+        # Map InsightFace landmarks to dlib's 68-point model
+        # Example: use specific indices for eyes and mouth
+        eye_left = np.mean(landmarks[36:42], axis=0)
+        eye_right = np.mean(landmarks[42:48], axis=0)
+        mouth_left = landmarks[48]
+        mouth_right = landmarks[54]
+
+        # Calculate transformation parameters
+        eye_center = (eye_left + eye_right) / 2
+        mouth_center = (mouth_left + mouth_right) / 2
+        eye_to_eye = eye_right - eye_left
+        eye_to_mouth = mouth_center - eye_center
+
+        # Define the transformation matrix
+        x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+        x /= np.hypot(*x)
+        x *= np.hypot(*eye_to_eye) * 2.0
+        y = np.flipud(x) * [-1, 1]
+        c = eye_center + eye_to_mouth * 0.1
+        quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+        qsize = np.hypot(*x) * 2
+
+        # Transform and crop the image
+        transform_size = 256
+        output_size = 256
+        img = Image.fromarray(image)
+        img = img.transform((transform_size, transform_size), Image.QUAD, (quad + 0.5).flatten(), Image.BILINEAR)
+        if output_size < transform_size:
+            img = img.resize((output_size, output_size), Image.ANTIALIAS)
+
+        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: