Gradio interface

3DDFA_V2_cropping/recrop_images.py

@@ -23,6 +23,8 @@
 from utils.functions import draw_landmarks, get_suffix
 from utils.tddfa_util import str2bool
 
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
+
 def eg3dcamparams(R_in):
     camera_dist = 2.7
     intrinsics = np.array([[4.2647, 0, 0.5], [0, 4.2647, 0.5], [0, 0, 1]])
@@ -288,6 +290,7 @@ def main(args):
         # Save cropped images
         cropped_img = crop_final(img_orig, size=size, quad=quad)
         os.makedirs(args.out_dir, exist_ok=True)
+        #cv2.imwrite(os.path.join(args.out_dir, os.path.basename(img_path)), cropped_img)
         cv2.imwrite(os.path.join(args.out_dir, os.path.basename(img_path).replace(".png",".jpg")), cropped_img)
 
     # Save quads
@@ -298,7 +301,7 @@ def main(args):
     # Save meta data
     results_new = []
     for img, P  in results_meta.items():
-        img = os.path.basename(img)
+        img = os.path.basename(img).replace(".png",".jpg")        
         res = [format(r, '.6f') for r in P]
         results_new.append((img,res))
     with open(os.path.join(args.out_dir, args.output_json), 'w') as outfile:
@@ -312,7 +315,7 @@ def main(args):
     parser.add_argument('-j', '--output_json', type=str, default='dataset.json')
     parser.add_argument('-p', '--prefix', type=str, default='')
     parser.add_argument('--size', type=int, default=1024)
-    parser.add_argument('--out_dir', type=str, default='./crop_samples/img')
+    parser.add_argument('--out_dir', type=str, default='./crop_samples')
     parser.add_argument('--mode', type=str, default='gpu', help='gpu or cpu mode')
     parser.add_argument('--config', type=str, default='configs/mb1_120x120.yml')
     parser.add_argument('--individual', action='store_true', default=False)

README.md

