Create vtoonify/style_transfer.py

vtoonify/style_transfer.py

@@ -0,0 +1,232 @@
+import os
+#os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+import argparse
+import numpy as np
+import cv2
+import dlib
+import torch
+from torchvision import transforms
+import torch.nn.functional as F
+from tqdm import tqdm
+from model.vtoonify import VToonify
+from model.bisenet.model import BiSeNet
+from model.encoder.align_all_parallel import align_face
+from util import save_image, load_image, visualize, load_psp_standalone, get_video_crop_parameter, tensor2cv2
+
+
+class TestOptions():
+    def __init__(self):
+
+        self.parser = argparse.ArgumentParser(description="Style Transfer")
+        self.parser.add_argument("--content", type=str, default='./data/077436.jpg', help="path of the content image/video")
+        self.parser.add_argument("--style_id", type=int, default=26, help="the id of the style image")
+        self.parser.add_argument("--style_degree", type=float, default=0.5, help="style degree for VToonify-D")
+        self.parser.add_argument("--color_transfer", action="store_true", help="transfer the color of the style")
+        self.parser.add_argument("--ckpt", type=str, default='./checkpoint/vtoonify_d_cartoon/vtoonify_s_d.pt', help="path of the saved model")
+        self.parser.add_argument("--output_path", type=str, default='./output/', help="path of the output images")
+        self.parser.add_argument("--scale_image", action="store_true", help="resize and crop the image to best fit the model")
+        self.parser.add_argument("--style_encoder_path", type=str, default='./checkpoint/encoder.pt', help="path of the style encoder")
+        self.parser.add_argument("--exstyle_path", type=str, default=None, help="path of the extrinsic style code")
+        self.parser.add_argument("--faceparsing_path", type=str, default='./checkpoint/faceparsing.pth', help="path of the face parsing model")
+        self.parser.add_argument("--video", action="store_true", help="if true, video stylization; if false, image stylization")
+        self.parser.add_argument("--cpu", action="store_true", help="if true, only use cpu")
+        self.parser.add_argument("--backbone", type=str, default='dualstylegan', help="dualstylegan | toonify")
+        self.parser.add_argument("--padding", type=int, nargs=4, default=[200,200,200,200], help="left, right, top, bottom paddings to the face center")
+        self.parser.add_argument("--batch_size", type=int, default=4, help="batch size of frames when processing video")
+        self.parser.add_argument("--parsing_map_path", type=str, default=None, help="path of the refined parsing map of the target video")
+        
+    def parse(self):
+        self.opt = self.parser.parse_args()
+        if self.opt.exstyle_path is None:
+            self.opt.exstyle_path = os.path.join(os.path.dirname(self.opt.ckpt), 'exstyle_code.npy')
+        args = vars(self.opt)
+        print('Load options')
+        for name, value in sorted(args.items()):
+            print('%s: %s' % (str(name), str(value)))
+        return self.opt
+    
+if __name__ == "__main__":
+
+    parser = TestOptions()
+    args = parser.parse()
+    print('*'*98)
+    
+    
+    device = "cpu" if args.cpu else "cuda"
+    
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5,0.5,0.5]),
+        ])
+    
+    vtoonify = VToonify(backbone = args.backbone)
+    vtoonify.load_state_dict(torch.load(args.ckpt, map_location=lambda storage, loc: storage)['g_ema'])
+    vtoonify.to(device)
+
+    parsingpredictor = BiSeNet(n_classes=19)
+    parsingpredictor.load_state_dict(torch.load(args.faceparsing_path, map_location=lambda storage, loc: storage))
+    parsingpredictor.to(device).eval()
+
+    modelname = './checkpoint/shape_predictor_68_face_landmarks.dat'
+    if not os.path.exists(modelname):
+        import wget, bz2
+        wget.download('http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2', modelname+'.bz2')
+        zipfile = bz2.BZ2File(modelname+'.bz2')
+        data = zipfile.read()
+        open(modelname, 'wb').write(data) 
+    landmarkpredictor = dlib.shape_predictor(modelname)
+
+    pspencoder = load_psp_standalone(args.style_encoder_path, device)    
+
+    if args.backbone == 'dualstylegan':
+        exstyles = np.load(args.exstyle_path, allow_pickle='TRUE').item()
+        stylename = list(exstyles.keys())[args.style_id]
+        exstyle = torch.tensor(exstyles[stylename]).to(device)
+        with torch.no_grad():  
+            exstyle = vtoonify.zplus2wplus(exstyle)
+
+    if args.video and args.parsing_map_path is not None:
+        x_p_hat = torch.tensor(np.load(args.parsing_map_path))          
+            
+    print('Load models successfully!')
+    
+    
+    filename = args.content
+    basename = os.path.basename(filename).split('.')[0]
+    scale = 1
+    kernel_1d = np.array([[0.125],[0.375],[0.375],[0.125]])
+    print('Processing ' + os.path.basename(filename) + ' with vtoonify_' + args.backbone[0])
+    if args.video:
+        cropname = os.path.join(args.output_path, basename + '_input.mp4')
+        savename = os.path.join(args.output_path, basename + '_vtoonify_' +  args.backbone[0] + '.mp4')
+
+        video_cap = cv2.VideoCapture(filename)
+        num = int(video_cap.get(7))
+
+        first_valid_frame = True
+        batch_frames = []
+        for i in tqdm(range(num)):
+            success, frame = video_cap.read()
+            if success == False:
+                assert('load video frames error')
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            # We proprocess the video by detecting the face in the first frame, 
+            # and resizing the frame so that the eye distance is 64 pixels.
+            # Centered on the eyes, we crop the first frame to almost 400x400 (based on args.padding).
+            # All other frames use the same resizing and cropping parameters as the first frame.
+            if first_valid_frame:
+                if args.scale_image:
+                    paras = get_video_crop_parameter(frame, landmarkpredictor, args.padding)
+                    if paras is None:
+                        continue
+                    h,w,top,bottom,left,right,scale = paras
+                    H, W = int(bottom-top), int(right-left)
+                    # for HR video, we apply gaussian blur to the frames to avoid flickers caused by bilinear downsampling
+                    # this can also prevent over-sharp stylization results. 
+                    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]
+                else:
+                    H, W = frame.shape[0], frame.shape[1]
+
+                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+                videoWriter = cv2.VideoWriter(cropname, fourcc, video_cap.get(5), (W, H))
+                videoWriter2 = cv2.VideoWriter(savename, fourcc, video_cap.get(5), (4*W, 4*H))
+                
+                # For each video, we detect and align the face in the first frame for pSp to obtain the style code. 
+                # This style code is used for all other frames.
+                with torch.no_grad():
+                    I = align_face(frame, landmarkpredictor)
+                    I = transform(I).unsqueeze(dim=0).to(device)
+                    s_w = pspencoder(I)
+                    s_w = vtoonify.zplus2wplus(s_w)
+                    if vtoonify.backbone == 'dualstylegan':
+                        if args.color_transfer:
+                            s_w = exstyle
+                        else:
+                            s_w[:,:7] = exstyle[:,:7]
+                first_valid_frame = False
+            elif args.scale_image:
+                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]
+
+            videoWriter.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+
+            batch_frames += [transform(frame).unsqueeze(dim=0).to(device)]
+
+            if len(batch_frames) == args.batch_size or (i+1) == num:
+                x = torch.cat(batch_frames, dim=0)
+                batch_frames = []
+                with torch.no_grad():
+                    # parsing network works best on 512x512 images, so we predict parsing maps on upsmapled frames
+                    # followed by downsampling the parsing maps
+                    if args.video and args.parsing_map_path is not None:
+                        x_p = x_p_hat[i+1-x.size(0):i+1].to(device)
+                    else:
+                        x_p = F.interpolate(parsingpredictor(2*(F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0], 
+                                        scale_factor=0.5, recompute_scale_factor=False).detach()
+                    # we give parsing maps lower weight (1/16)
+                    inputs = torch.cat((x, x_p/16.), dim=1)
+                    # d_s has no effect when backbone is toonify
+                    y_tilde = vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s = args.style_degree)       
+                    y_tilde = torch.clamp(y_tilde, -1, 1)
+                for k in range(y_tilde.size(0)):
+                    videoWriter2.write(tensor2cv2(y_tilde[k].cpu()))
+
+        videoWriter.release()
+        videoWriter2.release()
+        video_cap.release()
+
+    
+    else:
+        cropname = os.path.join(args.output_path, basename + '_input.jpg')
+        savename = os.path.join(args.output_path, basename + '_vtoonify_' +  args.backbone[0] + '.jpg')
+
+        frame = cv2.imread(filename)
+        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+
+        # We detect the face in the image, and resize the image so that the eye distance is 64 pixels.
+        # Centered on the eyes, we crop the image to almost 400x400 (based on args.padding).
+        if args.scale_image:
+            paras = get_video_crop_parameter(frame, landmarkpredictor, args.padding)
+            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
+                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, landmarkpredictor)
+            I = transform(I).unsqueeze(dim=0).to(device)
+            s_w = pspencoder(I)
+            s_w = vtoonify.zplus2wplus(s_w)
+            if vtoonify.backbone == 'dualstylegan':
+                if args.color_transfer:
+                    s_w = exstyle
+                else:
+                    s_w[:,:7] = exstyle[:,:7]
+
+            x = transform(frame).unsqueeze(dim=0).to(device)
+            # parsing network works best on 512x512 images, so we predict parsing maps on upsmapled frames
+            # followed by downsampling the parsing maps
+            x_p = F.interpolate(parsingpredictor(2*(F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0], 
+                                scale_factor=0.5, recompute_scale_factor=False).detach()
+            # we give parsing maps lower weight (1/16)
+            inputs = torch.cat((x, x_p/16.), dim=1)
+            # d_s has no effect when backbone is toonify
+            y_tilde = vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s = args.style_degree)        
+            y_tilde = torch.clamp(y_tilde, -1, 1)
+
+        cv2.imwrite(cropname, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+        save_image(y_tilde[0].cpu(), savename)
+        
+    print('Transfer style successfully!')