afruehstueck / vive3d Goto Github PK

    def generate(self, w, yaw=0.0, pitch=0.0, focal_length=None, noise_mode='const', output_all=False, grad=False, w_offset=None,):                 
        # convert angles to camera parameter(s)
        camera_params = self.get_camera_parameters(yaw=yaw, pitch=pitch, focal_length=focal_length)

        # convert w to correct size
        # (,512) -> (1, 512)
        if len(w.shape) == 1:
            w = w.unsqueeze(0)
        
        # (x, 512) -> (x, 1, 512)
        if len(w.shape) == 2:
            w = w.unsqueeze(1)
            
        # repeat to (1, 14, 512)
        if w.shape[1] == 1:
            w = w.repeat(1, self.active_G.backbone.synthesis.num_ws, 1)
        
        num_images = max(len(w), len(camera_params))
        
        # replicate if single w was passed
        if w.shape[0] < num_images:
            w = w.repeat(num_images, 1, 1)
        
        # replicate if single camera parameter was passed
        if camera_params.shape[0] < num_images:
            camera_params = camera_params.repeat(num_images, 1)

        if w_offset is not None:
            w = w + w_offset
        
        assert w.shape[0] == camera_params.shape[0], f'incompatible size of w ({w.shape}) and camera params ({camera_params.shape})'
        
        images = []
        images_raw = []
        depths = []
        self.active_G.to(self.device)