Initialization in texture sampling

xmeng525 commented 3 years ago

Hi there, I found that there is no initialization for:

"texture_k" in forward_sample_texture()

https://github.com/ShichenLiu/SoftRas/blob/85329af997629d705a2e7be8fb56862323354ba7/soft_renderer/cuda/soft_rasterize_cuda_kernel.cu#L180

"grad_texture_k" in backward_sample_texture()

https://github.com/ShichenLiu/SoftRas/blob/85329af997629d705a2e7be8fb56862323354ba7/soft_renderer/cuda/soft_rasterize_cuda_kernel.cu#L199

Would you please add initialization for these variables? Thanks! :) @chenweikai

chenweikai commented 3 years ago

Hi there, I found that there is no initialization for:

"texture_k" in forward_sample_texture()

https://github.com/ShichenLiu/SoftRas/blob/85329af997629d705a2e7be8fb56862323354ba7/soft_renderer/cuda/soft_rasterize_cuda_kernel.cu#L180

"grad_texture_k" in backward_sample_texture()

https://github.com/ShichenLiu/SoftRas/blob/85329af997629d705a2e7be8fb56862323354ba7/soft_renderer/cuda/soft_rasterize_cuda_kernel.cu#L199

Would you please add initialization for these variables? Thanks! :) @chenweikai

Thank you Xiaoxu for pointing it out! @ShichenLiu Could you look into it? We found that the lack of initialization would lead to incorrect texture generation in some cases.

ShichenLiu commented 3 years ago

Hi guys,

I just fixed this issue. Not entirely sure if it addresses your problems. But could you give some details or intuitions on why this would lead to a issue? @xmeng525 Thanks!

chenweikai commented 3 years ago

Hi Xiaoxu @xmeng525 , could you provide more details on the bug so that we can confirm the correctness of our code? Thanks a lot!

xmeng525 commented 3 years ago

Hi Shichen and Weikai, I use Softras for texture regression for a cube. Given target image like: target The regression process without initialization: deform The regression process with initialization: deform

Here's the code I used (100 lines):

"""
Demo texture regression with/without init. 
Xiaoxu Meng (xiaoxumeng1993@gmail.com)

"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import tqdm
import numpy as np
import imageio
import argparse

from PIL import Image
import soft_renderer as sr

current_dir = os.path.dirname(os.path.realpath(__file__))
data_dir = os.path.join(current_dir, '../data')

class Model(nn.Module):
    def __init__(self, template_path):
        super(Model, self).__init__()

        # set template mesh
        self.template_mesh = sr.Mesh.from_obj(template_path, load_texture=True, texture_res=4)
        self.register_parameter('textures', nn.Parameter(torch.zeros_like(self.template_mesh.textures)))

    def forward(self):
        textures = self.textures
        return sr.Mesh(self.template_mesh.vertices.clone().repeat(1, 1, 1), 
                       self.template_mesh.faces.clone().repeat(1, 1, 1),
                       textures=textures.repeat(1, 1, 1, 1), 
                       texture_res=4, texture_type='surface')

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--filename-input', type=str, 
        default=os.path.join(data_dir, 'obj/cube/cube2.obj'))
    parser.add_argument('-o', '--output-dir', type=str, 
        default=os.path.join(data_dir, 'results/output_regress_cube'))
    args = parser.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    image_size = 224
    camera_distance = 2.732
    elevation = 30
    azimuth = 0

    transform = sr.LookAt()
    lighting = sr.Lighting()
    rasterizer = sr.SoftRasterizer(image_size=image_size)
    ############################################

    model = Model(args.filename_input).cuda()
    optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99))

    # create renderer with SoftRas
    transform.set_eyes_from_angles(camera_distance, elevation, azimuth)
    model.template_mesh = transform(model.template_mesh)

    vis_textures = model.template_mesh.textures.detach().cpu().numpy()
    images = rasterizer(model.template_mesh)
    image_target = images[0,0:3,:,:].permute(1, 2, 0)
    imageio.imsave(os.path.join(args.output_dir, 'target.png'),\
        (255*image_target.detach().cpu().numpy()).astype(np.uint8))

    losses = []
    loop = tqdm.tqdm(list(range(0, 1000)))
    writer = imageio.get_writer(os.path.join(args.output_dir, 'deform.gif'), mode='I')
    if not os.path.exists(os.path.join(args.output_dir, 'steps')):
        os.mkdir(os.path.join(args.output_dir, 'steps'))

    for i in loop:
        optimizer.zero_grad()
        mesh = model()
        image_result = rasterizer(mesh)[0,0:3,:,:].permute(1, 2, 0)
        loss = torch.mean(((image_result - image_target)*255.0)**2)
        loss.backward(retain_graph=True)
        optimizer.step()

        losses.append(loss.item())

        if i %50 == 0:
            loop.set_description('Loss: %.8f'%(loss.item()))
            image = image_result.detach().cpu().numpy()
            writer.append_data((255*image).astype(np.uint8))
            imageio.imsave(os.path.join(args.output_dir, 'steps', 'regress_%05d.png'%i), \
                (255*image).astype(np.uint8))

        if loss.item() < 0.01:
            break

    writer.close()
    plt.plot(losses)
    plt.show()
    mesh.save_obj(os.path.join(args.output_dir, 'saved_cube.obj'), save_texture=True)

if __name__ == '__main__':
    main()

Envirionment: Windows10 with Anaconda Cuda 11.1 python 3.6.12 torch 1.8.0+cu111 torchaudio 0.8.0 torchvision 0.9.0+cu111

chenweikai commented 3 years ago

Hi Shichen and Weikai, I use Softras for texture regression for a cube. Given target image like: target The regression process without initialization: deform The regression process with initialization: deform

Here's the code I used (100 lines):

"""
Demo texture regression with/without init. 
Xiaoxu Meng (xiaoxumeng1993@gmail.com)

