Spaces:
Runtime error
Runtime error
| import sys | |
| import os | |
| os.system("git clone https://github.com/royorel/StyleSDF.git") | |
| sys.path.append("StyleSDF") | |
| os.system(f"{sys.executable} -m pip install -U fvcore") | |
| import torch | |
| pyt_version_str=torch.__version__.split("+")[0].replace(".", "") | |
| version_str="".join([ | |
| f"py3{sys.version_info.minor}_cu", | |
| torch.version.cuda.replace(".",""), | |
| f"_pyt{pyt_version_str}" | |
| ]) | |
| os.system(f"{sys.executable} -m pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html") | |
| from download_models import download_pretrained_models | |
| download_pretrained_models() | |
| import torch | |
| import trimesh | |
| import numpy as np | |
| from munch import * | |
| from PIL import Image | |
| from tqdm import tqdm | |
| from torch.nn import functional as F | |
| from torch.utils import data | |
| from torchvision import utils | |
| from torchvision import transforms | |
| from skimage.measure import marching_cubes | |
| from scipy.spatial import Delaunay | |
| from options import BaseOptions | |
| from model import Generator | |
| from utils import ( | |
| generate_camera_params, | |
| align_volume, | |
| extract_mesh_with_marching_cubes, | |
| xyz2mesh, | |
| ) | |
| from utils import ( | |
| generate_camera_params, align_volume, extract_mesh_with_marching_cubes, | |
| xyz2mesh, create_cameras, create_mesh_renderer, add_textures, | |
| ) | |
| from pytorch3d.structures import Meshes | |
| from pdb import set_trace as st | |
| import skvideo.io | |
| def generate(opt, g_ema, surface_g_ema, device, mean_latent, surface_mean_latent): | |
| g_ema.eval() | |
| if not opt.no_surface_renderings: | |
| surface_g_ema.eval() | |
| # set camera angles | |
| if opt.fixed_camera_angles: | |
| # These can be changed to any other specific viewpoints. | |
| # You can add or remove viewpoints as you wish | |
| locations = torch.tensor([[0, 0], | |
| [-1.5 * opt.camera.azim, 0], | |
| [-1 * opt.camera.azim, 0], | |
| [-0.5 * opt.camera.azim, 0], | |
| [0.5 * opt.camera.azim, 0], | |
| [1 * opt.camera.azim, 0], | |
| [1.5 * opt.camera.azim, 0], | |
| [0, -1.5 * opt.camera.elev], | |
| [0, -1 * opt.camera.elev], | |
| [0, -0.5 * opt.camera.elev], | |
| [0, 0.5 * opt.camera.elev], | |
| [0, 1 * opt.camera.elev], | |
| [0, 1.5 * opt.camera.elev]], device=device) | |
| # For zooming in/out change the values of fov | |
| # (This can be defined for each view separately via a custom tensor | |
| # like the locations tensor above. Tensor shape should be [locations.shape[0],1]) | |
| # reasonable values are [0.75 * opt.camera.fov, 1.25 * opt.camera.fov] | |
| fov = opt.camera.fov * torch.ones((locations.shape[0],1), device=device) | |
| num_viewdirs = locations.shape[0] | |
| else: # draw random camera angles | |
| locations = None | |
| # fov = None | |
| fov = opt.camera.fov | |
| num_viewdirs = opt.num_views_per_id | |
| # generate images | |
| for i in tqdm(range(opt.identities)): | |
| with torch.no_grad(): | |
| chunk = 8 | |
| sample_z = torch.randn(1, opt.style_dim, device=device).repeat(num_viewdirs,1) | |
| sample_cam_extrinsics, sample_focals, sample_near, sample_far, sample_locations = \ | |
| generate_camera_params(opt.renderer_output_size, device, batch=num_viewdirs, | |
| locations=locations, #input_fov=fov, | |
| uniform=opt.camera.uniform, azim_range=opt.camera.azim, | |
| elev_range=opt.camera.elev, fov_ang=fov, | |
| dist_radius=opt.camera.dist_radius) | |
| rgb_images = torch.Tensor(0, 3, opt.size, opt.size) | |
| rgb_images_thumbs = torch.Tensor(0, 3, opt.renderer_output_size, opt.renderer_output_size) | |
| for j in range(0, num_viewdirs, chunk): | |
| out = g_ema([sample_z[j:j+chunk]], | |
| sample_cam_extrinsics[j:j+chunk], | |
| sample_focals[j:j+chunk], | |
| sample_near[j:j+chunk], | |
| sample_far[j:j+chunk], | |
| truncation=opt.truncation_ratio, | |
| truncation_latent=mean_latent) | |
| rgb_images = torch.cat([rgb_images, out[0].cpu()], 0) | |
| rgb_images_thumbs = torch.cat([rgb_images_thumbs, out[1].cpu()], 0) | |
| utils.save_image(rgb_images, | |
| os.path.join(opt.results_dst_dir, 'images','{}.png'.format(str(i).zfill(7))), | |
| nrow=num_viewdirs, | |
| normalize=True, | |
| padding=0, | |
| value_range=(-1, 1),) | |
| utils.save_image(rgb_images_thumbs, | |
| os.path.join(opt.results_dst_dir, 'images','{}_thumb.png'.format(str(i).zfill(7))), | |
| nrow=num_viewdirs, | |
| normalize=True, | |
| padding=0, | |
| value_range=(-1, 1),) | |
| # this is done to fit to RTX2080 RAM size (11GB) | |
| del out | |
| torch.cuda.empty_cache() | |
| if not opt.no_surface_renderings: | |
| surface_chunk = 1 | |
| scale = surface_g_ema.renderer.out_im_res / g_ema.renderer.out_im_res | |
| surface_sample_focals = sample_focals * scale | |
| for j in range(0, num_viewdirs, surface_chunk): | |
| surface_out = surface_g_ema([sample_z[j:j+surface_chunk]], | |
| sample_cam_extrinsics[j:j+surface_chunk], | |
| surface_sample_focals[j:j+surface_chunk], | |
| sample_near[j:j+surface_chunk], | |
| sample_far[j:j+surface_chunk], | |
| truncation=opt.truncation_ratio, | |
| truncation_latent=surface_mean_latent, | |
| return_sdf=True, | |
| return_xyz=True) | |
| xyz = surface_out[2].cpu() | |
| sdf = surface_out[3].cpu() | |
| # this is done to fit to RTX2080 RAM size (11GB) | |
| del surface_out | |
| torch.cuda.empty_cache() | |
| # mesh extractions are done one at a time | |
| for k in range(surface_chunk): | |
| curr_locations = sample_locations[j:j+surface_chunk] | |
| loc_str = '_azim{}_elev{}'.format(int(curr_locations[k,0] * 180 / np.pi), | |
| int(curr_locations[k,1] * 180 / np.pi)) | |
| # Save depth outputs as meshes | |
| depth_mesh_filename = os.path.join(opt.results_dst_dir,'depth_map_meshes','sample_{}_depth_mesh{}.obj'.format(i, loc_str)) | |
| depth_mesh = xyz2mesh(xyz[k:k+surface_chunk]) | |
| if depth_mesh != None: | |
| with open(depth_mesh_filename, 'w') as f: | |
| depth_mesh.export(f,file_type='obj') | |
| # extract full geometry with marching cubes | |
| if j == 0: | |
| try: | |
| frostum_aligned_sdf = align_volume(sdf) | |
| marching_cubes_mesh = extract_mesh_with_marching_cubes(frostum_aligned_sdf[k:k+surface_chunk]) | |
| except ValueError: | |
| marching_cubes_mesh = None | |
| print('Marching cubes extraction failed.') | |
| print('Please check whether the SDF values are all larger (or all smaller) than 0.') | |
| return depth_mesh,marching_cubes_mesh | |
| # User options | |
| def get_generate_vars(model_type): | |
| opt = BaseOptions().parse() | |
| opt.camera.uniform = True | |
| opt.model.is_test = True | |
| opt.model.freeze_renderer = False | |
| opt.rendering.offset_sampling = True | |
| opt.rendering.static_viewdirs = True | |
| opt.rendering.force_background = True | |
| opt.rendering.perturb = 0 | |
| opt.inference.renderer_output_size = opt.model.renderer_spatial_output_dim | |
| opt.inference.style_dim = opt.model.style_dim | |
| opt.inference.project_noise = opt.model.project_noise | |
| # User options | |
| opt.inference.no_surface_renderings = False # When true, only RGB images will be created | |
| opt.inference.fixed_camera_angles = False # When true, each identity will be rendered from a specific set of 13 viewpoints. Otherwise, random views are generated | |
| opt.inference.identities = 1 # Number of identities to generate | |
| opt.inference.num_views_per_id = 1 # Number of viewpoints generated per identity. This option is ignored if opt.inference.fixed_camera_angles is true. | |
| opt.inference.camera = opt.camera | |
| # Load saved model | |
| if model_type == 'ffhq': | |
| model_path = 'ffhq1024x1024.pt' | |
| opt.model.size = 1024 | |
| opt.experiment.expname = 'ffhq1024x1024' | |
| else: | |
| opt.inference.camera.azim = 0.15 | |
| model_path = 'afhq512x512.pt' | |
| opt.model.size = 512 | |
| opt.experiment.expname = 'afhq512x512' | |
| # Create results directory | |
| result_model_dir = 'final_model' | |
| results_dir_basename = os.path.join(opt.inference.results_dir, opt.experiment.expname) | |
| opt.inference.results_dst_dir = os.path.join(results_dir_basename, result_model_dir) | |
| if opt.inference.fixed_camera_angles: | |
| opt.inference.results_dst_dir = os.path.join(opt.inference.results_dst_dir, 'fixed_angles') | |
| else: | |
| opt.inference.results_dst_dir = os.path.join(opt.inference.results_dst_dir, 'random_angles') | |
| os.makedirs(opt.inference.results_dst_dir, exist_ok=True) | |
| os.makedirs(os.path.join(opt.inference.results_dst_dir, 'images'), exist_ok=True) | |
| if not opt.inference.no_surface_renderings: | |
| os.makedirs(os.path.join(opt.inference.results_dst_dir, 'depth_map_meshes'), exist_ok=True) | |
| os.makedirs(os.path.join(opt.inference.results_dst_dir, 'marching_cubes_meshes'), exist_ok=True) | |
| opt.inference.size = opt.model.size | |
| checkpoint_path = os.path.join('full_models', model_path) | |
| checkpoint = torch.load(checkpoint_path) | |
| # Load image generation model | |
| g_ema = Generator(opt.model, opt.rendering).to(device) | |
| pretrained_weights_dict = checkpoint["g_ema"] | |
| model_dict = g_ema.state_dict() | |
| for k, v in pretrained_weights_dict.items(): | |
| if v.size() == model_dict[k].size(): | |
| model_dict[k] = v | |
| g_ema.load_state_dict(model_dict) | |
| # Load a second volume renderer that extracts surfaces at 128x128x128 (or higher) for better surface resolution | |
| if not opt.inference.no_surface_renderings: | |
| opt['surf_extraction'] = Munch() | |
| opt.surf_extraction.rendering = opt.rendering | |
| opt.surf_extraction.model = opt.model.copy() | |
| opt.surf_extraction.model.renderer_spatial_output_dim = 128 | |
| opt.surf_extraction.rendering.N_samples = opt.surf_extraction.model.renderer_spatial_output_dim | |
| opt.surf_extraction.rendering.return_xyz = True | |
| opt.surf_extraction.rendering.return_sdf = True | |
| surface_g_ema = Generator(opt.surf_extraction.model, opt.surf_extraction.rendering, full_pipeline=False).to(device) | |
| # Load weights to surface extractor | |
| surface_extractor_dict = surface_g_ema.state_dict() | |
| for k, v in pretrained_weights_dict.items(): | |
| if k in surface_extractor_dict.keys() and v.size() == surface_extractor_dict[k].size(): | |
| surface_extractor_dict[k] = v | |
| surface_g_ema.load_state_dict(surface_extractor_dict) | |
| else: | |
| surface_g_ema = None | |
| # Get the mean latent vector for g_ema | |
| if opt.inference.truncation_ratio < 1: | |
| with torch.no_grad(): | |
| mean_latent = g_ema.mean_latent(opt.inference.truncation_mean, device) | |
| else: | |
| surface_mean_latent = None | |
| # Get the mean latent vector for surface_g_ema | |
| if not opt.inference.no_surface_renderings: | |
| surface_mean_latent = mean_latent[0] | |
| else: | |
| surface_mean_latent = None | |
| return opt.inference, g_ema, surface_g_ema, mean_latent, surface_mean_latent,opt.inference.results_dst_dir | |
| def get_rendervideo_vars(model_type,number_frames): | |
| opt = BaseOptions().parse() | |
| opt.model.is_test = True | |
| opt.model.style_dim = 256 | |
| opt.model.freeze_renderer = False | |
| opt.inference.size = opt.model.size | |
| opt.inference.camera = opt.camera | |
| opt.inference.renderer_output_size = opt.model.renderer_spatial_output_dim | |
| opt.inference.style_dim = opt.model.style_dim | |
| opt.inference.project_noise = opt.model.project_noise | |
| opt.rendering.perturb = 0 | |
| opt.rendering.force_background = True | |
| opt.rendering.static_viewdirs = True | |
| opt.rendering.return_sdf = True | |
| opt.rendering.N_samples = 64 | |
| opt.inference.identities = 1 | |
| # Load saved model | |
| if model_type == 'ffhq': | |
| model_path = 'ffhq1024x1024.pt' | |
| opt.model.size = 1024 | |
| opt.experiment.expname = 'ffhq1024x1024' | |
| else: | |
| opt.inference.camera.azim = 0.15 | |
| model_path = 'afhq512x512.pt' | |
| opt.model.size = 512 | |
| opt.experiment.expname = 'afhq512x512' | |
| opt.inference.size = opt.model.size | |
| # Create results directory | |
| result_model_dir = 'final_model' | |
| results_dir_basename = os.path.join(opt.inference.results_dir, opt.experiment.expname) | |
| opt.inference.results_dst_dir = os.path.join(results_dir_basename, result_model_dir) | |
| os.makedirs(opt.inference.results_dst_dir, exist_ok=True) | |
| os.makedirs(os.path.join(opt.inference.results_dst_dir, 'videos'), exist_ok=True) | |
| checkpoints_dir = './full_models' | |
| checkpoint_path = os.path.join('full_models', model_path) | |
| if os.path.isfile(checkpoint_path): | |
| # define results directory name | |
| result_model_dir = 'final_model' | |
| results_dir_basename = os.path.join(opt.inference.results_dir, opt.experiment.expname) | |
| opt.inference.results_dst_dir = os.path.join(results_dir_basename, result_model_dir, 'videos') | |
| if opt.model.project_noise: | |
| opt.inference.results_dst_dir = os.path.join(opt.inference.results_dst_dir, 'with_noise_projection') | |
| os.makedirs(opt.inference.results_dst_dir, exist_ok=True) | |
| print(checkpoint_path) | |
| # load saved model | |
| checkpoint = torch.load(checkpoint_path) | |
| # load image generation model | |
| g_ema = Generator(opt.model, opt.rendering).to(device) | |
| # temp fix because of wrong noise sizes | |
| pretrained_weights_dict = checkpoint["g_ema"] | |
| model_dict = g_ema.state_dict() | |
| for k, v in pretrained_weights_dict.items(): | |
| if v.size() == model_dict[k].size(): | |
| model_dict[k] = v | |
| g_ema.load_state_dict(model_dict) | |
| # load a the volume renderee to a second that extracts surfaces at 128x128x128 | |
| if not opt.inference.no_surface_renderings or opt.model.project_noise: | |
| opt['surf_extraction'] = Munch() | |
| opt.surf_extraction.rendering = opt.rendering | |
| opt.surf_extraction.model = opt.model.copy() | |
| opt.surf_extraction.model.renderer_spatial_output_dim = 128 | |
| opt.surf_extraction.rendering.N_samples = opt.surf_extraction.model.renderer_spatial_output_dim | |
| opt.surf_extraction.rendering.return_xyz = True | |
| opt.surf_extraction.rendering.return_sdf = True | |
| opt.inference.surf_extraction_output_size = opt.surf_extraction.model.renderer_spatial_output_dim | |
| surface_g_ema = Generator(opt.surf_extraction.model, opt.surf_extraction.rendering, full_pipeline=False).to(device) | |
| # Load weights to surface extractor | |
| surface_extractor_dict = surface_g_ema.state_dict() | |
| for k, v in pretrained_weights_dict.items(): | |
| if k in surface_extractor_dict.keys() and v.size() == surface_extractor_dict[k].size(): | |
| surface_extractor_dict[k] = v | |
| surface_g_ema.load_state_dict(surface_extractor_dict) | |
| else: | |
| surface_g_ema = None | |
| # get the mean latent vector for g_ema | |
| if opt.inference.truncation_ratio < 1: | |
| with torch.no_grad(): | |
| mean_latent = g_ema.mean_latent(opt.inference.truncation_mean, device) | |
| else: | |
| mean_latent = None | |
| # get the mean latent vector for surface_g_ema | |
| if not opt.inference.no_surface_renderings or opt.model.project_noise: | |
| surface_mean_latent = mean_latent[0] | |
| else: | |
| surface_mean_latent = None | |
| return opt.inference, g_ema, surface_g_ema, mean_latent, surface_mean_latent,opt.inference.results_dst_dir | |
| def render_video(opt, g_ema, surface_g_ema, device, mean_latent, surface_mean_latent,numberofframes): | |
| g_ema.eval() | |
| if not opt.no_surface_renderings or opt.project_noise: | |
| surface_g_ema.eval() | |
| images = torch.Tensor(0, 3, opt.size, opt.size) | |
| num_frames = numberofframes | |
| # Generate video trajectory | |
| trajectory = np.zeros((num_frames,3), dtype=np.float32) | |
| # set camera trajectory | |
| # sweep azimuth angles (4 seconds) | |
| if opt.azim_video: | |
| t = np.linspace(0, 1, num_frames) | |
| elev = 0 | |
| fov = opt.camera.fov | |
| if opt.camera.uniform: | |
| azim = opt.camera.azim * np.cos(t * 2 * np.pi) | |
| else: | |
| azim = 1.5 * opt.camera.azim * np.cos(t * 2 * np.pi) | |
| trajectory[:num_frames,0] = azim | |
| trajectory[:num_frames,1] = elev | |
| trajectory[:num_frames,2] = fov | |
| # elipsoid sweep (4 seconds) | |
| else: | |
| t = np.linspace(0, 1, num_frames) | |
| fov = opt.camera.fov #+ 1 * np.sin(t * 2 * np.pi) | |
| if opt.camera.uniform: | |
| elev = opt.camera.elev / 2 + opt.camera.elev / 2 * np.sin(t * 2 * np.pi) | |
| azim = opt.camera.azim * np.cos(t * 2 * np.pi) | |
| else: | |
| elev = 1.5 * opt.camera.elev * np.sin(t * 2 * np.pi) | |
| azim = 1.5 * opt.camera.azim * np.cos(t * 2 * np.pi) | |
| trajectory[:num_frames,0] = azim | |
| trajectory[:num_frames,1] = elev | |
| trajectory[:num_frames,2] = fov | |
| trajectory = torch.from_numpy(trajectory).to(device) | |
| # generate input parameters for the camera trajectory | |
| # sample_cam_poses, sample_focals, sample_near, sample_far = \ | |
| # generate_camera_params(trajectory, opt.renderer_output_size, device, dist_radius=opt.camera.dist_radius) | |
| sample_cam_extrinsics, sample_focals, sample_near, sample_far, _ = \ | |
| generate_camera_params(opt.renderer_output_size, device, locations=trajectory[:,:2], | |
| fov_ang=trajectory[:,2:], dist_radius=opt.camera.dist_radius) | |
| # In case of noise projection, generate input parameters for the frontal position. | |
| # The reference mesh for the noise projection is extracted from the frontal position. | |
| # For more details see section C.1 in the supplementary material. | |
| if opt.project_noise: | |
| frontal_pose = torch.tensor([[0.0,0.0,opt.camera.fov]]).to(device) | |
| # frontal_cam_pose, frontal_focals, frontal_near, frontal_far = \ | |
| # generate_camera_params(frontal_pose, opt.surf_extraction_output_size, device, dist_radius=opt.camera.dist_radius) | |
| frontal_cam_pose, frontal_focals, frontal_near, frontal_far, _ = \ | |
| generate_camera_params(opt.surf_extraction_output_size, device, location=frontal_pose[:,:2], | |
| fov_ang=frontal_pose[:,2:], dist_radius=opt.camera.dist_radius) | |
| # create geometry renderer (renders the depth maps) | |
| cameras = create_cameras(azim=np.rad2deg(trajectory[0,0].cpu().numpy()), | |
| elev=np.rad2deg(trajectory[0,1].cpu().numpy()), | |
| dist=1, device=device) | |
| renderer = create_mesh_renderer(cameras, image_size=512, specular_color=((0,0,0),), | |
| ambient_color=((0.1,.1,.1),), diffuse_color=((0.75,.75,.75),), | |
| device=device) | |
| suffix = '_azim' if opt.azim_video else '_elipsoid' | |
| # generate videos | |
| for i in range(opt.identities): | |
| print('Processing identity {}/{}...'.format(i+1, opt.identities)) | |
| chunk = 1 | |
| sample_z = torch.randn(1, opt.style_dim, device=device).repeat(chunk,1) | |
| video_filename = 'sample_video_{}{}.mp4'.format(i,suffix) | |
| writer = skvideo.io.FFmpegWriter(os.path.join(opt.results_dst_dir, video_filename), | |
| outputdict={'-pix_fmt': 'yuv420p', '-crf': '10'}) | |
| if not opt.no_surface_renderings: | |
| depth_video_filename = 'sample_depth_video_{}{}.mp4'.format(i,suffix) | |
| depth_writer = skvideo.io.FFmpegWriter(os.path.join(opt.results_dst_dir, depth_video_filename), | |
| outputdict={'-pix_fmt': 'yuv420p', '-crf': '1'}) | |
| ####################### Extract initial surface mesh from the frontal viewpoint ############# | |
| # For more details see section C.1 in the supplementary material. | |
| if opt.project_noise: | |
| with torch.no_grad(): | |
| frontal_surface_out = surface_g_ema([sample_z], | |
| frontal_cam_pose, | |
| frontal_focals, | |
| frontal_near, | |
| frontal_far, | |
| truncation=opt.truncation_ratio, | |
| truncation_latent=surface_mean_latent, | |
| return_sdf=True) | |
| frontal_sdf = frontal_surface_out[2].cpu() | |
| print('Extracting Identity {} Frontal view Marching Cubes for consistent video rendering'.format(i)) | |
| frostum_aligned_frontal_sdf = align_volume(frontal_sdf) | |
| del frontal_sdf | |
| try: | |
| frontal_marching_cubes_mesh = extract_mesh_with_marching_cubes(frostum_aligned_frontal_sdf) | |
| except ValueError: | |
| frontal_marching_cubes_mesh = None | |
| if frontal_marching_cubes_mesh != None: | |
| frontal_marching_cubes_mesh_filename = os.path.join(opt.results_dst_dir,'sample_{}_frontal_marching_cubes_mesh{}.obj'.format(i,suffix)) | |
| with open(frontal_marching_cubes_mesh_filename, 'w') as f: | |
| frontal_marching_cubes_mesh.export(f,file_type='obj') | |
| del frontal_surface_out | |
| torch.cuda.empty_cache() | |
| ############################################################################################# | |
| for j in tqdm(range(0, num_frames, chunk)): | |
| with torch.no_grad(): | |
| out = g_ema([sample_z], | |
| sample_cam_extrinsics[j:j+chunk], | |
| sample_focals[j:j+chunk], | |
| sample_near[j:j+chunk], | |
| sample_far[j:j+chunk], | |
| truncation=opt.truncation_ratio, | |
| truncation_latent=mean_latent, | |
| randomize_noise=False, | |
| project_noise=opt.project_noise, | |
| mesh_path=frontal_marching_cubes_mesh_filename if opt.project_noise else None) | |
| rgb = out[0].cpu() | |
| utils.save_image(rgb, | |
| os.path.join(opt.results_dst_dir, '{}.png'.format(str(i).zfill(7))), | |
| nrow= trajectory[:,:2].shape[0], | |
| normalize=True, | |
| padding=0, | |
| value_range=(-1, 1),) | |
| # this is done to fit to RTX2080 RAM size (11GB) | |
| del out | |
| torch.cuda.empty_cache() | |
| # Convert RGB from [-1, 1] to [0,255] | |
| rgb = 127.5 * (rgb.clamp(-1,1).permute(0,2,3,1).cpu().numpy() + 1) | |
| # Add RGB, frame to video | |
| for k in range(chunk): | |
| writer.writeFrame(rgb[k]) | |
| ########## Extract surface ########## | |
| if not opt.no_surface_renderings: | |
| scale = surface_g_ema.renderer.out_im_res / g_ema.renderer.out_im_res | |
| surface_sample_focals = sample_focals * scale | |
| surface_out = surface_g_ema([sample_z], | |
| sample_cam_extrinsics[j:j+chunk], | |
| surface_sample_focals[j:j+chunk], | |
| sample_near[j:j+chunk], | |
| sample_far[j:j+chunk], | |
| truncation=opt.truncation_ratio, | |
| truncation_latent=surface_mean_latent, | |
| return_xyz=True) | |
| xyz = surface_out[2].cpu() | |
| # this is done to fit to RTX2080 RAM size (11GB) | |
| del surface_out | |
| torch.cuda.empty_cache() | |
| # Render mesh for video | |
| depth_mesh = xyz2mesh(xyz) | |
| mesh = Meshes( | |
| verts=[torch.from_numpy(np.asarray(depth_mesh.vertices)).to(torch.float32).to(device)], | |
| faces = [torch.from_numpy(np.asarray(depth_mesh.faces)).to(torch.float32).to(device)], | |
| textures=None, | |
| verts_normals=[torch.from_numpy(np.copy(np.asarray(depth_mesh.vertex_normals))).to(torch.float32).to(device)], | |
| ) | |
| mesh = add_textures(mesh) | |
| cameras = create_cameras(azim=np.rad2deg(trajectory[j,0].cpu().numpy()), | |
| elev=np.rad2deg(trajectory[j,1].cpu().numpy()), | |
| fov=2*trajectory[j,2].cpu().numpy(), | |
| dist=1, device=device) | |
| renderer = create_mesh_renderer(cameras, image_size=512, | |
| light_location=((0.0,1.0,5.0),), specular_color=((0.2,0.2,0.2),), | |
| ambient_color=((0.1,0.1,0.1),), diffuse_color=((0.65,.65,.65),), | |
| device=device) | |
| mesh_image = 255 * renderer(mesh).cpu().numpy() | |
| mesh_image = mesh_image[...,:3] | |
| # Add depth frame to video | |
| for k in range(chunk): | |
| depth_writer.writeFrame(mesh_image[k]) | |
| # Close video writers | |
| writer.close() | |
| if not opt.no_surface_renderings: | |
| depth_writer.close() | |
| return video_filename | |
| import gradio as gr | |
| import plotly.graph_objects as go | |
| from PIL import Image | |
| device='cuda' if torch.cuda.is_available() else 'cpu' | |
| def get_video(model_type,numberofframes,mesh_type): | |
| options,g_ema,surface_g_ema, mean_latent, surface_mean_latent,result_filename=get_rendervideo_vars(model_type,numberofframes) | |
| render_video(options, g_ema, surface_g_ema, device, mean_latent, surface_mean_latent,numberofframes) | |
| torch.cuda.empty_cache() | |
| del options,g_ema,surface_g_ema, mean_latent, surface_mean_latent | |
| path_img=os.path.join(result_filename,"0000000.png") | |
| image=Image.open(path_img) | |
| if mesh_type=="DepthMesh": | |
| path=os.path.join(result_filename,"sample_depth_video_0_elipsoid.mp4") | |
| else: | |
| path=os.path.join(result_filename,"sample_video_0_elipsoid.mp4") | |
| return path,image | |
| def get_mesh(model_type,mesh_type): | |
| options,g_ema,surface_g_ema, mean_latent, surface_mean_latent,result_filename=get_generate_vars(model_type) | |
| depth_mesh,mc_mesh=generate(options, g_ema, surface_g_ema, device, mean_latent, surface_mean_latent) | |
| torch.cuda.empty_cache() | |
| del options,g_ema,surface_g_ema, mean_latent, surface_mean_latent | |
| if mesh_type=="DepthMesh": | |
| mesh=depth_mesh | |
| else: | |
| mesh=mc_mesh | |
| x=np.asarray(mesh.vertices).T[0] | |
| y=np.asarray(mesh.vertices).T[1] | |
| z=np.asarray(mesh.vertices).T[2] | |
| i=np.asarray(mesh.faces).T[0] | |
| j=np.asarray(mesh.faces).T[1] | |
| k=np.asarray(mesh.faces).T[2] | |
| fig = go.Figure(go.Mesh3d(x=x, y=y, z=z, | |
| i=i, j=j, k=k, | |
| colorscale="Viridis", | |
| colorbar_len=0.75, | |
| flatshading=True, | |
| lighting=dict(ambient=0.5, | |
| diffuse=1, | |
| fresnel=4, | |
| specular=0.5, | |
| roughness=0.05, | |
| facenormalsepsilon=0, | |
| vertexnormalsepsilon=0), | |
| lightposition=dict(x=100, | |
| y=100, | |
| z=1000))) | |
| path=os.path.join(result_filename,"images/0000000.png") | |
| image=Image.open(path) | |
| return fig,image | |
| markdown=f''' | |
| # StyleSDF: High-Resolution 3D-Consistent Image and Geometry Generation | |
| [The space demo for the CVPR 2022 paper "StyleSDF: High-Resolution 3D-Consistent Image and Geometry Generation".](https://arxiv.org/abs/2112.11427) | |
| [For the official implementation.](https://github.com/royorel/StyleSDF) | |
| ### Future Work based on interest | |
| - Adding new models for new type objects | |
| - New Customization | |
| It is running on {device} | |
| The process can take long time.Especially ,To generate videos and the time of process depends the number of frames and current compiler device. | |
| Note : For RGB video , choose marching cubes mesh type | |
| ''' | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown(markdown) | |
| with gr.Column(): | |
| with gr.Row(): | |
| with gr.Column(): | |
| image=gr.Image(type="pil",shape=(512,512)) | |
| with gr.Column(): | |
| mesh = gr.Plot() | |
| with gr.Column(): | |
| video=gr.Video() | |
| with gr.Row(): | |
| numberoframes = gr.Slider( minimum=30, maximum=250,label='Number Of Frame For Video Generation') | |
| model_name=gr.Dropdown(choices=["ffhq","afhq"],label="Choose Model Type") | |
| mesh_type=gr.Dropdown(choices=["DepthMesh","Marching Cubes"],label="Choose Mesh Type") | |
| with gr.Row(): | |
| btn = gr.Button(value="Generate Mesh") | |
| btn_2=gr.Button(value="Generate Video") | |
| btn.click(get_mesh, [model_name,mesh_type],[ mesh,image]) | |
| btn_2.click(get_video,[model_name,numberoframes,mesh_type],[video,image]) | |
| demo.launch(debug=True) | |