diff --git a/3DDFA_V2_cropping/recrop_images.py b/3DDFA_V2_cropping/recrop_images.py index 4b157fb..8066cbb 100644 --- a/3DDFA_V2_cropping/recrop_images.py +++ b/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) diff --git a/README.md b/README.md index a019f96..b09c0e9 100644 --- a/README.md +++ b/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! \ No newline at end of file +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! diff --git a/gen_pti_script_with_mesh_noSeg.sh b/gen_pti_script_with_mesh_noSeg.sh new file mode 100644 index 0000000..202ba03 --- /dev/null +++ b/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 diff --git a/gen_videos_proj_withseg.py b/gen_videos_proj_withseg.py index c2bc364..cfa3715 100644 --- a/gen_videos_proj_withseg.py +++ b/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,6 +306,8 @@ def generate_images( sampling_multiplier: float, nrr: Optional[int], shapes: bool, + level: float, + frames: bool, interpolate: bool, ): """Render a latent vector interpolation video. @@ -271,7 +315,7 @@ def generate_images( """ 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) #---------------------------------------------------------------------------- diff --git a/gradiodemo.py b/gradiodemo.py new file mode 100644 index 0000000..083a1f6 --- /dev/null +++ b/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= "

Github Repo

" +examples= [ +] + +demo = gr.Interface( + inference, + gr.Image(type="filepath"), + outputs=gr.Video(label="Out"), + title=title, + description=description, + article=article, + examples=examples +) + +demo.launch() diff --git a/projector_withseg.py b/projector_withseg.py index 2480310..679bd99 100644 --- a/projector_withseg.py +++ b/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 -#---------------------------------------------------------------------------- \ No newline at end of file +#----------------------------------------------------------------------------