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!')