Update to 2026-09-17 pipeline snapshot; add weights, L20 assets and recording via Git LFS
Source: RGB-D -> Dyn-HaMR -> L20 retargeting -> FoundationPose -> reference repair -> SPIDER, documented in docs/PIPELINE_LATEST.md and docs/SETUP_AND_WEIGHTS.md. Adds FoundationPose and nvdiffrast upstream snapshots, requirements/pipeline_venv.txt and the FoundationPose weight manifest/downloader. Assets (Git LFS): weights/ (WiLoR detector, HandFlow denoiser, UniDepth-L), FoundationPose checkpoints, HaMeR checkpoint, Dyn-HaMR HMP model and BMC constraints, L20 URDF/meshes, the 20260915_171525 D405 recording and the two box CADs. MANO models are not redistributed (third_party/hamer/_DATA/data/mano/README.txt). Environments, caches and run outputs excluded. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
liyang
+202
@@ -0,0 +1,202 @@
|
||||
# Copyright (c) 2020, 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.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pathlib
|
||||
import numpy as np
|
||||
import torch
|
||||
import imageio
|
||||
|
||||
import util
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
# Transform vertex positions to clip space
|
||||
def transform_pos(mtx, pos):
|
||||
t_mtx = torch.from_numpy(mtx).cuda() if isinstance(mtx, np.ndarray) else mtx
|
||||
# (x,y,z) -> (x,y,z,1)
|
||||
posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()], axis=1)
|
||||
return torch.matmul(posw, t_mtx.t())[None, ...]
|
||||
|
||||
def render(glctx, mtx, pos, pos_idx, vtx_col, col_idx, resolution: int):
|
||||
pos_clip = transform_pos(mtx, pos)
|
||||
rast_out, _ = dr.rasterize(glctx, pos_clip, pos_idx, resolution=[resolution, resolution])
|
||||
color, _ = dr.interpolate(vtx_col[None, ...], rast_out, col_idx)
|
||||
color = dr.antialias(color, rast_out, pos_clip, pos_idx)
|
||||
return color
|
||||
|
||||
def make_grid(arr, ncols=2):
|
||||
n, height, width, nc = arr.shape
|
||||
nrows = n//ncols
|
||||
assert n == nrows*ncols
|
||||
return arr.reshape(nrows, ncols, height, width, nc).swapaxes(1,2).reshape(height*nrows, width*ncols, nc)
|
||||
|
||||
def fit_cube(max_iter = 5000,
|
||||
resolution = 4,
|
||||
discontinuous = False,
|
||||
repeats = 1,
|
||||
log_interval = 10,
|
||||
display_interval = None,
|
||||
display_res = 512,
|
||||
out_dir = None,
|
||||
log_fn = None,
|
||||
mp4save_interval = None,
|
||||
mp4save_fn = None):
|
||||
|
||||
log_file = None
|
||||
writer = None
|
||||
if out_dir:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
if log_fn:
|
||||
log_file = open(f'{out_dir}/{log_fn}', 'wt')
|
||||
if mp4save_interval != 0:
|
||||
writer = imageio.get_writer(f'{out_dir}/{mp4save_fn}', mode='I', fps=30, codec='libx264', bitrate='16M')
|
||||
else:
|
||||
mp4save_interval = None
|
||||
|
||||
datadir = f'{pathlib.Path(__file__).absolute().parents[1]}/data'
|
||||
fn = 'cube_%s.npz' % ('d' if discontinuous else 'c')
|
||||
with np.load(f'{datadir}/{fn}') as f:
|
||||
pos_idx, vtxp, col_idx, vtxc = f.values()
|
||||
print("Mesh has %d triangles and %d vertices." % (pos_idx.shape[0], vtxp.shape[0]))
|
||||
|
||||
# Create position/triangle index tensors
|
||||
pos_idx = torch.from_numpy(pos_idx.astype(np.int32)).cuda()
|
||||
col_idx = torch.from_numpy(col_idx.astype(np.int32)).cuda()
|
||||
vtx_pos = torch.from_numpy(vtxp.astype(np.float32)).cuda()
|
||||
vtx_col = torch.from_numpy(vtxc.astype(np.float32)).cuda()
|
||||
|
||||
# Rasterizer context
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
|
||||
# Repeats.
|
||||
for rep in range(repeats):
|
||||
|
||||
ang = 0.0
|
||||
gl_avg = []
|
||||
|
||||
vtx_pos_rand = np.random.uniform(-0.5, 0.5, size=vtxp.shape) + vtxp
|
||||
vtx_col_rand = np.random.uniform(0.0, 1.0, size=vtxc.shape)
|
||||
vtx_pos_opt = torch.tensor(vtx_pos_rand, dtype=torch.float32, device='cuda', requires_grad=True)
|
||||
vtx_col_opt = torch.tensor(vtx_col_rand, dtype=torch.float32, device='cuda', requires_grad=True)
|
||||
|
||||
# Adam optimizer for vertex position and color with a learning rate ramp.
|
||||
optimizer = torch.optim.Adam([vtx_pos_opt, vtx_col_opt], lr=1e-2)
|
||||
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: max(0.01, 10**(-x*0.0005)))
|
||||
|
||||
for it in range(max_iter + 1):
|
||||
# Random rotation/translation matrix for optimization.
|
||||
r_rot = util.random_rotation_translation(0.25)
|
||||
|
||||
# Smooth rotation for display.
|
||||
a_rot = np.matmul(util.rotate_x(0.4), util.rotate_y(ang))
|
||||
|
||||
# Modelview and modelview + projection matrices.
|
||||
proj = util.projection(x=0.4)
|
||||
r_mv = np.matmul(util.translate(0, 0, -3.5), r_rot)
|
||||
r_mvp = np.matmul(proj, r_mv).astype(np.float32)
|
||||
a_mv = np.matmul(util.translate(0, 0, -3.5), a_rot)
|
||||
a_mvp = np.matmul(proj, a_mv).astype(np.float32)
|
||||
|
||||
# Compute geometric error for logging.
|
||||
with torch.no_grad():
|
||||
geom_loss = torch.mean(torch.sum((torch.abs(vtx_pos_opt) - .5)**2, dim=1)**0.5)
|
||||
gl_avg.append(float(geom_loss))
|
||||
|
||||
# Print/save log.
|
||||
if log_interval and (it % log_interval == 0):
|
||||
gl_val = np.mean(np.asarray(gl_avg))
|
||||
gl_avg = []
|
||||
s = ("rep=%d," % rep) if repeats > 1 else ""
|
||||
s += "iter=%d,err=%f" % (it, gl_val)
|
||||
print(s)
|
||||
if log_file:
|
||||
log_file.write(s + "\n")
|
||||
|
||||
color = render(glctx, r_mvp, vtx_pos, pos_idx, vtx_col, col_idx, resolution)
|
||||
color_opt = render(glctx, r_mvp, vtx_pos_opt, pos_idx, vtx_col_opt, col_idx, resolution)
|
||||
|
||||
# Compute loss and train.
|
||||
loss = torch.mean((color - color_opt)**2) # L2 pixel loss.
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
|
||||
# Show/save image.
|
||||
display_image = display_interval and (it % display_interval == 0)
|
||||
save_mp4 = mp4save_interval and (it % mp4save_interval == 0)
|
||||
|
||||
if display_image or save_mp4:
|
||||
ang = ang + 0.01
|
||||
|
||||
img_b = color[0].cpu().numpy()[::-1]
|
||||
img_o = color_opt[0].detach().cpu().numpy()[::-1]
|
||||
img_d = render(glctx, a_mvp, vtx_pos_opt, pos_idx, vtx_col_opt, col_idx, display_res)[0]
|
||||
img_r = render(glctx, a_mvp, vtx_pos, pos_idx, vtx_col, col_idx, display_res)[0]
|
||||
|
||||
scl = display_res // img_o.shape[0]
|
||||
img_b = np.repeat(np.repeat(img_b, scl, axis=0), scl, axis=1)
|
||||
img_o = np.repeat(np.repeat(img_o, scl, axis=0), scl, axis=1)
|
||||
result_image = make_grid(np.stack([img_o, img_b, img_d.detach().cpu().numpy()[::-1], img_r.cpu().numpy()[::-1]]))
|
||||
|
||||
if display_image:
|
||||
util.display_image(result_image, size=display_res, title='%d / %d' % (it, max_iter))
|
||||
if save_mp4:
|
||||
writer.append_data(np.clip(np.rint(result_image*255.0), 0, 255).astype(np.uint8))
|
||||
|
||||
# Done.
|
||||
if writer is not None:
|
||||
writer.close()
|
||||
if log_file:
|
||||
log_file.close()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Cube fit example')
|
||||
parser.add_argument('--outdir', help='specify output directory', default='')
|
||||
parser.add_argument('--discontinuous', action='store_true', default=False)
|
||||
parser.add_argument('--resolution', type=int, default=0, required=True)
|
||||
parser.add_argument('--display-interval', type=int, default=0)
|
||||
parser.add_argument('--mp4save-interval', type=int, default=100)
|
||||
parser.add_argument('--max-iter', type=int, default=1000)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up logging.
|
||||
if args.outdir:
|
||||
ds = 'd' if args.discontinuous else 'c'
|
||||
out_dir = f'{args.outdir}/cube_{ds}_{args.resolution}'
|
||||
print (f'Saving results under {out_dir}')
|
||||
else:
|
||||
out_dir = None
|
||||
print ('No output directory specified, not saving log or images')
|
||||
|
||||
# Run.
|
||||
fit_cube(
|
||||
max_iter=args.max_iter,
|
||||
resolution=args.resolution,
|
||||
discontinuous=args.discontinuous,
|
||||
log_interval=10,
|
||||
display_interval=args.display_interval,
|
||||
out_dir=out_dir,
|
||||
log_fn='log.txt',
|
||||
mp4save_interval=args.mp4save_interval,
|
||||
mp4save_fn='progress.mp4',
|
||||
)
|
||||
|
||||
# Done.
|
||||
print("Done.")
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
+205
@@ -0,0 +1,205 @@
|
||||
# Copyright (c) 2020, 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.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pathlib
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import util
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Helpers.
|
||||
|
||||
def transform_pos(mtx, pos):
|
||||
t_mtx = torch.from_numpy(mtx).cuda() if isinstance(mtx, np.ndarray) else mtx
|
||||
posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()], axis=1)
|
||||
return torch.matmul(posw, t_mtx.t())[None, ...]
|
||||
|
||||
def render(glctx, mtx, pos, pos_idx, uv, uv_idx, tex, resolution, enable_mip, max_mip_level):
|
||||
pos_clip = transform_pos(mtx, pos)
|
||||
rast_out, rast_out_db = dr.rasterize(glctx, pos_clip, pos_idx, resolution=[resolution, resolution])
|
||||
|
||||
if enable_mip:
|
||||
texc, texd = dr.interpolate(uv[None, ...], rast_out, uv_idx, rast_db=rast_out_db, diff_attrs='all')
|
||||
color = dr.texture(tex[None, ...], texc, texd, filter_mode='linear-mipmap-linear', max_mip_level=max_mip_level)
|
||||
else:
|
||||
texc, _ = dr.interpolate(uv[None, ...], rast_out, uv_idx)
|
||||
color = dr.texture(tex[None, ...], texc, filter_mode='linear')
|
||||
|
||||
color = color * torch.clamp(rast_out[..., -1:], 0, 1) # Mask out background.
|
||||
return color
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def fit_earth(max_iter = 20000,
|
||||
log_interval = 10,
|
||||
display_interval = None,
|
||||
display_res = 1024,
|
||||
enable_mip = True,
|
||||
res = 512,
|
||||
ref_res = 2048, # Dropped from 4096 to 2048 to allow using the Cuda rasterizer.
|
||||
lr_base = 1e-2,
|
||||
lr_ramp = 0.1,
|
||||
out_dir = None,
|
||||
log_fn = None,
|
||||
texsave_interval = None,
|
||||
texsave_fn = None,
|
||||
imgsave_interval = None,
|
||||
imgsave_fn = None):
|
||||
|
||||
log_file = None
|
||||
if out_dir:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
if log_fn:
|
||||
log_file = open(out_dir + '/' + log_fn, 'wt')
|
||||
else:
|
||||
imgsave_interval, texsave_interval = None, None
|
||||
|
||||
# Mesh and texture adapted from "3D Earth Photorealistic 2K" model at
|
||||
# https://www.turbosquid.com/3d-models/3d-realistic-earth-photorealistic-2k-1279125
|
||||
datadir = f'{pathlib.Path(__file__).absolute().parents[1]}/data'
|
||||
with np.load(f'{datadir}/earth.npz') as f:
|
||||
pos_idx, pos, uv_idx, uv, tex = f.values()
|
||||
tex = tex.astype(np.float32)/255.0
|
||||
max_mip_level = 9 # Texture is a 4x3 atlas of 512x512 maps.
|
||||
print("Mesh has %d triangles and %d vertices." % (pos_idx.shape[0], pos.shape[0]))
|
||||
|
||||
# Some input geometry contains vertex positions in (N, 4) (with v[:,3]==1). Drop
|
||||
# the last column in that case.
|
||||
if pos.shape[1] == 4: pos = pos[:, 0:3]
|
||||
|
||||
# Create position/triangle index tensors
|
||||
pos_idx = torch.from_numpy(pos_idx.astype(np.int32)).cuda()
|
||||
vtx_pos = torch.from_numpy(pos.astype(np.float32)).cuda()
|
||||
uv_idx = torch.from_numpy(uv_idx.astype(np.int32)).cuda()
|
||||
vtx_uv = torch.from_numpy(uv.astype(np.float32)).cuda()
|
||||
|
||||
tex = torch.from_numpy(tex.astype(np.float32)).cuda()
|
||||
tex_opt = torch.full(tex.shape, 0.2, device='cuda', requires_grad=True)
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
|
||||
ang = 0.0
|
||||
|
||||
# Adam optimizer for texture with a learning rate ramp.
|
||||
optimizer = torch.optim.Adam([tex_opt], lr=lr_base)
|
||||
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: lr_ramp**(float(x)/float(max_iter)))
|
||||
|
||||
# Render.
|
||||
ang = 0.0
|
||||
texloss_avg = []
|
||||
for it in range(max_iter + 1):
|
||||
# Random rotation/translation matrix for optimization.
|
||||
r_rot = util.random_rotation_translation(0.25)
|
||||
|
||||
# Smooth rotation for display.
|
||||
a_rot = np.matmul(util.rotate_x(0.4), util.rotate_y(ang))
|
||||
dist = np.random.uniform(0.0, 48.5)
|
||||
|
||||
# Modelview and modelview + projection matrices.
|
||||
proj = util.projection(x=0.4, n=1.0, f=200.0)
|
||||
r_mv = np.matmul(util.translate(0, 0, -1.5-dist), r_rot)
|
||||
r_mvp = np.matmul(proj, r_mv).astype(np.float32)
|
||||
a_mv = np.matmul(util.translate(0, 0, -3.5), a_rot)
|
||||
a_mvp = np.matmul(proj, a_mv).astype(np.float32)
|
||||
|
||||
# Measure texture-space RMSE loss
|
||||
with torch.no_grad():
|
||||
texmask = torch.zeros_like(tex)
|
||||
tr = tex.shape[1]//4
|
||||
texmask[tr+13:2*tr-13, 25:-25, :] += 1.0
|
||||
texmask[25:-25, tr+13:2*tr-13, :] += 1.0
|
||||
# Measure only relevant portions of texture when calculating texture
|
||||
# PSNR.
|
||||
texloss = (torch.sum(texmask * (tex - tex_opt)**2)/torch.sum(texmask))**0.5 # RMSE within masked area.
|
||||
texloss_avg.append(float(texloss))
|
||||
|
||||
# Render reference and optimized frames. Always enable mipmapping for reference.
|
||||
color = render(glctx, r_mvp, vtx_pos, pos_idx, vtx_uv, uv_idx, tex, ref_res, True, max_mip_level)
|
||||
color_opt = render(glctx, r_mvp, vtx_pos, pos_idx, vtx_uv, uv_idx, tex_opt, res, enable_mip, max_mip_level)
|
||||
|
||||
# Reduce the reference to correct size.
|
||||
while color.shape[1] > res:
|
||||
color = util.bilinear_downsample(color)
|
||||
|
||||
# Compute loss and perform a training step.
|
||||
loss = torch.mean((color - color_opt)**2) # L2 pixel loss.
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
|
||||
# Print/save log.
|
||||
if log_interval and (it % log_interval == 0):
|
||||
texloss_val = np.mean(np.asarray(texloss_avg))
|
||||
texloss_avg = []
|
||||
psnr = -10.0 * np.log10(texloss_val**2) # PSNR based on average RMSE.
|
||||
s = "iter=%d,loss=%f,psnr=%f" % (it, texloss_val, psnr)
|
||||
print(s)
|
||||
if log_file:
|
||||
log_file.write(s + '\n')
|
||||
|
||||
# Show/save image.
|
||||
display_image = display_interval and (it % display_interval == 0)
|
||||
save_image = imgsave_interval and (it % imgsave_interval == 0)
|
||||
save_texture = texsave_interval and (it % texsave_interval) == 0
|
||||
|
||||
if display_image or save_image:
|
||||
ang = ang + 0.1
|
||||
|
||||
with torch.no_grad():
|
||||
result_image = render(glctx, a_mvp, vtx_pos, pos_idx, vtx_uv, uv_idx, tex_opt, res, enable_mip, max_mip_level)[0].cpu().numpy()[::-1]
|
||||
|
||||
if display_image:
|
||||
util.display_image(result_image, size=display_res, title='%d / %d' % (it, max_iter))
|
||||
if save_image:
|
||||
util.save_image(out_dir + '/' + (imgsave_fn % it), result_image)
|
||||
|
||||
if save_texture:
|
||||
texture = tex_opt.cpu().numpy()[::-1]
|
||||
util.save_image(out_dir + '/' + (texsave_fn % it), texture)
|
||||
|
||||
|
||||
# Done.
|
||||
if log_file:
|
||||
log_file.close()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Earth texture fitting example')
|
||||
parser.add_argument('--outdir', help='specify output directory', default='')
|
||||
parser.add_argument('--mip', help='enable mipmapping', action='store_true', default=False)
|
||||
parser.add_argument('--display-interval', type=int, default=0)
|
||||
parser.add_argument('--max-iter', type=int, default=10000)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up logging.
|
||||
if args.outdir:
|
||||
ms = 'mip' if args.mip else 'nomip'
|
||||
out_dir = f'{args.outdir}/earth_{ms}'
|
||||
print (f'Saving results under {out_dir}')
|
||||
else:
|
||||
out_dir = None
|
||||
print ('No output directory specified, not saving log or images')
|
||||
|
||||
# Run.
|
||||
fit_earth(max_iter=args.max_iter, log_interval=10, display_interval=args.display_interval, enable_mip=args.mip, out_dir=out_dir, log_fn='log.txt', texsave_interval=1000, texsave_fn='tex_%06d.png', imgsave_interval=1000, imgsave_fn='img_%06d.png')
|
||||
|
||||
# Done.
|
||||
print("Done.")
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
+227
@@ -0,0 +1,227 @@
|
||||
# Copyright (c) 2020, 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.
|
||||
|
||||
import argparse
|
||||
import numpy as np
|
||||
import torch
|
||||
import os
|
||||
import pathlib
|
||||
import imageio
|
||||
|
||||
import util
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Environment map and Phong BRDF learning.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def fit_env_phong(max_iter = 1000,
|
||||
log_interval = 10,
|
||||
display_interval = None,
|
||||
display_res = 1024,
|
||||
res = 1024,
|
||||
lr_base = 1e-2,
|
||||
lr_ramp = 1.0,
|
||||
out_dir = None,
|
||||
log_fn = None,
|
||||
mp4save_interval = None,
|
||||
mp4save_fn = None):
|
||||
|
||||
log_file = None
|
||||
writer = None
|
||||
if out_dir:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
if log_fn:
|
||||
log_file = open(out_dir + '/' + log_fn, 'wt')
|
||||
if mp4save_interval != 0:
|
||||
writer = imageio.get_writer(f'{out_dir}/{mp4save_fn}', mode='I', fps=30, codec='libx264', bitrate='16M')
|
||||
else:
|
||||
mp4save_interval = None
|
||||
|
||||
# Texture adapted from https://github.com/WaveEngine/Samples/tree/master/Materials/EnvironmentMap/Content/Assets/CubeMap.cubemap
|
||||
datadir = f'{pathlib.Path(__file__).absolute().parents[1]}/data'
|
||||
with np.load(f'{datadir}/envphong.npz') as f:
|
||||
pos_idx, pos, normals, env = f.values()
|
||||
env = env.astype(np.float32)/255.0
|
||||
env = np.stack(env)[:, ::-1].copy()
|
||||
print("Mesh has %d triangles and %d vertices." % (pos_idx.shape[0], pos.shape[0]))
|
||||
|
||||
# Move all the stuff to GPU.
|
||||
pos_idx = torch.as_tensor(pos_idx, dtype=torch.int32, device='cuda')
|
||||
pos = torch.as_tensor(pos, dtype=torch.float32, device='cuda')
|
||||
normals = torch.as_tensor(normals, dtype=torch.float32, device='cuda')
|
||||
env = torch.as_tensor(env, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Target Phong parameters.
|
||||
phong_rgb = np.asarray([1.0, 0.8, 0.6], np.float32)
|
||||
phong_exp = 25.0
|
||||
phong_rgb_t = torch.as_tensor(phong_rgb, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Learned variables: environment maps, phong color, phong exponent.
|
||||
env_var = torch.ones_like(env) * .5
|
||||
env_var.requires_grad_()
|
||||
phong_var_raw = torch.as_tensor(np.random.uniform(size=[4]), dtype=torch.float32, device='cuda')
|
||||
phong_var_raw.requires_grad_()
|
||||
phong_var_mul = torch.as_tensor([1.0, 1.0, 1.0, 10.0], dtype=torch.float32, device='cuda')
|
||||
|
||||
# Render.
|
||||
ang = 0.0
|
||||
imgloss_avg, phong_avg = [], []
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
zero_tensor = torch.as_tensor(0.0, dtype=torch.float32, device='cuda')
|
||||
one_tensor = torch.as_tensor(1.0, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Adam optimizer for environment map and phong with a learning rate ramp.
|
||||
optimizer = torch.optim.Adam([env_var, phong_var_raw], lr=lr_base)
|
||||
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: lr_ramp**(float(x)/float(max_iter)))
|
||||
|
||||
for it in range(max_iter + 1):
|
||||
phong_var = phong_var_raw * phong_var_mul
|
||||
|
||||
# Random rotation/translation matrix for optimization.
|
||||
r_rot = util.random_rotation_translation(0.25)
|
||||
|
||||
# Smooth rotation for display.
|
||||
ang = ang + 0.01
|
||||
a_rot = np.matmul(util.rotate_x(0.4), util.rotate_y(ang))
|
||||
|
||||
# Modelview and modelview + projection matrices.
|
||||
proj = util.projection(x=0.4, n=1.0, f=200.0)
|
||||
r_mv = np.matmul(util.translate(0, 0, -3.5), r_rot)
|
||||
r_mvp = np.matmul(proj, r_mv).astype(np.float32)
|
||||
a_mv = np.matmul(util.translate(0, 0, -3.5), a_rot)
|
||||
a_mvp = np.matmul(proj, a_mv).astype(np.float32)
|
||||
a_mvc = a_mvp
|
||||
r_mvp = torch.as_tensor(r_mvp, dtype=torch.float32, device='cuda')
|
||||
a_mvp = torch.as_tensor(a_mvp, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Solve camera positions.
|
||||
a_campos = torch.as_tensor(np.linalg.inv(a_mv)[:3, 3], dtype=torch.float32, device='cuda')
|
||||
r_campos = torch.as_tensor(np.linalg.inv(r_mv)[:3, 3], dtype=torch.float32, device='cuda')
|
||||
|
||||
# Random light direction.
|
||||
lightdir = np.random.normal(size=[3])
|
||||
lightdir /= np.linalg.norm(lightdir) + 1e-8
|
||||
lightdir = torch.as_tensor(lightdir, dtype=torch.float32, device='cuda')
|
||||
|
||||
def render_refl(ldir, cpos, mvp):
|
||||
# Transform and rasterize.
|
||||
viewvec = pos[..., :3] - cpos[np.newaxis, np.newaxis, :] # View vectors at vertices.
|
||||
reflvec = viewvec - 2.0 * normals[np.newaxis, ...] * torch.sum(normals[np.newaxis, ...] * viewvec, -1, keepdim=True) # Reflection vectors at vertices.
|
||||
reflvec = reflvec / torch.sum(reflvec**2, -1, keepdim=True)**0.5 # Normalize.
|
||||
pos_clip = torch.matmul(pos, mvp.t())[np.newaxis, ...]
|
||||
rast_out, rast_out_db = dr.rasterize(glctx, pos_clip, pos_idx, [res, res])
|
||||
refl, refld = dr.interpolate(reflvec, rast_out, pos_idx, rast_db=rast_out_db, diff_attrs='all') # Interpolated reflection vectors.
|
||||
|
||||
# Phong light.
|
||||
refl = refl / (torch.sum(refl**2, -1, keepdim=True) + 1e-8)**0.5 # Normalize.
|
||||
ldotr = torch.sum(-ldir * refl, -1, keepdim=True) # L dot R.
|
||||
|
||||
# Return
|
||||
return refl, refld, ldotr, (rast_out[..., -1:] == 0)
|
||||
|
||||
# Render the reflections.
|
||||
refl, refld, ldotr, mask = render_refl(lightdir, r_campos, r_mvp)
|
||||
|
||||
# Reference color. No need for AA because we are not learning geometry.
|
||||
color = dr.texture(env[np.newaxis, ...], refl, uv_da=refld, filter_mode='linear-mipmap-linear', boundary_mode='cube')
|
||||
color = color + phong_rgb_t * torch.max(zero_tensor, ldotr) ** phong_exp # Phong.
|
||||
color = torch.where(mask, one_tensor, color) # White background.
|
||||
|
||||
# Candidate rendering same up to this point, but uses learned texture and Phong parameters instead.
|
||||
color_opt = dr.texture(env_var[np.newaxis, ...], refl, uv_da=refld, filter_mode='linear-mipmap-linear', boundary_mode='cube')
|
||||
color_opt = color_opt + phong_var[:3] * torch.max(zero_tensor, ldotr) ** phong_var[3] # Phong.
|
||||
color_opt = torch.where(mask, one_tensor, color_opt) # White background.
|
||||
|
||||
# Compute loss and train.
|
||||
loss = torch.mean((color - color_opt)**2) # L2 pixel loss.
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
|
||||
# Collect losses.
|
||||
imgloss_avg.append(loss.detach().cpu().numpy())
|
||||
phong_avg.append(phong_var.detach().cpu().numpy())
|
||||
|
||||
# Print/save log.
|
||||
if log_interval and (it % log_interval == 0):
|
||||
imgloss_val, imgloss_avg = np.mean(np.asarray(imgloss_avg, np.float32)), []
|
||||
phong_val, phong_avg = np.mean(np.asarray(phong_avg, np.float32), axis=0), []
|
||||
phong_rgb_rmse = np.mean((phong_val[:3] - phong_rgb)**2)**0.5
|
||||
phong_exp_rel_err = np.abs(phong_val[3] - phong_exp)/phong_exp
|
||||
s = "iter=%d,phong_rgb_rmse=%f,phong_exp_rel_err=%f,img_rmse=%f" % (it, phong_rgb_rmse, phong_exp_rel_err, imgloss_val)
|
||||
print(s)
|
||||
if log_file:
|
||||
log_file.write(s + '\n')
|
||||
|
||||
# Show/save result image.
|
||||
display_image = display_interval and (it % display_interval == 0)
|
||||
save_mp4 = mp4save_interval and (it % mp4save_interval == 0)
|
||||
|
||||
if display_image or save_mp4:
|
||||
lightdir = np.asarray([.8, -1., .5, 0.0])
|
||||
lightdir = np.matmul(a_mvc, lightdir)[:3]
|
||||
lightdir /= np.linalg.norm(lightdir)
|
||||
lightdir = torch.as_tensor(lightdir, dtype=torch.float32, device='cuda')
|
||||
refl, refld, ldotr, mask = render_refl(lightdir, a_campos, a_mvp)
|
||||
color_opt = dr.texture(env_var[np.newaxis, ...], refl, uv_da=refld, filter_mode='linear-mipmap-linear', boundary_mode='cube')
|
||||
color_opt = color_opt + phong_var[:3] * torch.max(zero_tensor, ldotr) ** phong_var[3]
|
||||
color_opt = torch.where(mask, one_tensor, color_opt)
|
||||
result_image = color_opt.detach()[0].cpu().numpy()[::-1]
|
||||
if display_image:
|
||||
util.display_image(result_image, size=display_res, title='%d / %d' % (it, max_iter))
|
||||
if save_mp4:
|
||||
writer.append_data(np.clip(np.rint(result_image*255.0), 0, 255).astype(np.uint8))
|
||||
|
||||
# Done.
|
||||
if writer is not None:
|
||||
writer.close()
|
||||
if log_file:
|
||||
log_file.close()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Main function.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Environment map fitting example')
|
||||
parser.add_argument('--outdir', help='specify output directory', default='')
|
||||
parser.add_argument('--display-interval', type=int, default=0)
|
||||
parser.add_argument('--mp4save-interval', type=int, default=10)
|
||||
parser.add_argument('--max-iter', type=int, default=5000)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up logging.
|
||||
if args.outdir:
|
||||
out_dir = f'{args.outdir}/env_phong'
|
||||
print (f'Saving results under {out_dir}')
|
||||
else:
|
||||
out_dir = None
|
||||
print ('No output directory specified, not saving log or images')
|
||||
|
||||
# Run.
|
||||
fit_env_phong(
|
||||
max_iter=args.max_iter,
|
||||
log_interval=100,
|
||||
display_interval=args.display_interval,
|
||||
out_dir=out_dir,
|
||||
mp4save_interval=args.mp4save_interval,
|
||||
mp4save_fn='progress.mp4',
|
||||
)
|
||||
|
||||
# Done.
|
||||
print("Done.")
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
+290
@@ -0,0 +1,290 @@
|
||||
# Copyright (c) 2020, 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.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pathlib
|
||||
import numpy as np
|
||||
import torch
|
||||
import imageio
|
||||
|
||||
import util
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Quaternion math.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
# Unit quaternion.
|
||||
def q_unit():
|
||||
return np.asarray([1, 0, 0, 0], np.float32)
|
||||
|
||||
# Get a random normalized quaternion.
|
||||
def q_rnd():
|
||||
u, v, w = np.random.uniform(0.0, 1.0, size=[3])
|
||||
v *= 2.0 * np.pi
|
||||
w *= 2.0 * np.pi
|
||||
return np.asarray([(1.0-u)**0.5 * np.sin(v), (1.0-u)**0.5 * np.cos(v), u**0.5 * np.sin(w), u**0.5 * np.cos(w)], np.float32)
|
||||
|
||||
# Get a random quaternion from the octahedral symmetric group S_4.
|
||||
_r2 = 0.5**0.5
|
||||
_q_S4 = [[ 1.0, 0.0, 0.0, 0.0], [ 0.0, 1.0, 0.0, 0.0], [ 0.0, 0.0, 1.0, 0.0], [ 0.0, 0.0, 0.0, 1.0],
|
||||
[-0.5, 0.5, 0.5, 0.5], [-0.5,-0.5,-0.5, 0.5], [ 0.5,-0.5, 0.5, 0.5], [ 0.5, 0.5,-0.5, 0.5],
|
||||
[ 0.5, 0.5, 0.5, 0.5], [-0.5, 0.5,-0.5, 0.5], [ 0.5,-0.5,-0.5, 0.5], [-0.5,-0.5, 0.5, 0.5],
|
||||
[ _r2,-_r2, 0.0, 0.0], [ _r2, _r2, 0.0, 0.0], [ 0.0, 0.0, _r2, _r2], [ 0.0, 0.0,-_r2, _r2],
|
||||
[ 0.0, _r2, _r2, 0.0], [ _r2, 0.0, 0.0,-_r2], [ _r2, 0.0, 0.0, _r2], [ 0.0,-_r2, _r2, 0.0],
|
||||
[ _r2, 0.0, _r2, 0.0], [ 0.0, _r2, 0.0, _r2], [ _r2, 0.0,-_r2, 0.0], [ 0.0,-_r2, 0.0, _r2]]
|
||||
def q_rnd_S4():
|
||||
return np.asarray(_q_S4[np.random.randint(24)], np.float32)
|
||||
|
||||
# Quaternion slerp.
|
||||
def q_slerp(p, q, t):
|
||||
d = np.dot(p, q)
|
||||
if d < 0.0:
|
||||
q = -q
|
||||
d = -d
|
||||
if d > 0.999:
|
||||
a = p + t * (q-p)
|
||||
return a / np.linalg.norm(a)
|
||||
t0 = np.arccos(d)
|
||||
tt = t0 * t
|
||||
st = np.sin(tt)
|
||||
st0 = np.sin(t0)
|
||||
s1 = st / st0
|
||||
s0 = np.cos(tt) - d*s1
|
||||
return s0*p + s1*q
|
||||
|
||||
# Quaterion scale (slerp vs. identity quaternion).
|
||||
def q_scale(q, scl):
|
||||
return q_slerp(q_unit(), q, scl)
|
||||
|
||||
# Quaternion product.
|
||||
def q_mul(p, q):
|
||||
s1, V1 = p[0], p[1:]
|
||||
s2, V2 = q[0], q[1:]
|
||||
s = s1*s2 - np.dot(V1, V2)
|
||||
V = s1*V2 + s2*V1 + np.cross(V1, V2)
|
||||
return np.asarray([s, V[0], V[1], V[2]], np.float32)
|
||||
|
||||
# Angular difference between two quaternions in degrees.
|
||||
def q_angle_deg(p, q):
|
||||
p = p.detach().cpu().numpy()
|
||||
q = q.detach().cpu().numpy()
|
||||
d = np.abs(np.dot(p, q))
|
||||
d = min(d, 1.0)
|
||||
return np.degrees(2.0 * np.arccos(d))
|
||||
|
||||
# Quaternion product
|
||||
def q_mul_torch(p, q):
|
||||
a = p[0]*q[0] - p[1]*q[1] - p[2]*q[2] - p[3]*q[3]
|
||||
b = p[0]*q[1] + p[1]*q[0] + p[2]*q[3] - p[3]*q[2]
|
||||
c = p[0]*q[2] + p[2]*q[0] + p[3]*q[1] - p[1]*q[3]
|
||||
d = p[0]*q[3] + p[3]*q[0] + p[1]*q[2] - p[2]*q[1]
|
||||
return torch.stack([a, b, c, d])
|
||||
|
||||
# Convert quaternion to 4x4 rotation matrix.
|
||||
def q_to_mtx(q):
|
||||
r0 = torch.stack([1.0-2.0*q[1]**2 - 2.0*q[2]**2, 2.0*q[0]*q[1] - 2.0*q[2]*q[3], 2.0*q[0]*q[2] + 2.0*q[1]*q[3]])
|
||||
r1 = torch.stack([2.0*q[0]*q[1] + 2.0*q[2]*q[3], 1.0 - 2.0*q[0]**2 - 2.0*q[2]**2, 2.0*q[1]*q[2] - 2.0*q[0]*q[3]])
|
||||
r2 = torch.stack([2.0*q[0]*q[2] - 2.0*q[1]*q[3], 2.0*q[1]*q[2] + 2.0*q[0]*q[3], 1.0 - 2.0*q[0]**2 - 2.0*q[1]**2])
|
||||
rr = torch.transpose(torch.stack([r0, r1, r2]), 1, 0)
|
||||
rr = torch.cat([rr, torch.tensor([[0], [0], [0]], dtype=torch.float32).cuda()], dim=1) # Pad right column.
|
||||
rr = torch.cat([rr, torch.tensor([[0, 0, 0, 1]], dtype=torch.float32).cuda()], dim=0) # Pad bottom row.
|
||||
return rr
|
||||
|
||||
# Transform vertex positions to clip space
|
||||
def transform_pos(mtx, pos):
|
||||
t_mtx = torch.from_numpy(mtx).cuda() if isinstance(mtx, np.ndarray) else mtx
|
||||
# (x,y,z) -> (x,y,z,1)
|
||||
posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()], axis=1)
|
||||
return torch.matmul(posw, t_mtx.t())[None, ...]
|
||||
|
||||
def render(glctx, mtx, pos, pos_idx, col, col_idx, resolution: int):
|
||||
# Setup TF graph for reference.
|
||||
pos_clip = transform_pos(mtx, pos)
|
||||
rast_out, _ = dr.rasterize(glctx, pos_clip, pos_idx, resolution=[resolution, resolution])
|
||||
color , _ = dr.interpolate(col[None, ...], rast_out, col_idx)
|
||||
color = dr.antialias(color, rast_out, pos_clip, pos_idx)
|
||||
return color
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Cube pose fitter.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def fit_pose(max_iter = 10000,
|
||||
repeats = 1,
|
||||
log_interval = 10,
|
||||
display_interval = None,
|
||||
display_res = 512,
|
||||
lr_base = 0.01,
|
||||
lr_falloff = 1.0,
|
||||
nr_base = 1.0,
|
||||
nr_falloff = 1e-4,
|
||||
grad_phase_start = 0.5,
|
||||
resolution = 256,
|
||||
out_dir = None,
|
||||
log_fn = None,
|
||||
mp4save_interval = None,
|
||||
mp4save_fn = None):
|
||||
|
||||
log_file = None
|
||||
writer = None
|
||||
if out_dir:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
if log_fn:
|
||||
log_file = open(out_dir + '/' + log_fn, 'wt')
|
||||
if mp4save_interval != 0:
|
||||
writer = imageio.get_writer(f'{out_dir}/{mp4save_fn}', mode='I', fps=30, codec='libx264', bitrate='16M')
|
||||
else:
|
||||
mp4save_interval = None
|
||||
|
||||
datadir = f'{pathlib.Path(__file__).absolute().parents[1]}/data'
|
||||
with np.load(f'{datadir}/cube_p.npz') as f:
|
||||
pos_idx, pos, col_idx, col = f.values()
|
||||
print("Mesh has %d triangles and %d vertices." % (pos_idx.shape[0], pos.shape[0]))
|
||||
|
||||
# Some input geometry contains vertex positions in (N, 4) (with v[:,3]==1). Drop
|
||||
# the last column in that case.
|
||||
if pos.shape[1] == 4: pos = pos[:, 0:3]
|
||||
|
||||
# Create position/triangle index tensors
|
||||
pos_idx = torch.from_numpy(pos_idx.astype(np.int32)).cuda()
|
||||
vtx_pos = torch.from_numpy(pos.astype(np.float32)).cuda()
|
||||
col_idx = torch.from_numpy(col_idx.astype(np.int32)).cuda()
|
||||
vtx_col = torch.from_numpy(col.astype(np.float32)).cuda()
|
||||
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
|
||||
for rep in range(repeats):
|
||||
pose_target = torch.tensor(q_rnd(), device='cuda')
|
||||
pose_init = q_rnd()
|
||||
pose_opt = torch.tensor(pose_init / np.sum(pose_init**2)**0.5, dtype=torch.float32, device='cuda', requires_grad=True)
|
||||
|
||||
loss_best = np.inf
|
||||
pose_best = pose_opt.detach().clone()
|
||||
|
||||
# Modelview + projection matrix.
|
||||
mvp = torch.tensor(np.matmul(util.projection(x=0.4), util.translate(0, 0, -3.5)).astype(np.float32), device='cuda')
|
||||
|
||||
# Adam optimizer for texture with a learning rate ramp.
|
||||
optimizer = torch.optim.Adam([pose_opt], betas=(0.9, 0.999), lr=lr_base)
|
||||
# Render.
|
||||
for it in range(max_iter + 1):
|
||||
# Set learning rate.
|
||||
itf = 1.0 * it / max_iter
|
||||
nr = nr_base * nr_falloff**itf
|
||||
lr = lr_base * lr_falloff**itf
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
# Noise input.
|
||||
if itf >= grad_phase_start:
|
||||
noise = q_unit()
|
||||
else:
|
||||
noise = q_scale(q_rnd(), nr)
|
||||
noise = q_mul(noise, q_rnd_S4()) # Orientation noise.
|
||||
|
||||
# Render.
|
||||
color = render(glctx, torch.matmul(mvp, q_to_mtx(pose_target)), vtx_pos, pos_idx, vtx_col, col_idx, resolution)
|
||||
pose_total_opt = q_mul_torch(pose_opt, noise)
|
||||
mtx_total_opt = torch.matmul(mvp, q_to_mtx(pose_total_opt))
|
||||
color_opt = render(glctx, mtx_total_opt, vtx_pos, pos_idx, vtx_col, col_idx, resolution)
|
||||
|
||||
# Image-space loss.
|
||||
diff = (color_opt - color)**2 # L2 norm.
|
||||
diff = torch.tanh(5.0 * torch.max(diff, dim=-1)[0])
|
||||
loss = torch.mean(diff)
|
||||
|
||||
# Measure image-space loss and update best found pose.
|
||||
loss_val = float(loss)
|
||||
if (loss_val < loss_best) and (loss_val > 0.0):
|
||||
pose_best = pose_total_opt.detach().clone()
|
||||
loss_best = loss_val
|
||||
if itf < grad_phase_start:
|
||||
with torch.no_grad(): pose_opt[:] = pose_best
|
||||
|
||||
# Print/save log.
|
||||
if log_interval and (it % log_interval == 0):
|
||||
err = q_angle_deg(pose_opt, pose_target)
|
||||
ebest = q_angle_deg(pose_best, pose_target)
|
||||
s = "rep=%d,iter=%d,err=%f,err_best=%f,loss=%f,loss_best=%f,lr=%f,nr=%f" % (rep, it, err, ebest, loss_val, loss_best, lr, nr)
|
||||
print(s)
|
||||
if log_file:
|
||||
log_file.write(s + "\n")
|
||||
|
||||
# Run gradient training step.
|
||||
if itf >= grad_phase_start:
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
with torch.no_grad():
|
||||
pose_opt /= torch.sum(pose_opt**2)**0.5
|
||||
|
||||
# Show/save image.
|
||||
display_image = display_interval and (it % display_interval == 0)
|
||||
save_mp4 = mp4save_interval and (it % mp4save_interval == 0)
|
||||
|
||||
if display_image or save_mp4:
|
||||
img_ref = color[0].detach().cpu().numpy()
|
||||
img_opt = color_opt[0].detach().cpu().numpy()
|
||||
img_best = render(glctx, torch.matmul(mvp, q_to_mtx(pose_best)), vtx_pos, pos_idx, vtx_col, col_idx, resolution)[0].detach().cpu().numpy()
|
||||
result_image = np.concatenate([img_ref, img_best, img_opt], axis=1)[::-1]
|
||||
|
||||
if display_image:
|
||||
util.display_image(result_image, size=display_res, title='(%d) %d / %d' % (rep, it, max_iter))
|
||||
if save_mp4:
|
||||
writer.append_data(np.clip(np.rint(result_image*255.0), 0, 255).astype(np.uint8))
|
||||
|
||||
# Done.
|
||||
if writer is not None:
|
||||
writer.close()
|
||||
if log_file:
|
||||
log_file.close()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Cube pose fitting example')
|
||||
parser.add_argument('--outdir', help='specify output directory', default='')
|
||||
parser.add_argument('--display-interval', type=int, default=0)
|
||||
parser.add_argument('--mp4save-interval', type=int, default=10)
|
||||
parser.add_argument('--max-iter', type=int, default=1000)
|
||||
parser.add_argument('--repeats', type=int, default=1)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up logging.
|
||||
if args.outdir:
|
||||
out_dir = f'{args.outdir}/pose'
|
||||
print (f'Saving results under {out_dir}')
|
||||
else:
|
||||
out_dir = None
|
||||
print ('No output directory specified, not saving log or images')
|
||||
|
||||
# Run.
|
||||
fit_pose(
|
||||
max_iter=args.max_iter,
|
||||
repeats=args.repeats,
|
||||
log_interval=100,
|
||||
display_interval=args.display_interval,
|
||||
out_dir=out_dir,
|
||||
log_fn='log.txt',
|
||||
mp4save_interval=args.mp4save_interval,
|
||||
mp4save_fn='progress.mp4',
|
||||
)
|
||||
|
||||
# Done.
|
||||
print("Done.")
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -0,0 +1,30 @@
|
||||
# Copyright (c) 2020, 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.
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
def tensor(*args, **kwargs):
|
||||
return torch.tensor(*args, device='cuda', **kwargs)
|
||||
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
|
||||
pos = tensor([[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]], dtype=torch.float32)
|
||||
col = tensor([[[1, 0, 0], [0, 1, 0], [0, 0, 1]]], dtype=torch.float32)
|
||||
tri = tensor([[0, 1, 2]], dtype=torch.int32)
|
||||
|
||||
rast, _ = dr.rasterize(glctx, pos, tri, resolution=[256, 256])
|
||||
out, _ = dr.interpolate(col, rast, tri)
|
||||
|
||||
img = out.cpu().numpy()[0, ::-1, :, :] # Flip vertically.
|
||||
img = np.clip(np.rint(img * 255), 0, 255).astype(np.uint8) # Quantize to np.uint8
|
||||
|
||||
print("Saving to 'tri.png'.")
|
||||
imageio.imsave('tri.png', img)
|
||||
+120
@@ -0,0 +1,120 @@
|
||||
# Copyright (c) 2020, 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.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Projection and transformation matrix helpers.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def projection(x=0.1, n=1.0, f=50.0):
|
||||
return np.array([[n/x, 0, 0, 0],
|
||||
[ 0, n/x, 0, 0],
|
||||
[ 0, 0, -(f+n)/(f-n), -(2*f*n)/(f-n)],
|
||||
[ 0, 0, -1, 0]]).astype(np.float32)
|
||||
|
||||
def translate(x, y, z):
|
||||
return np.array([[1, 0, 0, x],
|
||||
[0, 1, 0, y],
|
||||
[0, 0, 1, z],
|
||||
[0, 0, 0, 1]]).astype(np.float32)
|
||||
|
||||
def rotate_x(a):
|
||||
s, c = np.sin(a), np.cos(a)
|
||||
return np.array([[1, 0, 0, 0],
|
||||
[0, c, -s, 0],
|
||||
[0, s, c, 0],
|
||||
[0, 0, 0, 1]]).astype(np.float32)
|
||||
|
||||
def rotate_y(a):
|
||||
s, c = np.sin(a), np.cos(a)
|
||||
return np.array([[ c, 0, s, 0],
|
||||
[ 0, 1, 0, 0],
|
||||
[-s, 0, c, 0],
|
||||
[ 0, 0, 0, 1]]).astype(np.float32)
|
||||
|
||||
def random_rotation_translation(t):
|
||||
m = np.random.normal(size=[3, 3])
|
||||
m[1] = np.cross(m[0], m[2])
|
||||
m[2] = np.cross(m[0], m[1])
|
||||
m = m / np.linalg.norm(m, axis=1, keepdims=True)
|
||||
m = np.pad(m, [[0, 1], [0, 1]], mode='constant')
|
||||
m[3, 3] = 1.0
|
||||
m[:3, 3] = np.random.uniform(-t, t, size=[3])
|
||||
return m
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Bilinear downsample by 2x.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def bilinear_downsample(x):
|
||||
w = torch.tensor([[1, 3, 3, 1], [3, 9, 9, 3], [3, 9, 9, 3], [1, 3, 3, 1]], dtype=torch.float32, device=x.device) / 64.0
|
||||
w = w.expand(x.shape[-1], 1, 4, 4)
|
||||
x = torch.nn.functional.conv2d(x.permute(0, 3, 1, 2), w, padding=1, stride=2, groups=x.shape[-1])
|
||||
return x.permute(0, 2, 3, 1)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Image display function using OpenGL.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
_glfw_window = None
|
||||
def display_image(image, zoom=None, size=None, title=None): # HWC
|
||||
# Import OpenGL and glfw.
|
||||
import OpenGL.GL as gl
|
||||
import glfw
|
||||
|
||||
# Zoom image if requested.
|
||||
image = np.asarray(image)
|
||||
if size is not None:
|
||||
assert zoom is None
|
||||
zoom = max(1, size // image.shape[0])
|
||||
if zoom is not None:
|
||||
image = image.repeat(zoom, axis=0).repeat(zoom, axis=1)
|
||||
height, width, channels = image.shape
|
||||
|
||||
# Initialize window.
|
||||
if title is None:
|
||||
title = 'Debug window'
|
||||
global _glfw_window
|
||||
if _glfw_window is None:
|
||||
glfw.init()
|
||||
_glfw_window = glfw.create_window(width, height, title, None, None)
|
||||
glfw.make_context_current(_glfw_window)
|
||||
glfw.show_window(_glfw_window)
|
||||
glfw.swap_interval(0)
|
||||
else:
|
||||
glfw.make_context_current(_glfw_window)
|
||||
glfw.set_window_title(_glfw_window, title)
|
||||
glfw.set_window_size(_glfw_window, width, height)
|
||||
|
||||
# Update window.
|
||||
glfw.poll_events()
|
||||
gl.glClearColor(0, 0, 0, 1)
|
||||
gl.glClear(gl.GL_COLOR_BUFFER_BIT)
|
||||
gl.glWindowPos2f(0, 0)
|
||||
gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
|
||||
gl_format = {3: gl.GL_RGB, 2: gl.GL_RG, 1: gl.GL_LUMINANCE}[channels]
|
||||
gl_dtype = {'uint8': gl.GL_UNSIGNED_BYTE, 'float32': gl.GL_FLOAT}[image.dtype.name]
|
||||
gl.glDrawPixels(width, height, gl_format, gl_dtype, image[::-1])
|
||||
glfw.swap_buffers(_glfw_window)
|
||||
if glfw.window_should_close(_glfw_window):
|
||||
return False
|
||||
return True
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Image save helper.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def save_image(fn, x):
|
||||
import imageio
|
||||
x = np.rint(x * 255.0)
|
||||
x = np.clip(x, 0, 255).astype(np.uint8)
|
||||
imageio.imsave(fn, x)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
Reference in New Issue
Block a user