@@ -68,6 +68,13 @@ python gen_interpolation.py --network models/easy-khair-180-gpc0.8-trans10-02500
         --trunc 0.7 --outdir interpolation_out
 ```
 
+## Gradio demo
+
+```.bash
+#path_3DDFA = "E:/3DDFA_V2-master/" at line 15 should be modified to fit your 3DDFA folder setup
+python gradiodemo.py
+
+```
 
 
 ## Using networks from Python
@@ -170,4 +177,4 @@ This is a research reference implementation and is treated as a one-time code dr
 
 We thank Shuhong Chen for the discussion during Sizhe's internship.
 
-This repo is heavily based off the [NVlabs/eg3d](https://github.com/NVlabs/eg3d) repo; Huge thanks to the EG3D authors for releasing their code!
+This repo is heavily based off the [NVlabs/eg3d](https://github.com/NVlabs/eg3d) repo; Huge thanks to the EG3D authors for releasing their code!

gen_pti_script_with_mesh_noSeg.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+
+model="easy-khair-180-gpc0.8-trans10-025000.pkl"
+
+input_dir="models"
+output_dir="pti_out"
+
+target_img="dataset/testdata_img"
+#target_seg="dataset/testdata_seg"
+
+
+# Perform the pti and save w
+python projector_withseg.py --outdir="${output_dir}" --target_img="${target_img}" --network "${input_dir}/${model}" --idx "0" --shapes=True --save-video=False
+# Generate .mp4 before finetune
+python gen_videos_proj_withseg.py --output="${output_dir}/${model}/0/PTI_render/pre.mp4" --latent="${output_dir}/${model}/0/projected_w.npz" --trunc 0.7 --network "${input_dir}/${model}" --cfg Head
+# Generate .mp4, .ply mesh and frame images after finetune
+    python gen_videos_proj_withseg.py --output="${output_dir}/${model}/0/PTI_render/post.mp4" --latent="${output_dir}/${model}/0/projected_w.npz" --trunc 0.7 --network "${output_dir}/${model}/0/fintuned_generator.pkl" --cfg Head --shapes True --frames True --level 42
+
+done

gen_videos_proj_withseg.py

@@ -1,8 +1,14 @@
-''' Generate videos using pretrained network pickle.
-Code adapted from following paper
-"Efficient Geometry-aware 3D Generative Adversarial Networks."
-See LICENSES/LICENSE_EG3D for original license.
-'''
+# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+"""Generate lerp videos using pretrained network pickle."""
 
 import os
 import re
@@ -16,13 +22,38 @@
 import torch
 from tqdm import tqdm
 import mrcfile
+from PIL import Image
 
 import legacy
 
 from camera_utils import LookAtPoseSampler
 from torch_utils import misc
+
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
 #----------------------------------------------------------------------------
 
+def save_image(img, name, output_dir):
+    os.makedirs(output_dir, exist_ok=True)  # Create the output directory if it doesn't exist
+
+    image_path = os.path.join(output_dir, f"video_frame_{name}.png")  # Define the path and filename for the image
+
+    img = (img + 1) / 2  # Convert the image from [-1, 1] range to [0, 1] range
+    img = img.permute(1, 2, 0)  # Rearrange the dimensions of the image tensor
+    img = (img * 255).clamp(0, 255).byte()  # Scale the image values to [0, 255] and convert to byte tensor
+
+    if img.shape[-1] == 1:  # Grayscale image with single channel
+        img = img.squeeze(dim=1)  # Remove the single channel channel dimension
+        PIL_image = Image.fromarray(img.cpu().numpy(), mode='L')  # Convert to PIL image with 'L' mode (grayscale)
+    elif img.shape[-1] == 3:  # RGB image with three channels
+        PIL_image = Image.fromarray(img.cpu().numpy(), mode='RGB')  # Convert to PIL image with 'RGB' mode
+    else:
+        # For images with more than 3 channels, convert to RGBA format by adding an alpha channel
+        img = torch.cat((img, torch.full_like(img[..., :1], 255, dtype=torch.uint8)), dim=-1)
+        PIL_image = Image.fromarray(img.cpu().numpy(), mode='RGBA')  # Convert to PIL image with 'RGBA' mode
+
+    PIL_image.save(image_path)  # Save the PIL image to disk
+
+
 def layout_grid(img, grid_w=None, grid_h=1, float_to_uint8=True, chw_to_hwc=True, to_numpy=True):
     batch_size, channels, img_h, img_w = img.shape
     if grid_w is None:
@@ -65,7 +96,7 @@ def create_samples(N=256, voxel_origin=[0, 0, 0], cube_length=2.0):
 
 #----------------------------------------------------------------------------
 
-def gen_interp_video(G, mp4: str, ws, w_frames=60*4, kind='cubic', grid_dims=(1,1), num_keyframes=None, wraps=2, psi=1, truncation_cutoff=14, cfg='FFHQ', image_mode='image', gen_shapes=False, device=torch.device('cuda'), **video_kwargs):
+def gen_interp_video(G, mp4: str, ws, w_frames=60*4, kind='cubic', grid_dims=(1,1), num_keyframes=None, wraps=2, psi=1, truncation_cutoff=14, cfg='FFHQ', image_mode='image', gen_shapes=False, gen_frames=False, iso_level=10.0, device=torch.device('cuda'), **video_kwargs):
     grid_w = grid_dims[0]
     grid_h = grid_dims[1]
 
@@ -148,6 +179,9 @@ def gen_interp_video(G, mp4: str, ws, w_frames=60*4, kind='cubic', grid_dims=(1,
 
                 imgs.append(img)
 
+                if gen_frames:
+                    save_image(img, frame_idx, mp4.replace('.mp4', '/'))
+
                 if gen_shapes and frame_idx == 0:
                     # generate shapes
                     print('Generating shape for frame %d / %d ...' % (frame_idx, num_keyframes * w_frames))
@@ -165,6 +199,12 @@ def gen_interp_video(G, mp4: str, ws, w_frames=60*4, kind='cubic', grid_dims=(1,
                                 torch.manual_seed(0)
                                 sigma = G.sample_mixed(samples[:, head:head+max_batch], transformed_ray_directions_expanded[:, :samples.shape[1]-head], w.unsqueeze(0), truncation_psi=psi, noise_mode='const')['sigma']
                                 sigmas[:, head:head+max_batch] = sigma
+                                '''
+                                sample_result = G.sample_mixed(samples[:, head:head+max_batch], transformed_ray_directions_expanded[:, :samples.shape[1]-head], w.unsqueeze(0), truncation_psi=psi, noise_mode='const')
+                                sigmas[:, head:head+max_batch] = sample_result['sigma']
+                                color_batch = G.torgb(sample_result['rgb'].transpose(1,2)[...,None], ws[0,0,0,:1])
+                                colors[:, head:head+max_batch] = np.transpose(color_batch[...,0], (2, 1, 0))
+                                '''
                                 head += max_batch
                                 pbar.update(max_batch)
 
@@ -180,10 +220,10 @@ def gen_interp_video(G, mp4: str, ws, w_frames=60*4, kind='cubic', grid_dims=(1,
                     sigmas[:, :, :pad] = 0
                     sigmas[:, :, -pad:] = 0
 
-                    output_ply = False
+                    output_ply = True
                     if output_ply:
                         from shape_utils import convert_sdf_samples_to_ply
-                        convert_sdf_samples_to_ply(np.transpose(sigmas, (2, 1, 0)), [0, 0, 0], 1, os.path.join(outdirs, mp4.replace('.mp4', '.ply')), level=10)
+                        convert_sdf_samples_to_ply(np.transpose(sigmas, (2, 1, 0)), [0, 0, 0], 1, mp4.replace('.mp4', '.ply'), level=iso_level)
                     else: # output mrc
                         with mrcfile.new_mmap(mp4.replace('.mp4', '.mrc'), overwrite=True, shape=sigmas.shape, mrc_mode=2) as mrc:
                             mrc.data[:] = sigmas
@@ -247,6 +287,8 @@ def parse_tuple(s: Union[str, Tuple[int,int]]) -> Tuple[int, int]:
 @click.option('--sample_mult', 'sampling_multiplier', type=float, help='Multiplier for depth sampling in volume rendering', default=1, show_default=True)
 @click.option('--nrr', type=int, help='Neural rendering resolution override', default=None, show_default=True)
 @click.option('--shapes', type=bool, help='Gen shapes for shape interpolation', default=False, show_default=True)
+@click.option('--level', type=float, help='Iso surface level for mesh generation', default=10, show_default=True)
+@click.option('--frames', type=bool, help='Save frames as images', default=False, show_default=True)
 @click.option('--interpolate', type=bool, help='Interpolate between seeds', default=True, show_default=True)
 
 def generate_images(
@@ -264,14 +306,16 @@ def generate_images(
     sampling_multiplier: float,
     nrr: Optional[int],
     shapes: bool,
+    level: float,
+    frames: bool,
     interpolate: bool,
 ):
     """Render a latent vector interpolation video.
 
     """
 
     print('Loading networks from "%s"...' % network_pkl)
-    device = torch.device('cuda:1')
+    device = torch.device('cuda')
     with dnnlib.util.open_url(network_pkl) as f:
         G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
 
@@ -287,7 +331,7 @@ def generate_images(
         truncation_cutoff = 14 # no truncation so doesn't matter where we cutoff
 
     ws = torch.tensor(np.load(latent)['w']).to(device)
-    gen_interp_video(G=G, mp4=output, ws=ws, bitrate='100M', grid_dims=grid, num_keyframes=num_keyframes, w_frames=w_frames, psi=truncation_psi, truncation_cutoff=truncation_cutoff, cfg=cfg, image_mode=image_mode, gen_shapes=shapes, device=device)
+    gen_interp_video(G=G, mp4=output, ws=ws, bitrate='100M', grid_dims=grid, num_keyframes=num_keyframes, w_frames=w_frames, psi=truncation_psi, truncation_cutoff=truncation_cutoff, cfg=cfg, image_mode=image_mode, gen_shapes=shapes, iso_level=level, gen_frames=frames, device=device)
 
 #----------------------------------------------------------------------------
 

gradiodemo.py

@@ -0,0 +1,90 @@
+from codecs import ignore_errors
+from genericpath import isdir
+import sys
+from subprocess import call
+import os
+import torch
+import random
+import string
+
+import shutil
+
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
+
+
+path_3DDFA = "E:/3DDFA_V2-master/"
+path_crop_targets = "test/original/"
+path_crop_results = "crop_samples/"
+path_target_img = "dataset/testdata_img/"
+
+def run_cmd(command):
+    try:
+        print(command)
+        call(command, shell=True)
+    except Exception as e:
+        print(f"Error: {e}!")
+
+
+import gradio as gr
+
+
+
+def inference (img):    
+    filename = os.path.basename(img)    
+    serialname = f''.join(random.choices(string.digits, k=6))
+    os.chdir(path_3DDFA)
+
+    #flush
+    shutil.rmtree(path_crop_results, ignore_errors= True)
+    shutil.rmtree(path_3DDFA + path_crop_targets, ignore_errors= True)    
+    os.makedirs(f'./' + path_crop_targets, exist_ok = True)
+    #copy img
+    shutil.copyfile(img, path_3DDFA + path_crop_targets + serialname + os.path.splitext(filename)[1])
+
+    run_cmd(f'python dlib_kps.py')
+    run_cmd(f'python recrop_images.py -i data.pkl -j dataset.json')
+    if os.path.isfile( path_3DDFA + path_crop_results + serialname + f'.jpg' ) == False:
+        print("Error: no face")
+        return
+    os.chdir(os.path.dirname(__file__)) #return to python file path
+    #os.chdir(os.path.dirname(os.getcwd())) #return to root    
+
+
+    #flush2
+    shutil.rmtree(path_target_img, ignore_errors= True)
+    #os.makedirs(f'./' + path_target_img, exist_ok = True)
+    #copy2
+    shutil.copytree(path_3DDFA + path_crop_results , path_target_img)
+
+    run_cmd(f'gen_pti_script_with_mesh_noSeg.sh')
+
+    #store .ply and post.mp4 separately under a new folder
+    if os.path.isdir('out') == False:
+        os.mkdir('out')
+    os.mkdir(f'out/' + serialname)
+
+    shutil.copyfile('pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/PTI_render/post.mp4', f'out/' + serialname + '/post.mp4')
+    shutil.copyfile('pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/PTI_render/post.ply', f'out/' + serialname + '/post.ply')
+    shutil.copyfile(path_target_img + serialname + '.jpg', f'out/' + serialname + '/'+ serialname + '.jpg' )
+    shutil.copyfile(path_target_img + 'dataset.json', f'out/' + serialname + '/dataset.json' )
+
+    return f'pti_out/easy-khair-180-gpc0.8-trans10-025000.pkl/0/PTI_render/post.mp4'
+
+
+title= "Panohead demo"
+description= "Panohead demo for gradio"
+article= "<p style='text-align: center'><a href='https://github.com/SizheAn/PanoHead'>Github Repo</a></p>"
+examples= [       
+]
+
+demo = gr.Interface(
+     inference,     
+     gr.Image(type="filepath"), 
+     outputs=gr.Video(label="Out"),
+    title=title,
+    description=description,
+    article=article,
+    examples=examples
+)
+
+demo.launch()

projector_withseg.py

@@ -1,4 +1,12 @@
-""" Projecting input images into latent spaces. """
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Project given image to the latent space of pretrained network pickle."""
 
 import copy
 import os
@@ -19,6 +27,8 @@
 
 from camera_utils import LookAtPoseSampler
 
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
+
 def create_samples(N=256, voxel_origin=[0, 0, 0], cube_length=2.0):
     # NOTE: the voxel_origin is actually the (bottom, left, down) corner, not the middle
     voxel_origin = np.array(voxel_origin) - cube_length/2
@@ -430,4 +440,4 @@ def run_projection(
 if __name__ == "__main__":
     run_projection() # pylint: disable=no-value-for-parameter
 
-#----------------------------------------------------------------------------
+#----------------------------------------------------------------------------