Add files using upload-large-folder tool

112139c verified 3 months ago

7.56 kB

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import time
import torch
import argparse
import numpy as np
from PIL import Image
from torchvision import transforms
from collections import OrderedDict

import sys
sys.path.append(".")
sys.path.append("..")

import imageio
from tqdm import tqdm
from configs.paths_config import model_paths
from models.triencoder import TriEncoder
from models.eg3d.shape_utils import convert_sdf_samples_to_ply

device = "cuda:0"

def gen_rand_pose(device):
    return (torch.rand(1, device=device) - 0.5) * torch.pi/6 + torch.pi/2

def layout_grid(img: np.ndarray, 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:
        grid_w = batch_size // grid_h
    assert batch_size == grid_w * grid_h
    if float_to_uint8:
        img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8)
    img = img.reshape(grid_h, grid_w, channels, img_h, img_w)
    img = img.permute(2, 0, 3, 1, 4)
    img = img.reshape(channels, grid_h * img_h, grid_w * img_w)
    if chw_to_hwc:
        img = img.permute(1, 2, 0)
    if to_numpy:
        img = img.cpu().numpy()
    return img

def generate_video(net, gen_triplanes, savepath, batch_size=1):
    voxel_resolution = 512
    video_kwargs = {}
    video_out = imageio.get_writer(savepath, mode='I', fps=60, codec='libx264', **video_kwargs)
    grid_h, grid_w = 1, 1
    num_keyframes, w_frames = 12, 24

    image_mode = "image"
    for frame_idx in tqdm(range(num_keyframes * w_frames)):
        imgs = []
        for yi in range(grid_h):
            for xi in range(grid_w):
                pitch_range = 0.25
                yaw_range = 0.35
                camera_p = net.eural_to_camera(batch_size, 
                        3.14/2 -0.05 + pitch_range * np.cos(2 * 3.14 * frame_idx / (num_keyframes * w_frames)),
                        3.14/2 + yaw_range * np.sin(2 * 3.14 * frame_idx / (num_keyframes * w_frames)),
                        )
                render_res = net.triplane_renderer(gen_triplanes, camera_p)
                img = render_res[image_mode][0]
                if image_mode == 'image_depth':
                    img = -img
                    img = (img - img.min()) / (img.max() - img.min()) * 2 - 1

                imgs.append(img)

        video_out.append_data(layout_grid(torch.stack(imgs), grid_h=1, grid_w=1))
    video_out.close()

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
    voxel_size = cube_length / (N - 1)

    overall_index = torch.arange(0, N ** 3, 1, out=torch.LongTensor())
    samples = torch.zeros(N ** 3, 3)

    # transform first 3 columns
    # to be the x, y, z index
    samples[:, 2] = overall_index % N
    samples[:, 1] = (overall_index.float() / N) % N
    samples[:, 0] = ((overall_index.float() / N) / N) % N

    # transform first 3 columns
    # to be the x, y, z coordinate
    samples[:, 0] = (samples[:, 0] * voxel_size) + voxel_origin[2]
    samples[:, 1] = (samples[:, 1] * voxel_size) + voxel_origin[1]
    samples[:, 2] = (samples[:, 2] * voxel_size) + voxel_origin[0]

    num_samples = N ** 3

    return samples.unsqueeze(0), voxel_origin, voxel_size

