新智元报道
来源:GitHub
编辑:白峰
让马斯克加入你的Zoom会议
项目核心组件:一阶运动模型
如何在Zoom会议中使用替身
    
    
      
     
     
       def load_checkpoints(config_path, checkpoint_path, device='cuda'):
     
     
       
     
     
           with open(config_path) as f:
     
     
               config = yaml.load(f)
     
     
       
     
     
           generator = OcclusionAwareGenerator(**config['model_params']['generator_params'],
     
     
                                               **config['model_params']['common_params'])
     
     
           generator.to(device)
     
     
       
     
     
           kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
     
     
                                    **config['model_params']['common_params'])
     
     
           kp_detector.to(device)
     
     
       
     
     
           checkpoint = torch.load(checkpoint_path, map_location=device)
     
     
           generator.load_state_dict(checkpoint['generator'])
     
     
           kp_detector.load_state_dict(checkpoint['kp_detector'])
     
     
       
     
     
           generator = DataParallelWithCallback(generator)
     
     
           kp_detector = DataParallelWithCallback(kp_detector)
     
     
       
     
     
           generator.eval()
     
     
           kp_detector.eval()
     
     
       
     
     
           return generator, kp_detector
    
    
        
        
    
      
     
     
       def normalize_alignment_kp(kp):
     
     
           kp = kp - kp.mean(axis=0, keepdims=True)
     
     
           area = ConvexHull(kp[:, :2]).volume
     
     
           area = np.sqrt(area)
     
     
           kp[:, :2] = kp[:, :2] / area
     
     
           return kp
    
    
        
        
    
      
     
     
       def crop(img, p=0.7):
     
     
           h, w = img.shape[:2]
     
     
           x = int(min(w, h) * p)
     
     
           l = (w - x) // 2
     
     
           r = w - l
     
     
           u = (h - x) // 2
     
     
           d = h - u
     
     
           return img[u:d, l:r], (l,r,u,d)
     
     
       
     
     
       def pad_img(img, orig):
     
     
           h, w = orig.shape[:2]
     
     
           pad = int(256 * (w / h) - 256)
     
     
           out = np.pad(img, [[0,0], [pad//2, pad//2], [0,0]], 'constant')
     
     
           out = cv2.resize(out, (w, h))
     
     
           return out
    
    
        
        
    
      
     
     
       def predict(driving_frame, source_image, relative, adapt_movement_scale, fa, device='cuda'):
     
     
           global start_frame
     
     
           global start_frame_kp
     
     
           global kp_driving_initial
     
     
       
     
     
           with torch.no_grad():
     
     
               source = torch.tensor(source_image[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2).to(device)
     
     
               driving = torch.tensor(driving_frame[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2).to(device)
     
     
               kp_source = kp_detector(source)
     
     
       
     
     
               if kp_driving_initial is None:
     
     
                   kp_driving_initial = kp_detector(driving)
     
     
                   start_frame = driving_frame.copy()
     
     
                   start_frame_kp = get_frame_kp(fa, driving_frame)
     
     
       
     
     
               kp_driving = kp_detector(driving)
     
     
               kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving,
     
     
                                   kp_driving_initial=kp_driving_initial, use_relative_movement=relative,
     
     
                                   use_relative_jacobian=relative, adapt_movement_scale=adapt_movement_scale)
     
     
               out = generator(source, kp_source=kp_source, kp_driving=kp_norm)
     
     
       
     
     
               out = np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0]
     
     
               out = out[..., ::-1]
     
     
               out = (np.clip(out, 0, 1) * 255).astype(np.uint8)
     
     
       
     
     
               return out
    
    
        
    http://papers.nips.cc/paper/8935-first-order-motion-model-for-image-animation
https://aliaksandrsiarohin.github.io/first-order-model-website/
https://github.com/alievk/avatarify