uci-rendering / psdr-cuda

Path-space differentiable renderer
BSD 3-Clause "New" or "Revised" License
155 stars 11 forks source link

Optimizing vertex positions #4

Closed hyeonjang closed 2 years ago

hyeonjang commented 2 years ago

Hello,

I am now trying to test the optimization process w.r.t. vertex positions and BSDF, and the python code is like below.

import enoki as ek
from enoki.cuda_autodiff import Float32 as FloatD, Vector3f as Vector3fD, Matrix4f as Matrix4fD
import psdr_cuda

from utils.adam import Adam

import imageio

# Load the scene without configuring it
scene = psdr_cuda.Scene()
scene.load_file('data/scenes/cbox_bunny.xml', auto_configure=True)

# target rendering
image_ref = psdr_cuda.DirectIntegrator().renderD(scene, sensor_id=0)
np_image_ref = image_ref.numpy().reshape(scene.opts.height, scene.opts.width, 3)
imageio.imwrite(f"./tmp/out_ref.exr", np_image_ref)
del scene

scene = psdr_cuda.Scene()
scene.load_file('data/scenes/cbox_bunny.xml', auto_configure=False)

# Construct a rotation matrix
P = FloatD(1.)
mat = Matrix4fD.rotate(Vector3fD(0., 0.0, 1.), P)

# Left-multiply this matrix to the to-world transformation of the first mesh of the scene
scene.param_map["Mesh[1]"].append_transform(mat)
ek.set_requires_gradient(scene.param_map["Mesh[1]"].vertex_positions)

# Configure the scene
scene.configure()

# optimization process
opt = Adam(scene.param_map, scene.param_map, ["BSDF[0]"], ["Mesh[1]"], lr=0.2)
for i in range(20):
    image  = psdr_cuda.DirectIntegrator().renderD(scene, sensor_id=0)

    # image writing
    np_image = image.numpy().reshape(scene.opts.height, scene.opts.width, 3)
    imageio.imwrite(f"./tmp/out_{i}.exr", np_image)

    loss = ek.hmean(ek.squared_norm(image_ref - image))

    # Backward Process
    ek.backward(loss)
    opt.step()

But this optimization process works for BSDF only.
What did I miss?

Can you help me please?

andyyankai commented 2 years ago

Are you trying to rotate the mesh?

hyeonjang commented 2 years ago

Yes, I am try to rotate mesh. However, in the end, I want to also move vertex positions individually, as shown as EGSR 2021 "Unified shape and SVBRDF recovery using differentiable Monte Carlo rendering" paper

shuangz25 commented 2 years ago

As demonstrated in our EGSR 2021 paper, optimizing vertex position in a robust way is nontrivial: You need, for example, proper regularizations of the mesh (e.g., by minimizing the Laplacian). Currently, psdr-cuda is meant to provide a differentiable rendering tool, not a full inverse-rendering optimization pipeline.

hyeonjang commented 2 years ago

Thank you for replying!

Yes, I know that the mesh regularizers takes a big role through your EGSR 2021 paper. It showed the awesome results! I left the question because any changes didn't occur about the mesh in my toy example code (about rotation matrix or vertex positions). Therefore, I guess that I am missing some parts.

Please, I wish to get an advice.

andyyankai commented 2 years ago

Let me share an example for rotation, hope this can help. I can check your code in detail later as well:

def psdr_optimize():
    print("psdr_optimize!")
    sc = psdr_cuda.Scene()
    sc.load_file(scene_file, False)
    sc.param_map["Mesh[1]"].enable_edges = False
    ro = sc.opts
    ro.spp = 16
    ro.sppe = 16
    ro.sppse = 16
    sc.opts.log_level = 0
    num_sensors = sc.num_sensors
    sc.configure()
    axis = Vector3fD(0.0, 0.0, 1.0)
    valid_camera = 0
    print('Loading target images')
    tars = []
    target = cv2.imread(target_path+"exr/sensor_%d.exr" % valid_camera, cv2.IMREAD_UNCHANGED)
    target = torch.from_numpy(cv2.cvtColor(target, cv2.COLOR_RGB2BGR)).float()
    target = target.reshape((-1, 3))
    tars.append(target)'

   class PSDRRender(torch.autograd.Function):
        @staticmethod 
        def forward(ctx, A, iter):
            _dir_vector = FloatD(A)
            ek.set_requires_gradient(_dir_vector, A.requires_grad)
            ctx.in1 = _dir_vector
            sc.param_map['Mesh[0]'].set_transform(Matrix4fD.rotate(axis, _dir_vector))
            sc.configure()

            integrator = psdr_cuda.DirectIntegrator(1, 1, 0)
            img = integrator.renderD(sc, valid_camera)
            tar = Vector3fD(tars[0].cuda())

            loss = ek.abs(img - tar)
            loss = ek.hsum(ek.hsum(loss)) / (ro.height * ro.width * 3)
            ctx.out = loss 
            out_torch = ctx.out.torch() 
            return out_torch

        @staticmethod
        def backward(ctx, grad_out):
            ek.set_gradient(ctx.out, FloatC(grad_out))
            FloatD.backward() 
            gradA = ek.gradient(ctx.in1).torch()
            result = gradA
            del ctx.out, ctx.in1
            return (result, None)

    P = FloatD(rotate_value)
    A = Variable(P.torch(), requires_grad=True)
    params = [
        {'params': A, 'lr':  0.03},
    ]
    optimizer = torch.optim.Adam(params)

    render = PSDRRender.apply
    npass = 100
    print('Starting optimization')
    loss_graph = np.zeros(npass)
    para_result = np.array([rotate_value])
    para_graph = np.zeros(npass)

    for i in range(npass):
        start = time.time()

        average_loss = 0
        average_grad = 0
        optimizer.zero_grad()
        gpass = 12
        for g in range(gpass):
            loss_temp = render(A, g)
            loss_temp.backward()
            average_loss += loss_temp.item() / float(gpass)
            average_grad += A.grad / float(gpass)
        A.grad = average_grad
        optimizer.step()
        end = time.time()

        loss_graph[i] = average_loss
        curr_para = A.cpu().detach().numpy()[0]
        para_graph[i] = abs(curr_para)
        print('i:', i, ' loss:', average_loss, ' para_diff:', curr_para, ' Time:', end - start) 
        torch.cuda.empty_cache() 
        ek.cuda_malloc_trim()

    image_loss(npass, loss_graph, output_dir)
    para_loss(npass, para_graph, output_dir)
hyeonjang commented 2 years ago

Thanks a lot for your sharing! From your code, I could have figured out the problem in my code. I didn't call scene.configure() function after a iteration. So, the mesh buffer didn't be updated.

xdobetter commented 2 years ago

Thanks a lot for your sharing! From your code, I could have figured out the problem in my code. I didn't call scene.configure() function after a iteration. So, the mesh buffer didn't be updated.

Sorry to bother you, I'm also using psdr-cuda for mesh optimization recently, I refer to the code you shared and your latest suggestion ( call scene.configure() function after a iteration) but it doesn't work, Could you share your corrected code?Thanks.