def extract_mesh(net, gen_triplanes, savepath):
    shape_res =  512
    max_batch=1000000
    my_plane = gen_triplanes.reshape(1,3,32,256,256)
    samples, voxel_origin, voxel_size = create_samples(N=shape_res, voxel_origin=[0, 0, 0], cube_length=net.triplane_generator.rendering_kwargs['box_warp'] * 1)#.reshape(1, -1, 3)
    samples = samples.to(device)
    sigmas = torch.zeros((samples.shape[0], samples.shape[1], 1), device=device)
    transformed_ray_directions_expanded = torch.zeros((samples.shape[0], max_batch, 3), device=device)
    transformed_ray_directions_expanded[..., -1] = -1

    head = 0
    with tqdm(total = samples.shape[1]) as pbar:
        with torch.no_grad():
            while head < samples.shape[1]:
                torch.manual_seed(0)
                sigma = net.triplane_renderer.sample_density(my_plane, samples[:, head:head+max_batch], transformed_ray_directions_expanded[:, :samples.shape[1]-head])
                sigmas[:, head:head+max_batch] = sigma['sigma']
                head += max_batch
                pbar.update(max_batch)

    sigmas = sigmas.reshape((shape_res, shape_res, shape_res)).cpu().numpy()
    sigmas = np.flip(sigmas, 0)

    # Trim the border of the extracted cube
    pad = int(30 * shape_res / 256)
    pad_value = -1000
    sigmas[:pad] = pad_value
    sigmas[-pad:] = pad_value
    sigmas[:, :pad] = pad_value
    sigmas[:, -pad:] = pad_value
    sigmas[:, :, :pad] = pad_value
    sigmas[:, :, -pad:] = pad_value

    convert_sdf_samples_to_ply(np.transpose(sigmas, (2, 1, 0)), [0, 0, 0], 1, savepath, level=10)