"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import tqdm
import numpy as np
import imageio
import argparse

from PIL import Image
import soft_renderer as sr

current_dir = os.path.dirname(os.path.realpath(__file__))
data_dir = os.path.join(current_dir, '../data')

class Model(nn.Module):
  def __init__(self, template_path):
      super(Model, self).__init__()

      # set template mesh
      self.template_mesh = sr.Mesh.from_obj(template_path, load_texture=True, texture_res=4)
      self.register_parameter('textures', nn.Parameter(torch.zeros_like(self.template_mesh.textures)))

  def forward(self):
      textures = self.textures
      return sr.Mesh(self.template_mesh.vertices.clone().repeat(1, 1, 1), 
                       self.template_mesh.faces.clone().repeat(1, 1, 1),
                       textures=textures.repeat(1, 1, 1, 1), 
                       texture_res=4, texture_type='surface')

def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('-i', '--filename-input', type=str, 
      default=os.path.join(data_dir, 'obj/cube/cube2.obj'))
  parser.add_argument('-o', '--output-dir', type=str, 
      default=os.path.join(data_dir, 'results/output_regress_cube'))
  args = parser.parse_args()
  os.makedirs(args.output_dir, exist_ok=True)

  image_size = 224
  camera_distance = 2.732
  elevation = 30
  azimuth = 0

  transform = sr.LookAt()
  lighting = sr.Lighting()
  rasterizer = sr.SoftRasterizer(image_size=image_size)
  ############################################

  model = Model(args.filename_input).cuda()
  optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99))

  # create renderer with SoftRas
  transform.set_eyes_from_angles(camera_distance, elevation, azimuth)
  model.template_mesh = transform(model.template_mesh)

  vis_textures = model.template_mesh.textures.detach().cpu().numpy()
  images = rasterizer(model.template_mesh)
  image_target = images[0,0:3,:,:].permute(1, 2, 0)
  imageio.imsave(os.path.join(args.output_dir, 'target.png'),\
      (255*image_target.detach().cpu().numpy()).astype(np.uint8))

  losses = []
  loop = tqdm.tqdm(list(range(0, 1000)))
  writer = imageio.get_writer(os.path.join(args.output_dir, 'deform.gif'), mode='I')
  if not os.path.exists(os.path.join(args.output_dir, 'steps')):
      os.mkdir(os.path.join(args.output_dir, 'steps'))

  for i in loop:
      optimizer.zero_grad()
      mesh = model()
      image_result = rasterizer(mesh)[0,0:3,:,:].permute(1, 2, 0)
      loss = torch.mean(((image_result - image_target)*255.0)**2)
      loss.backward(retain_graph=True)
      optimizer.step()

      losses.append(loss.item())

      if i %50 == 0:
          loop.set_description('Loss: %.8f'%(loss.item()))
          image = image_result.detach().cpu().numpy()
          writer.append_data((255*image).astype(np.uint8))
          imageio.imsave(os.path.join(args.output_dir, 'steps', 'regress_%05d.png'%i), \
              (255*image).astype(np.uint8))

      if loss.item() < 0.01:
          break

  writer.close()
  plt.plot(losses)
  plt.show()
  mesh.save_obj(os.path.join(args.output_dir, 'saved_cube.obj'), save_texture=True)

if __name__ == '__main__':
  main()

Envirionment: Windows10 with Anaconda Cuda 11.1 python 3.6.12 torch 1.8.0+cu111 torchaudio 0.8.0 torchvision 0.9.0+cu111

Thank you so much for the detailed information!:-) Really appreciate it!

ShichenLiu / SoftRas

Initialization in texture sampling #88