def multi_view_output(net, gen_triplanes, savepath, batch_size=1):
    # observe in multiple views
    image_mode = "image" # or "depth"
    offset = np.linspace(-np.pi/6, np.pi/6, 8)
    imgs_list = []
    for p in offset:
        camera_p = net.eural_to_camera(batch_size, np.pi/2, np.pi/2 + p) # pitch yawl
        # print(camera_p)
        render_res = net.triplane_renderer(gen_triplanes, camera_p)
        img = render_res[image_mode][0]
        imgs_list.append(img.clamp(-1,1).squeeze(0).cpu().detach().permute((1,2,0)).numpy())
    image_list = (np.concatenate(imgs_list, axis=1) + 1) / 2 * 255
    Image.fromarray(image_list.astype(np.uint8)).save(savepath)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default=model_paths["encoder_render"], help='model path')
    parser.add_argument('--image', type=str, default='./imgs/3.png', help='image path')
    parser.add_argument('--output', type=str, default='./output', help='output path')
    parser.add_argument('--save_mesh', type=bool, default=True, help='If true, save mesh')
    parser.add_argument('--save_video', type=bool, default=True, help='If true, save video')
    parser.add_argument('--mode', type=str, default='Normal', help='mode of triplane generation', choices=['LT', 'Normal'])
    args = parser.parse_args()

    model_path = args.model
    image_path = args.image
    output_path = args.output
    mode = args.mode
    if not os.path.exists(output_path):
        os.mkdir(output_path)

    # load model
    model_dict = torch.load(model_path)
    net = TriEncoder(mode=mode)
    net.encoder.load_state_dict(model_dict['encoder_state_dict'])
    net.triplane_renderer.load_state_dict(model_dict['renderer_state_dict'])

    # load image
    img = Image.open(image_path).resize((512, 512))
    img_tensor = transforms.ToTensor()(img).unsqueeze(0).to(device)
    if img_tensor.shape[1] == 4: # remove alpha channel
        img_tensor = img_tensor[:, :3, :, :]    
    img_tensor = img_tensor * 2 - 1

    # inference
    net.encoder.to(device)
    net.encoder.eval()
    basename = os.path.basename(image_path).split('.')[0]
    with torch.no_grad():
        gen_triplanes = net.encoder(img_tensor)
        multi_view_output(net, gen_triplanes, os.path.join(output_path, 'multi_view_' + basename + '.png'))
        if args.save_video:
            generate_video(net, gen_triplanes, os.path.join(output_path, 'video_' + basename + '.mp4'))
        if args.save_mesh:
            extract_mesh(net, gen_triplanes, os.path.join(output_path, 'mesh_' + basename + '.ply'))

	import os
	os.environ["CUDA_VISIBLE_DEVICES"] = "3"
	import time
	import torch
	import argparse
	import numpy as np
	from PIL import Image
	from torchvision import transforms
	from collections import OrderedDict

	import sys
	sys.path.append(".")
	sys.path.append("..")

	import imageio
	from tqdm import tqdm
	from configs.paths_config import model_paths
	from models.triencoder import TriEncoder
	from models.eg3d.shape_utils import convert_sdf_samples_to_ply

	device = "cuda:0"

	def gen_rand_pose(device):
	return (torch.rand(1, device=device) - 0.5) * torch.pi/6 + torch.pi/2

	def layout_grid(img: np.ndarray, 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:
	grid_w = batch_size // grid_h
	assert batch_size == grid_w * grid_h
	if float_to_uint8:
	img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8)
	img = img.reshape(grid_h, grid_w, channels, img_h, img_w)
	img = img.permute(2, 0, 3, 1, 4)
	img = img.reshape(channels, grid_h * img_h, grid_w * img_w)
	if chw_to_hwc:
	img = img.permute(1, 2, 0)
	if to_numpy:
	img = img.cpu().numpy()
	return img

	def generate_video(net, gen_triplanes, savepath, batch_size=1):
	voxel_resolution = 512
	video_kwargs = {}
	video_out = imageio.get_writer(savepath, mode='I', fps=60, codec='libx264', **video_kwargs)
	grid_h, grid_w = 1, 1
	num_keyframes, w_frames = 12, 24

	image_mode = "image"
	for frame_idx in tqdm(range(num_keyframes * w_frames)):
	imgs = []
	for yi in range(grid_h):
	for xi in range(grid_w):
	pitch_range = 0.25
	yaw_range = 0.35
	camera_p = net.eural_to_camera(batch_size,
	3.14/2 -0.05 + pitch_range * np.cos(2 * 3.14 * frame_idx / (num_keyframes * w_frames)),
	3.14/2 + yaw_range * np.sin(2 * 3.14 * frame_idx / (num_keyframes * w_frames)),
	)
	render_res = net.triplane_renderer(gen_triplanes, camera_p)
	img = render_res[image_mode][0]
	if image_mode == 'image_depth':
	img = -img
	img = (img - img.min()) / (img.max() - img.min()) * 2 - 1

	imgs.append(img)

	video_out.append_data(layout_grid(torch.stack(imgs), grid_h=1, grid_w=1))
	video_out.close()

	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
	voxel_size = cube_length / (N - 1)

	overall_index = torch.arange(0, N ** 3, 1, out=torch.LongTensor())
	samples = torch.zeros(N ** 3, 3)

	# transform first 3 columns
	# to be the x, y, z index
	samples[:, 2] = overall_index % N
	samples[:, 1] = (overall_index.float() / N) % N
	samples[:, 0] = ((overall_index.float() / N) / N) % N

	# transform first 3 columns
	# to be the x, y, z coordinate
	samples[:, 0] = (samples[:, 0] * voxel_size) + voxel_origin[2]
	samples[:, 1] = (samples[:, 1] * voxel_size) + voxel_origin[1]
	samples[:, 2] = (samples[:, 2] * voxel_size) + voxel_origin[0]

	num_samples = N ** 3

	return samples.unsqueeze(0), voxel_origin, voxel_size

	def extract_mesh(net, gen_triplanes, savepath):
	shape_res = 512
	max_batch=1000000
	my_plane = gen_triplanes.reshape(1,3,32,256,256)
	samples, voxel_origin, voxel_size = create_samples(N=shape_res, voxel_origin=[0, 0, 0], cube_length=net.triplane_generator.rendering_kwargs['box_warp'] * 1)#.reshape(1, -1, 3)
	samples = samples.to(device)
	sigmas = torch.zeros((samples.shape[0], samples.shape[1], 1), device=device)
	transformed_ray_directions_expanded = torch.zeros((samples.shape[0], max_batch, 3), device=device)
	transformed_ray_directions_expanded[..., -1] = -1

	head = 0
	with tqdm(total = samples.shape[1]) as pbar:
	with torch.no_grad():
	while head < samples.shape[1]:
	torch.manual_seed(0)
	sigma = net.triplane_renderer.sample_density(my_plane, samples[:, head:head+max_batch], transformed_ray_directions_expanded[:, :samples.shape[1]-head])
	sigmas[:, head:head+max_batch] = sigma['sigma']
	head += max_batch
	pbar.update(max_batch)

	sigmas = sigmas.reshape((shape_res, shape_res, shape_res)).cpu().numpy()
	sigmas = np.flip(sigmas, 0)

	# Trim the border of the extracted cube
	pad = int(30 * shape_res / 256)
	pad_value = -1000
	sigmas[:pad] = pad_value
	sigmas[-pad:] = pad_value
	sigmas[:, :pad] = pad_value
	sigmas[:, -pad:] = pad_value
	sigmas[:, :, :pad] = pad_value
	sigmas[:, :, -pad:] = pad_value

	convert_sdf_samples_to_ply(np.transpose(sigmas, (2, 1, 0)), [0, 0, 0], 1, savepath, level=10)

	def multi_view_output(net, gen_triplanes, savepath, batch_size=1):
	# observe in multiple views
	image_mode = "image" # or "depth"
	offset = np.linspace(-np.pi/6, np.pi/6, 8)
	imgs_list = []
	for p in offset:
	camera_p = net.eural_to_camera(batch_size, np.pi/2, np.pi/2 + p) # pitch yawl
	# print(camera_p)
	render_res = net.triplane_renderer(gen_triplanes, camera_p)
	img = render_res[image_mode][0]
	imgs_list.append(img.clamp(-1,1).squeeze(0).cpu().detach().permute((1,2,0)).numpy())
	image_list = (np.concatenate(imgs_list, axis=1) + 1) / 2 * 255
	Image.fromarray(image_list.astype(np.uint8)).save(savepath)


	if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('--model', type=str, default=model_paths["encoder_render"], help='model path')
	parser.add_argument('--image', type=str, default='./imgs/3.png', help='image path')
	parser.add_argument('--output', type=str, default='./output', help='output path')
	parser.add_argument('--save_mesh', type=bool, default=True, help='If true, save mesh')
	parser.add_argument('--save_video', type=bool, default=True, help='If true, save video')
	parser.add_argument('--mode', type=str, default='Normal', help='mode of triplane generation', choices=['LT', 'Normal'])
	args = parser.parse_args()

	model_path = args.model
	image_path = args.image
	output_path = args.output
	mode = args.mode
	if not os.path.exists(output_path):
	os.mkdir(output_path)

	# load model
	model_dict = torch.load(model_path)
	net = TriEncoder(mode=mode)
	net.encoder.load_state_dict(model_dict['encoder_state_dict'])
	net.triplane_renderer.load_state_dict(model_dict['renderer_state_dict'])

	# load image
	img = Image.open(image_path).resize((512, 512))
	img_tensor = transforms.ToTensor()(img).unsqueeze(0).to(device)
	if img_tensor.shape[1] == 4: # remove alpha channel
	img_tensor = img_tensor[:, :3, :, :]
	img_tensor = img_tensor * 2 - 1

	# inference
	net.encoder.to(device)
	net.encoder.eval()
	basename = os.path.basename(image_path).split('.')[0]
	with torch.no_grad():
	gen_triplanes = net.encoder(img_tensor)
	multi_view_output(net, gen_triplanes, os.path.join(output_path, 'multi_view_' + basename + '.png'))
	if args.save_video:
	generate_video(net, gen_triplanes, os.path.join(output_path, 'video_' + basename + '.mp4'))
	if args.save_mesh:
	extract_mesh(net, gen_triplanes, os.path.join(output_path, 'mesh_' + basename + '.ply'))