def __init__(self, filename_obj, cfg): super(Reconstructor, self).__init__() feature_dim = cfg['Reconstruction']['f_dim'] img_size = cfg['input_size'] if cfg['Reconstruction']['encoder_type'] == 'resnet': self.encoder = ResnetEncoder(c_dim=feature_dim) elif ['Reconstruction']['encoder_type'] == 'simple': self.encoder = SimpleEncoder(dim_out=feature_dim, im_size=img_size) else: raise Exception('encoder type must be resnet or simple') self.color_rec = cfg['Reconstruction']['color_rec'] self.decoder = Decoder(filename_obj, self.color_rec, color_num=cfg['Reconstruction']['color_num']) if self.color_rec: self.renderer = sr.SoftRenderer( image_size=img_size, sigma_val=cfg['SoftRender']['SIGMA_VAL'], aggr_func_rgb='softmax', camera_mode='look_at', viewing_angle=15, dist_eps=1e-10) else: self.renderer = sr.SoftRenderer( image_size=img_size, sigma_val=cfg['SoftRender']['SIGMA_VAL'], aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=15, dist_eps=1e-10) self.laplacian_loss = sr.LaplacianLoss(self.decoder.vertices_base, self.decoder.faces) self.flatten_loss = sr.FlattenLoss(self.decoder.faces)
def __init__(self, img_size=256, render_type='softmax', background_color=[0, 0, 0], sigma_val=1e-5, gamma_val=1e-4, dist_eps=1e-10, anti_aliasing=True): super(SoftRenderer, self).__init__() self.renderer = sr.SoftRenderer(image_size=img_size, aggr_func_rgb=render_type, camera_mode='look_at', sigma_val=sigma_val, dist_eps=dist_eps, gamma_val=gamma_val, background_color=background_color, anti_aliasing=anti_aliasing, perspective=False) # Set a default camera to be at (0, 0, -2.732) self.renderer.transform.transformer._eye = [0, 0, -2.732] # Make it a bit brighter for vis self.renderer.lighting.ambient.light_intensity = 0.8 self.proj_fn = geom_utils.orthographic_proj_withz self.offset_z = 5.
def setup_renderer(): # renderer = sr.SoftRenderer(camera_mode="look", viewing_scale=2 / res, far=10000, perspective=False, image_size=res, # camera_direction=[0, 0, -1], camera_up=[0, 1, 0], light_intensity_ambient=1) # renderer.transform.set_eyes([res / 2, res / 2, 6000]) renderer = sr.SoftRenderer(camera_mode="look_at") return renderer
def __init__(self, filename_obj, args): super(Model, self).__init__() self.encoder = VPLEncoder() self.decoder = VPLDecoder(filename_obj) self.renderer = sr.SoftRenderer(image_size=args.image_size, sigma_val=args.sigma_val, aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=15, dist_eps=1e-10) self.laplacian_loss = sr.LaplacianLoss(self.decoder.vertices_base, self.decoder.faces) self.flatten_loss = sr.FlattenLoss(self.decoder.faces)
def main(): parser = argparse.ArgumentParser() parser.add_argument('-i', '--filename-input', type=str, default=os.path.join(data_dir, 'obj/spot/spot_triangulated.obj')) parser.add_argument('-o', '--output-dir', type=str, default=os.path.join(data_dir, 'results/output_render')) args = parser.parse_args() # other settings camera_distance = 2.732 elevation = 30 azimuth = 0 # load from Wavefront .obj file mesh = sr.Mesh.from_obj(args.filename_input, load_texture=True, texture_res=5, texture_type='surface') # create renderer with SoftRas renderer = sr.SoftRenderer(camera_mode='look_at') os.makedirs(args.output_dir, exist_ok=True) # draw object from different view loop = tqdm.tqdm(list(range(0, 360, 4))) writer = imageio.get_writer(os.path.join(args.output_dir, 'rotation.gif'), mode='I') for num, azimuth in enumerate(loop): # rest mesh to initial state mesh.reset_() loop.set_description('Drawing rotation') renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) images = renderer.render_mesh(mesh) image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255*image).astype(np.uint8)) writer.close() # draw object from different sigma and gamma loop = tqdm.tqdm(list(np.arange(-4, -2, 0.2))) renderer.transform.set_eyes_from_angles(camera_distance, elevation, 45) writer = imageio.get_writer(os.path.join(args.output_dir, 'bluring.gif'), mode='I') for num, gamma_pow in enumerate(loop): # rest mesh to initial state mesh.reset_() renderer.set_gamma(10**gamma_pow) renderer.set_sigma(10**(gamma_pow - 1)) loop.set_description('Drawing blurring') images = renderer.render_mesh(mesh) image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255*image).astype(np.uint8)) writer.close() # save to textured obj mesh.reset_() mesh.save_obj(os.path.join(args.output_dir, 'saved_spot.obj'), save_texture=True)
def main(): filename_input = os.path.join(data_dir, 'banana.obj') filename_output = os.path.join(output_directory, 'example1.gif') ########################### # camera settings ########################### camera_distance = 2 elevation = 30 ########################### # load object ########################### mesh = TriangleMesh.from_obj(filename_input) vertices = mesh.vertices faces = mesh.faces.int() face_textures = (faces).clone() vertices = vertices[None, :, :].cuda() faces = faces[None, :, :].cuda() face_textures[None, :, :].cuda() ########################### # normalize verts ########################### vertices_max = vertices.max() vertices_min = vertices.min() vertices_middle = (vertices_max + vertices_min) / 2. vertices = vertices - vertices_middle coef = 5 vertices = vertices * coef ########################### # Soft Rasterizer ########################### textures = torch.ones(1, faces.shape[1], 2, 3, dtype=torch.float32).cuda() mesh = sr.Mesh(vertices, faces, textures) renderer = sr.SoftRenderer(camera_mode='look_at') loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing SR') writer = imageio.get_writer(os.path.join(output_directory_sr, 'rotation.gif'), mode='I') for azimuth in loop: mesh.reset_() renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) images = renderer.render_mesh(mesh) image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) writer.close()
def __init__(self, filename_obj, cfg): super(Unconditional_Generator, self).__init__() z_dim = cfg['NormalGan']['G']['z_dim'] self.decoder = Decoder(filename_obj, color_rec=False, dim_in=z_dim) self.renderer = sr.SoftRenderer( image_size=cfg['input_size'], sigma_val=cfg['SoftRender']['SIGMA_VAL'], aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=15, dist_eps=1e-10) self.laplacian_loss = sr.LaplacianLoss(self.decoder.vertices_base, self.decoder.faces) self.flatten_loss = sr.FlattenLoss(self.decoder.faces)
def __init__(self, args): super(VSLModel, self).__init__() self.batch_size = args.batch_size self.inf_model = InfModel(args=args) self.gen_model = GenModel(args=args) self.image_decoder = ImageDecoder(args=args) self.renderer = sr.SoftRenderer(image_size=args.image_size, sigma_val=args.sigma_val, aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=15, dist_eps=1e-10) self.laplacian_loss = sr.LaplacianLoss(self.gen_model.vertices_base, self.gen_model.faces) self.flatten_loss = sr.FlattenLoss(self.gen_model.faces)
def __init__(self, img_size=256, perspective=True, **kwargs): super(SoftRas, self).__init__() import soft_renderer self.renderer = soft_renderer.SoftRenderer(image_size=img_size, camera_mode='look_at', perspective=perspective, eye=[0, 0, -2.732], **kwargs) self.renderer = self.renderer.cuda() self.viewing_angle = 30 self.perspective = perspective self.eye = [0, 0, -2.732] self.proj_fn = geom_utils.orthographic_proj_withz self.offset_z = 5.
def main(): parser = argparse.ArgumentParser() parser.add_argument('-a', '--cls', type=str, default='02958343') parser.add_argument('-b', '--id', type=str, default='') args = parser.parse_args() args.id = args.id.split('/')[-1] root = '/home/grisw/Downloads/ShapeNetCore.v2' output_root = os.path.join(data_dir, 'dataset_256_36_1') # create renderer with SoftRas renderer = sr.SoftRenderer(image_size=256, camera_mode='look_at') os.makedirs(os.path.join(output_root, args.cls), exist_ok=True) orientations = generate( os.path.join(root, args.cls, args.id, 'models', 'model_normalized.obj'), renderer) np.savez_compressed(os.path.join(output_root, args.cls, args.id), arr_0=np.array(orientations))
def render_mesh_rt(vertices, triangles, colors, rot=None, tran=None): ''' Apply the rt(rotation and translation) to a mesh then render it to an image Args: vertices: torch.Tensor, shape = [nver, 3] (on gpu) triangles: torch.Tensor, shape = [ntri, 3] (on gpu) colors: torch.Tensor, shape = [nver, 3] (on gpu) r: torch.Tensor, shape=[3,3] (on gpu) t: torch.Tensor, shape=[3,1] (on gpu) Return: image: torch.Tensor, shape=[4, width, height] (on gpu) ''' # vertices = torch.Tensor(vertices).cuda() # triangles = torch.Tensor(triangles).cuda().int() # colors = torch.Tensor(colors).cuda() # rot = torch.Tensor(rot).cuda() # tran = torch.Tensor(trans).cuda() # apply r and t if type(rot) == torch.Tensor and type(tran) == torch.Tensor: # new_vertices = (rot @ vertices.transpose(1,0) + tran).transpose(1,0) new_vertices = (rot.T @ (vertices.transpose(1,0) - tran)).transpose(1,0) else: new_vertices = vertices # prepare the mesh face_mesh = sr.Mesh(new_vertices, triangles, colors, texture_type="vertex") # prepare the renderer # renderer = sr.SoftRenderer(camera_mode="look_at", far=10000, image_size=512, camera_direction=[0,1,0], light_directions=[0,0,1], dist_func="hard") # the dist_func can be 'hard', 'barycentric' or 'euclidean' # renderer.transform.set_eyes_from_angles(-400, 0, 0) # this three number can be trainable parameters as well renderer = sr.SoftRenderer(camera_mode="look", far=10000, image_size=224, viewing_angle=15, camera_direction=[0,0,-1], camera_up=[-0.3,1,0], light_intensity_ambient=1, light_color_ambient=[1,1,1], dist_func="hard") renderer.transform.set_eyes([0,-50,680]) # do the rendering images = renderer.render_mesh(face_mesh) image = images[0] image = torch.flip(image, [2]) return image
os.environ["CUDA_VISIBLE_DEVICES"] = "7" RENDER_IMAGE_NAME_RGB = 'RGB' RENDER_IMAGE_NAME_D = 'depth' RENDER_IMAGE_NAME_NORMAL = 'normal' camera_distance = 10 elevation = 30 azimuth = 0 obj_file_i = os.path.join( "/mnt/zhengwen/model_synthesis/SF_temp/data/obj/sphere/sphere_642.obj") img_file_rgb = "/mnt/zhengwen/model_synthesis/SF_temp/data/obj/sphere/sphere_642.obj" + RENDER_IMAGE_NAME_RGB img_file_depth = "/mnt/zhengwen/model_synthesis/SF_temp/data/obj/sphere/sphere_642.obj" + RENDER_IMAGE_NAME_D img_file_normal = "/mnt/zhengwen/model_synthesis/SF_temp/data/obj/sphere/sphere_642.obj" + RENDER_IMAGE_NAME_NORMAL mesh = sr.Mesh.from_obj(obj_file_i) renderer = sr.SoftRenderer(camera_mode='look_at') for azimuth in range(0, 360, 15): count = azimuth // 15 # rest mesh to initial state mesh.reset_() renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) images = renderer.render_mesh(mesh) image_rgb = images[0].detach().cpu().numpy()[0] imageio.imwrite( img_file_rgb + '_' + str(count) + '.png', (255 * image_rgb[:, 128 - 32:128 + 32, 128 - 32:128 + 32]).transpose( (1, 2, 0)).astype(np.uint8))
def main(): ########################### # camera settings ########################### camera_distance = 2.732 elevation = 0 azimuth = 0 ########################### # load object ########################### filename_input = os.path.join(data_dir, 'banana.obj') filename_ref = os.path.join(data_dir, 'example2_ref.png') image_gt = torch.from_numpy(imread(filename_ref).astype(np.float32).mean(-1) / 255.)[None, ::].cuda() mesh = TriangleMesh.from_obj(filename_input) vertices = mesh.vertices faces = mesh.faces.int() uvs = torch.FloatTensor(get_spherical_coords_x(vertices.data.numpy())) face_textures = (faces).clone() pmax = vertices.max() pmin = vertices.min() pmiddle = (pmax + pmin) / 2 vertices = vertices - pmiddle coef = 10 vertices = vertices * coef vertices = vertices[None, :, :].cuda() faces = faces[None, :, :].cuda() uvs = uvs[None, :, :].cuda() face_textures[None, :, :].cuda() ########################## # normalize verts ########################## vertices_max = vertices.max() vertices_min = vertices.min() vertices_middle = (vertices_max + vertices_min)/2. vertices = vertices - vertices_middle # coef = 5 # vertices = vertices * coef ########################### # NMR ########################### textures = torch.ones(1, faces.shape[1], 2,2,2, 3, dtype=torch.float32).cuda() model = Model(vertices.clone()).cuda() renderer = nr.Renderer(camera_mode='look_at') renderer.eye = nr.get_points_from_angles(camera_distance, elevation, azimuth) optimizer = torch.optim.Adam(model.parameters(), 0.001, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing NMR') writer = imageio.get_writer(os.path.join(output_directory_nmr, 'deform.gif'), mode='I') for i in loop: optimizer.zero_grad() new_vertices = model() image_pred ,_, _= renderer(new_vertices, faces, textures) loss = torch.sum((image_pred - image_gt[None, :, :])**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f'%(loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) other_image = image_gt.detach().cpu().numpy().transpose((1, 2, 0)) pass_image = image + other_image writer.append_data((128*pass_image).astype(np.uint8)) ########################### # Soft Rasterizer ########################### textures = torch.ones(1, faces.shape[1], 2, 3, dtype=torch.float32).cuda() model = Model(vertices.clone()).cuda() mesh = sr.Mesh(vertices, faces, textures) renderer = sr.SoftRenderer(image_size=256, sigma_val=3e-5, aggr_func_rgb='hard', camera_mode='look_at') renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) optimizer = torch.optim.Adam(model.parameters(), 0.001, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing SR') writer = imageio.get_writer(os.path.join(output_directory_sr, 'deform.gif'), mode='I') for i in loop: optimizer.zero_grad() new_vertices = model() new_mesh = sr.Mesh(new_vertices, faces, textures) image_pred = renderer.render_mesh(new_mesh) loss = torch.sum((image_pred - image_gt[None, :, :])**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f'%(loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) other_image = image_gt.detach().cpu().numpy().transpose((1, 2, 0)) pass_image = image + other_image writer.append_data((128*pass_image).astype(np.uint8)) ################################ # Dib-Renderer - Vertex Colours ################################ model = Model(vertices.clone()).cuda() textures = torch.ones(1, vertices.shape[1], 3).cuda() renderer = Dib_Renderer(256, 256, mode = 'VertexColor') renderer.set_look_at_parameters([90-azimuth], [elevation], [camera_distance]) optimizer = torch.optim.Adam(model.parameters(), 0.001, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing Dib_Renderer VertexColor') writer = imageio.get_writer(os.path.join(output_directory_dib, 'deform_VertexColor.gif'), mode='I') for i in loop: optimizer.zero_grad() new_vertices = model() image_pred, alpha, _ = renderer.forward(points=[new_vertices, faces[0].long()], colors=[textures]) image_pred = torch.cat((image_pred, alpha), dim = 3) image_pred = image_pred.permute(0,3,1,2) loss = torch.sum((image_pred - image_gt[None, :, :])**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f'%(loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) other_image = image_gt.detach().cpu().numpy().transpose((1, 2, 0)) pass_image = image + other_image writer.append_data((127*pass_image).astype(np.uint8)) ################################ # Dib-Renderer - Lambertian ################################ model = Model(vertices.clone()).cuda() textures = torch.ones(1, 3, 256, 256).cuda() renderer = Dib_Renderer(256, 256, mode = 'Lambertian') renderer.set_look_at_parameters([90-azimuth], [elevation], [camera_distance]) optimizer = torch.optim.Adam(model.parameters(), 0.001, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing Dib_Renderer Lambertian') writer = imageio.get_writer(os.path.join(output_directory_dib, 'deform_Lambertian.gif'), mode='I') for i in loop: optimizer.zero_grad() new_vertices = model() image_pred, alpha, _ = renderer.forward(points=[new_vertices, faces[0].long()], colors=[uvs, face_textures.long(), textures]) image_pred = torch.cat((image_pred, alpha), dim = 3) image_pred = image_pred.permute(0,3,1,2) loss = torch.sum((image_pred - image_gt[None, :, :])**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f'%(loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) other_image = image_gt.detach().cpu().numpy().transpose((1, 2, 0)) pass_image = image + other_image writer.append_data((127*pass_image).astype(np.uint8)) ################################ # Dib-Renderer - Phong ################################ model = Model(vertices.clone()).cuda() textures = torch.ones(1, 3, 256, 256).cuda() renderer = Dib_Renderer(256, 256, mode = 'Phong') renderer.set_look_at_parameters([90-azimuth], [elevation], [camera_distance]) optimizer = torch.optim.Adam(model.parameters(), 0.001, betas=(0.5, 0.99)) ### Lighting info ### material = np.array([[0.1, 0.1, 0.1], [1.0, 1.0, 1.0], [0.4, 0.4, 0.4]], dtype=np.float32).reshape(-1, 3, 3) material = torch.from_numpy(material).repeat(1, 1, 1).cuda() shininess = np.array([100], dtype=np.float32).reshape(-1, 1) shininess = torch.from_numpy(shininess).repeat(1, 1).cuda() lightdirect = 2 * np.random.rand(1, 3).astype(np.float32) - 1 lightdirect[:, 2] += 2 lightdirect = torch.from_numpy(lightdirect).cuda() loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing Dib_Renderer Phong') writer = imageio.get_writer(os.path.join(output_directory_dib, 'deform_Phong.gif'), mode='I') for i in loop: optimizer.zero_grad() new_vertices = model() image_pred, alpha, _ = renderer.forward(points=[new_vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), textures],\ light= lightdirect, \ material=material, \ shininess=shininess) image_pred = torch.cat((image_pred, alpha), dim = 3) image_pred = image_pred.permute(0,3,1,2) loss = torch.sum((image_pred - image_gt[None, :, :])**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f'%(loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) other_image = image_gt.detach().cpu().numpy().transpose((1, 2, 0)) pass_image = image + other_image writer.append_data((127*pass_image).astype(np.uint8)) ################################ # Dib-Renderer - SphericalHarmonics ################################ model = Model(vertices.clone()).cuda() textures = torch.ones(1, 3, 256, 256).cuda() renderer = Dib_Renderer(256, 256, mode = 'SphericalHarmonics') renderer.set_look_at_parameters([90-azimuth], [elevation], [camera_distance]) optimizer = torch.optim.Adam(model.parameters(), 0.001, betas=(0.5, 0.99)) lightparam = np.random.rand(1, 9).astype(np.float32) lightparam[:, 0] += 2 lightparam = torch.from_numpy(lightparam).cuda() loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing Dib_Renderer SH') writer = imageio.get_writer(os.path.join(output_directory_dib, 'deform_SH.gif'), mode='I') for i in loop: optimizer.zero_grad() new_vertices = model() image_pred, alpha, _ = renderer.forward(points=[new_vertices, faces[0].long()],\ colors=[uvs, face_textures.long(), textures], light =lightparam) image_pred = torch.cat((image_pred, alpha), dim = 3) image_pred = image_pred.permute(0,3,1,2) loss = torch.sum((image_pred - image_gt[None, :, :])**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f'%(loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) other_image = image_gt.detach().cpu().numpy().transpose((1, 2, 0)) pass_image = image + other_image writer.append_data((127*pass_image).astype(np.uint8))
renderer.transform.set_eyes([0,-50,680]) # do the rendering images = renderer.render_mesh(face_mesh) image = images[0] image = torch.flip(image, [2]) return image if __name__ == "__main__": face_id = "00336" mesh_file = "00025/00025.obj" rt_file = "00025/00025_sRT.npy" face_mesh = sr.Mesh.from_obj(mesh_file, load_texture=True, texture_type="vertex") renderer = sr.SoftRenderer(camera_mode="look", far=10000, image_size=224, viewing_angle=15, camera_direction=[0,0,-1], camera_up=[-0.3,1,0], light_intensity_ambient=1, light_color_ambient=[1,1,1], dist_func="hard") renderer.transform.set_eyes([0,-50,680]) # do the rendering images = renderer.render_mesh(face_mesh) image = images[0] image = torch.flip(image, [2]) image = image.detach().cpu().numpy() image = image.transpose((1,2,0)) plt.imshow(image) plt.show() """ # load mesh and rt from files to np.array vertices, triangles, colors = load_obj(mesh_file) rots, trans = recover(np.load(rt_file)) # convert them to torch.Tensor and send to gpu
def main(): filename_input = os.path.join(data_dir, 'banana.obj') filename_output = os.path.join(output_directory, 'example1.gif') ########################### # camera settings ########################### camera_distance = 2 elevation = 30 ########################### # load object ########################### mesh = TriangleMesh.from_obj(filename_input) vertices = mesh.vertices faces = mesh.faces.int() uvs = torch.FloatTensor(get_spherical_coords_x(vertices.data.numpy())) face_textures = (faces).clone() vertices = vertices[None, :, :].cuda() faces = faces[None, :, :].cuda() uvs = uvs[None, :, :].cuda() face_textures[None, :, :].cuda() ########################### # normalize verts ########################### vertices_max = vertices.max() vertices_min = vertices.min() vertices_middle = (vertices_max + vertices_min) / 2. vertices = vertices - vertices_middle coef = 5 vertices = vertices * coef ########################### # NMR ########################### textures = torch.ones(1, faces.shape[1], 2, 2, 2, 3, dtype=torch.float32).cuda() renderer = nr.Renderer(camera_mode='look_at') loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing NMR') writer = imageio.get_writer(os.path.join(output_directory_nmr, 'rotation.gif'), mode='I') for num, azimuth in enumerate(loop): renderer.eye = nr.get_points_from_angles(camera_distance, elevation, azimuth) images, _, _ = renderer(vertices, faces, textures) image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) writer.close() ########################### # Soft Rasterizer ########################### textures = torch.ones(1, faces.shape[1], 2, 3, dtype=torch.float32).cuda() mesh = sr.Mesh(vertices, faces, textures) renderer = sr.SoftRenderer(camera_mode='look_at') loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing SR') writer = imageio.get_writer(os.path.join(output_directory_sr, 'rotation.gif'), mode='I') for azimuth in loop: mesh.reset_() renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) images = renderer.render_mesh(mesh) image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) writer.close() ################################ # Dib-Renderer - Vertex Colours ################################ renderer = Dib_Renderer(256, 256, mode='VertexColor') textures = torch.ones(1, vertices.shape[1], 3).cuda() loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing Dib_Renderer VertexColor') writer = imageio.get_writer(os.path.join(output_directory_dib, 'rotation_VertexColor.gif'), mode='I') for azimuth in loop: renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) predictions, _, _ = renderer.forward( points=[vertices, faces[0].long()], colors=[textures]) image = predictions.detach().cpu().numpy()[0] writer.append_data((image * 255).astype(np.uint8)) writer.close() ################################ # Dib-Renderer - Lambertian ################################ renderer = Dib_Renderer(256, 256, mode='Lambertian') textures = torch.ones(1, 3, 256, 256).cuda() loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing Dib_Renderer Lambertian') writer = imageio.get_writer(os.path.join(output_directory_dib, 'rotation_Lambertian.gif'), mode='I') for azimuth in loop: renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) predictions, _, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), textures]) image = predictions.detach().cpu().numpy()[0] writer.append_data((image * 255).astype(np.uint8)) writer.close() ################################ # Dib-Renderer - Phong ################################ renderer = Dib_Renderer(256, 256, mode='Phong') textures = torch.ones(1, 3, 256, 256).cuda() ### Lighting info ### material = np.array([[0.1, 0.1, 0.1], [1.0, 1.0, 1.0], [0.4, 0.4, 0.4]], dtype=np.float32).reshape(-1, 3, 3) material = torch.from_numpy(material).repeat(1, 1, 1).cuda() shininess = np.array([100], dtype=np.float32).reshape(-1, 1) shininess = torch.from_numpy(shininess).repeat(1, 1).cuda() lightdirect = 2 * np.random.rand(1, 3).astype(np.float32) - 1 lightdirect[:, 2] += 2 lightdirect = torch.from_numpy(lightdirect).cuda() loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing Dib_Renderer Phong') writer = imageio.get_writer(os.path.join(output_directory_dib, 'rotation_Phong.gif'), mode='I') for azimuth in loop: renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) predictions, _, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), textures],\ light= lightdirect, \ material=material, \ shininess=shininess ) image = predictions.detach().cpu().numpy()[0] writer.append_data((image * 255).astype(np.uint8)) writer.close() ################################ # Dib-Renderer - SH ################################ renderer = Dib_Renderer(256, 256, mode='SphericalHarmonics') textures = torch.ones(1, 3, 256, 256).cuda() ### Lighting info ### lightparam = np.random.rand(1, 9).astype(np.float32) lightparam[:, 0] += 2 lightparam = torch.from_numpy(lightparam).cuda() loop = tqdm.tqdm(list(range(0, 360, 4))) loop.set_description('Drawing Dib_Renderer SH') writer = imageio.get_writer(os.path.join(output_directory_dib, 'rotation_SH.gif'), mode='I') for azimuth in loop: renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) predictions, _, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), textures],\ light=lightparam) image = predictions.detach().cpu().numpy()[0] writer.append_data((image * 255).astype(np.uint8)) writer.close()
outCols = colorInit(features) #print (outPos.shape) if createMesh: templateVertex = dataBatch['TemplVertex'].cuda(opt.gpuId) templateFaces = dataBatch['TemplFaces'].cuda(opt.gpuId) meshM = models.MeshModel(templateFaces, templateVertex).cuda(opt.gpuId) createMesh = False # TODO : calculate lap and flat loss here.. meshDeformed, lapLoss, fltLoss = meshM.forward( outPos[:, :-1, :], torch.zeros_like(outPos[:, -1:, :]).cuda(opt.gpuId), opt.numViews, currBatchSize, outCols) renderer = sr.SoftRenderer(image_size=opt.imageSize, sigma_val=1e-4, aggr_func_rgb='hard', camera_mode='projection', P=projViews, orig_size=opt.origImageSize) imagesPred = renderer.render_mesh(meshDeformed) SS = losses.SilhouetteLoss(imagesPred[:, 3], imgViews) pixelL = pixelLoss( imagesPred[:, 0:3, :, :] * (imgViews.unsqueeze(1)), (colImgViews / 255.0) * (imgViews.unsqueeze(1))) loss = lamS*SS + \ lamL*lapLoss + \ lamF*fltLoss +\ lamP*pixelL # Train net..
def main(): parser = argparse.ArgumentParser() parser.add_argument('-o', '--output-dir', type=str, default=os.path.join(data_dir, 'results/output_deform')) args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) cameras = [] # eye, direction if OPTIMIZE_SILHOUETTE: cameras.append(([-2, 0, 0], [1, 0, 0])) cameras.append(([0, 0, -2], [0, 0, 1])) else: cameras.append(([2, 0, 0], [1, 0, 0])) cameras.append(([0, 0, 2], [0, 0, 1])) # Use wider angle if placing the camera between the object and the wall angle = 30 if OPTIMIZE_SILHOUETTE else 50 model = Model(os.path.join(data_dir, 'obj/sphere/colored_sphere.obj'), os.path.join(data_dir, 'obj/bg/bg_subdiv.obj')).cuda() renderer = sr.SoftRenderer(image_size=64, sigma_val=1e-4, aggr_func_rgb='softmax', camera_mode='look', viewing_angle=angle, light_intensity_ambient=1., light_intensity_directionals=0., background_color=[1., 1., 1.], light_width=3) renderer.transform.transformer._eye = [eye for eye, _ in cameras] renderer.transform.transformer.camera_direction = [ direction for _, direction in cameras ] # read training images and camera poses images = torch.stack([ load_png(os.path.join(data_dir, 'target/t.png')), load_png(os.path.join(data_dir, 'target/f.png')) ]) optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99)) loop = tqdm.tqdm(list(range(0, 1000))) writer = imageio.get_writer(os.path.join(args.output_dir, 'deform.gif'), mode='I') images_gt = images.cuda() image1 = images_gt[:, 0].detach().cpu().numpy()[0] #.transpose((1, 2, 0)) image2 = images_gt[:, 0].detach().cpu().numpy()[1] #.transpose((1, 2, 0)) imageio.imsave(os.path.join(args.output_dir, 'gt1.png'), (255 * image1).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'gt2.png'), (255 * image2).astype(np.uint8)) for i in loop: mesh, laplacian_loss, flatten_loss = model(len(cameras)) images_pred = renderer.render_mesh(mesh) # optimize mesh with silhouette reprojection error and # geometry constraints loss = 0.0003 * flatten_loss if OPTIMIZE_SILHOUETTE: loss += neg_iou_loss(images_pred[:, -1], 1. - images_gt[:, 0]) else: loss += torch.nn.functional.l1_loss(images_pred[:, 0], images_gt[:, 0]) loop.set_description('Loss: %.8f' % (loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() if i % 100 == 0: image1 = images_pred[:, 0].detach().cpu().numpy()[ 0] #.transpose((1, 2, 0)) image2 = images_pred[:, 0].detach().cpu().numpy()[ 1] #.transpose((1, 2, 0)) writer.append_data((255 * np.concatenate( (image1, image2), axis=1)).astype(np.uint8)) imageio.imsave( os.path.join(args.output_dir, 'deform1_%05d.png' % i), (255 * image1).astype(np.uint8)) imageio.imsave( os.path.join(args.output_dir, 'deform2_%05d.png' % i), (255 * image2).astype(np.uint8)) # save optimized mesh model(1)[0].save_obj(os.path.join(args.output_dir, 'tf.obj'), save_texture=False)
def main(_args=None): if _args is None: parser = argparse.ArgumentParser() parser.add_argument('-c', '--classes', type=str, default='02958343,04256520') # , parser.add_argument('-l', '--load', type=bool, default=False) parser.add_argument('-t', '--template-mesh', type=str, default=os.path.join(data_dir, 'obj/sphere/sphere_1352.obj')) parser.add_argument('-o', '--output-dir', type=str, default=os.path.join(data_dir, 'results/output_reconstruct')) parser.add_argument('-b', '--batch-size', type=int, default=64) parser.add_argument('-n', '--train-num', type=int, default=5000) parser.add_argument('-i', '--image', type=str) args = parser.parse_args() else: args = _args os.makedirs(args.output_dir, exist_ok=True) model = nn.Sequential( resnet.resnet18(num_classes=512), Decoder(args.template_mesh) ).cuda() if args.load: model.load_state_dict(torch.load(os.path.join(args.output_dir, 'v3-%d.pth' % args.train_num))) renderer = sr.SoftRenderer(image_size=64, sigma_val=1e-4, aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=15) # read training images and camera poses dataset_train = ShapeNet(os.path.join(data_dir, 'dataset'), args.classes.split(','), 'train') dataset_val = ShapeNet(os.path.join(data_dir, 'dataset'), args.classes.split(','), 'val') optimizer = torch.optim.Adam(model.parameters(), 0.0001, betas=(0.9, 0.999)) if not args.load: writer = imageio.get_writer(os.path.join(args.output_dir, 'train.gif'), mode='I') train_num = args.train_num loop = tqdm.tqdm(list(range(0, train_num))) Loss_list = np.zeros(train_num) for i in loop: images, distances, elevations, viewpoints = dataset_train.get_random_batch(args.batch_size) images_gt = images.cuda() mesh, laplacian_loss, flatten_loss = model(images_gt) renderer.transform.set_eyes_from_angles(distances, elevations, viewpoints) images_pred = renderer.render_mesh(mesh) # optimize mesh with silhouette reprojection error and # geometry constraints loss = neg_iou_loss(images_pred[:, 3], images_gt[:, 3]) + \ 0.05 * laplacian_loss + \ 0.001 * flatten_loss Loss_list[i]=loss loop.set_description('Loss: %.4f'%(loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() if i % 100 == 0: image = images_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'deform_%05d.png'%i), (255*image[..., -1]).astype(np.uint8)) torch.save(model.state_dict(), os.path.join(args.output_dir, 'v3-%d.pth' % args.train_num)) fig = plt.figure() x = range(0, train_num) y = Loss_list plt.plot(x, y) plt.xlabel('Iteration') plt.ylabel('Test loss') plt.show() fig.savefig("loss.jpg") np.set_printoptions(threshold=np.inf) print(Loss_list) # print(Loss_list) val_iou_loss = 0 for i in tqdm.tqdm(range(0, dataset_val.images.shape[0], 24)): val_imgs = torch.from_numpy(dataset_val.images[i:i+24].astype('float32') / 255.) val_imgs = val_imgs.cuda() val_distances = torch.ones(val_imgs.size(0)) * dataset_val.distance val_elevations = torch.ones(val_imgs.size(0)) * dataset_val.elevation val_viewpoint_ids = torch.ones(24) for v in range(24): val_viewpoint_ids[v] = v val_mesh, val_laplacian_loss, val_flatten_loss = model(val_imgs) renderer.transform.set_eyes_from_angles(val_distances, val_elevations, -val_viewpoint_ids * 15) val_images_pred = renderer.render_mesh(val_mesh) val_iou_loss += neg_iou_loss(val_images_pred[:, 3], val_imgs[:, 3]).item() print('Val IOU loss: ', val_iou_loss / (dataset_val.images.shape[0] / 24)) if args.image: img = imageio.imread(args.image) model(torch.from_numpy(np.array([img.astype('float32') / 255.]).reshape((-1, 4, 64, 64))).cuda())[0].save_obj(os.path.join(args.output_dir, 'a.obj'), save_texture=False) # save optimized mesh model(dataset_val.get_val(args.classes.split(',')[0], 0, 4).cuda())[0].save_obj(os.path.join(args.output_dir, '1.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[0], 1, 16).cuda())[0].save_obj(os.path.join(args.output_dir, '2.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[0], 2, 8).cuda())[0].save_obj(os.path.join(args.output_dir, '3.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[0], 2, 10).cuda())[0].save_obj(os.path.join(args.output_dir, '3_0.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[0], 2, 12).cuda())[0].save_obj(os.path.join(args.output_dir, '3_1.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[0], 2, 1).cuda())[0].save_obj(os.path.join(args.output_dir, '3_2.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[0], 3, 20).cuda())[0].save_obj(os.path.join(args.output_dir, '4.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[0], 3, 1).cuda())[0].save_obj(os.path.join(args.output_dir, '4_0.obj'), save_texture=False) img3_0 = dataset_val.get_val(args.classes.split(',')[0], 2, 10).detach().cpu().numpy()[0].transpose((1, 2, 0)) img3_1 = dataset_val.get_val(args.classes.split(',')[0], 2, 12).detach().cpu().numpy()[0].transpose((1, 2, 0)) img3_2 = dataset_val.get_val(args.classes.split(',')[0], 2, 1).detach().cpu().numpy()[0].transpose((1, 2, 0)) # img1 = dataset_val.get_val(args.classes.split(',')[0], 0, 4).detach().cpu().numpy()[0].transpose((1, 2, 0)) # img2 = dataset_val.get_val(args.classes.split(',')[0], 1, 16).detach().cpu().numpy()[0].transpose((1, 2, 0)) img3 = dataset_val.get_val(args.classes.split(',')[0], 2, 8).detach().cpu().numpy()[0].transpose((1, 2, 0)) img4 = dataset_val.get_val(args.classes.split(',')[0], 3, 20).detach().cpu().numpy()[0].transpose((1, 2, 0)) # imageio.imsave(os.path.join(args.output_dir, '233.png'), (255 * img233[..., -1]).astype(np.uint8)) # imageio.imsave(os.path.join(args.output_dir, '235.png'), (255 * img235[..., -1]).astype(np.uint8)) # imageio.imsave(os.path.join(args.output_dir, 'car_1.png'), (255 * img1).astype(np.uint8)) # imageio.imsave(os.path.join(args.output_dir, 'car_2.png'), (255 * img2).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'car_3.png'), (255 * img3).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'car_4.png'), (255 * img4).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'car_3_0.png'), (255 * img3_0).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'car_3_1.png'), (255 * img3_1).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'car_3_2.png'), (255 * img3_2).astype(np.uint8)) "Plane dataset" model(dataset_val.get_val(args.classes.split(',')[1], 0, 4).cuda())[0].save_obj(os.path.join(args.output_dir, 'plane_1.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[1], 1, 16).cuda())[0].save_obj(os.path.join(args.output_dir, 'plane_2.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[1], 2, 8).cuda())[0].save_obj(os.path.join(args.output_dir, 'plane_3.obj'), save_texture=False) model(dataset_val.get_val(args.classes.split(',')[1], 3, 20).cuda())[0].save_obj(os.path.join(args.output_dir, 'plane_4.obj'), save_texture=False) img1 = dataset_val.get_val(args.classes.split(',')[1], 0, 4).detach().cpu().numpy()[0].transpose((1, 2, 0)) img2 = dataset_val.get_val(args.classes.split(',')[1], 1, 16).detach().cpu().numpy()[0].transpose((1, 2, 0)) img3 = dataset_val.get_val(args.classes.split(',')[1], 2, 8).detach().cpu().numpy()[0].transpose((1, 2, 0)) img4 = dataset_val.get_val(args.classes.split(',')[1], 3, 20).detach().cpu().numpy()[0].transpose((1, 2, 0)) # imageio.imsave(os.path.join(args.output_dir, '233.png'), (255 * img233[..., -1]).astype(np.uint8)) # imageio.imsave(os.path.join(args.output_dir, '235.png'), (255 * img235[..., -1]).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'plane_1.png'), (255 * img1).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'plane_2.png'), (255 * img2).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'plane_3.png'), (255 * img3).astype(np.uint8)) imageio.imsave(os.path.join(args.output_dir, 'plane_4.png'), (255 * img4).astype(np.uint8))
# face_model = BFM.BFM_torch() # bs=2 # camera_distance = 2.732 # elevation = 0 # azimuth = 0 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") renderer = sr.SoftRenderer(image_size=250, sigma_val=1e-4, aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=30, fill_back=False, perspective=True, light_intensity_ambient=1.0, light_intensity_directionals=0) renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) # face_loss = BFMFaceLoss(renderer, 20, device) # d3d_param =torch.Tensor([[-1.82443619e+00, 4.34013456e-02, -7.67325342e-01, # -1.35242629e+00, -5.11713028e-01, -4.10691090e-02, # 5.26137531e-01, -9.85630751e-01, -4.09944296e-01, # -4.43508774e-01, -1.52504385e+00, -4.70615327e-01, # -2.29804009e-01, 9.98994768e-01, 2.24300519e-01, # -2.83951283e-01, -9.16861773e-01, -2.70456612e-01, # -2.22548082e-01, 1.06795633e+00, -2.08728766e+00,
def main(): parser = argparse.ArgumentParser() parser.add_argument( '-i', '--data_dir', type=str, default='/home/mridul/Code/FyuseMesh/uxgpbaxotx/uxgpbaxotx/') parser.add_argument('-t', '--template_mesh', type=str, default=os.path.join(current_dir, '../data/obj/car/meshS124.obj')) parser.add_argument('-o', '--output_dir', type=str, default=os.path.join(current_dir, '../data/results/output_deform')) parser.add_argument('-is', '--image_size', type=int, default=256) parser.add_argument('-os', '--orig_image_size', type=int, default=1920) args = parser.parse_args() torch.set_printoptions(profile="full") os.makedirs(args.output_dir, exist_ok=True) model = Model(args.template_mesh).cuda() #Read in data ./ poseData = PoseReader(args.data_dir) numFrames = len(poseData) print('No of frames : ', numFrames) imagesGT = [] projMatrices = [] distCoeffs = [] flagSavedGT = False lPoseData = list(poseData.poses) for ii in range(numFrames): idx = lPoseData[ii] projMat = torch.mm(poseData.poses[idx]['K'], poseData.poses[idx]['Rt'][0:3, :]) try: imgGt = imageio.imread(args.data_dir + 'NewGt/{0:05d}'.format(idx) + '.png').astype('float32') / 255.0 except: try: imgGt = imageio.imread(args.data_dir + 'NewGt/{0:05d}'.format(idx) + '.jpg').astype('float32') / 255.0 except: raise imagesGT.append(imgGt) projMatrices.append(projMat) distCoeffs.append(poseData.poses[idx]['distortion']) projMatrices = torch.stack(projMatrices) distCoeffs = torch.stack(distCoeffs) if __debug__: print(projMatrices) projMatrices = projMatrices.cuda() distCoeffs = distCoeffs.cuda() #print (projMatrices) renderer = sr.SoftRenderer(image_size=args.image_size, sigma_val=1e-4, aggr_func_rgb='hard', camera_mode='projection', P=projMatrices, orig_size=args.orig_image_size) optimizer = torch.optim.Adam(model.parameters(), 0.001, betas=(0.5, 0.99)) loop = tqdm.tqdm(list(range(0, 2000))) writer = imageio.get_writer(os.path.join(args.output_dir, 'deform.gif'), mode='I') images_gt = torch.from_numpy(np.array(imagesGT)).cuda() for i in loop: mesh, laplacian_loss, flatten_loss = model(numFrames) images_pred = renderer.render_mesh(mesh) if __debug__: print(images_pred.shape) print(SilhouetteLoss(images_pred[:, 3], images_gt), laplacian_loss, flatten_loss) for ii in range(numFrames): imageio.imsave( os.path.join(args.output_dir, 'Debug/pred_%05d.png' % ii), (255 * images_pred[ii, 3, :, :].detach().cpu().numpy()).astype( np.uint8)) imageio.imsave( os.path.join(args.output_dir, 'Debug/gt_%05d.png' % ii), (255 * imagesGT[ii]).astype(np.uint8)) break # optimize mesh with silhouette reprojection error and geometry constraints loss = SilhouetteLoss( images_pred[:, 3], images_gt) + 0.03 * laplacian_loss + 0.0003 * flatten_loss loop.set_description('Loss: %.4f' % (loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() if not flagSavedGT: globalImgGt = np.zeros((args.image_size * int(numFrames / 10 + 1), args.image_size * 10), dtype=np.uint8) if i % 100 == 0: images = images_pred.detach().cpu().numpy() globalImg = 255 * np.ones( (args.image_size * int(numFrames / 10 + 1), args.image_size * 10), dtype=np.uint8) for ii in range(numFrames): col = int(ii % 10) row = int(ii / 10) image = images[ii].transpose((1, 2, 0)) globalImg[row * args.image_size:row * args.image_size + args.image_size, col * args.image_size:col * args.image_size + args.image_size] = (255 - 255 * image[..., -1]).astype( np.uint8) if not flagSavedGT: globalImgGt[row * args.image_size:row * args.image_size + args.image_size, col * args.image_size:col * args.image_size + args.image_size] = (128 * imagesGT[ii]).astype( np.uint8) writer.append_data(globalImg + globalImgGt) imageio.imsave( os.path.join(args.output_dir, 'deform_%05d.png' % i), globalImg + globalImgGt) # save optimized mesh model(1)[0].save_obj(os.path.join(args.output_dir, 'car.obj'), save_texture=False) if not flagSavedGT: imageio.imsave( os.path.join(args.output_dir, 'groundT_%05d.png' % i), globalImgGt) flagSavedGT = True
def main(): parser = argparse.ArgumentParser() parser.add_argument('-c', '--classes', type=str, default='02958343') # ,02691156 parser.add_argument('-l', '--load', type=bool, default=False) parser.add_argument('-t', '--template-mesh', type=str, default=os.path.join(data_dir, 'obj/sphere/sphere_1352.obj')) parser.add_argument('-o', '--output-dir', type=str, default=os.path.join(data_dir, 'results/output_reconstruct')) parser.add_argument('-b', '--batch-size', type=int, default=64) parser.add_argument('-n', '--train-num', type=int, default=5000) args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) model = nn.Sequential(Encoder(), Decoder(args.template_mesh)).cuda() if args.load: model.load_state_dict( torch.load( os.path.join(args.output_dir, 'v6-%d.pth' % args.train_num))) renderer = sr.SoftRenderer(image_size=256, sigma_val=1e-4, aggr_func_rgb='hard', camera_mode='look_at') # read training images and camera poses dataset_train = ShapeNet(os.path.join(data_dir, 'dataset_256_36_1'), args.classes) optimizer = torch.optim.Adam(model.parameters()) if not args.load: writer = imageio.get_writer(os.path.join(args.output_dir, 'train.gif'), mode='I') train_num = args.train_num loop = tqdm.tqdm(list(range(0, train_num))) Loss_list = np.zeros(train_num) for i in loop: images_a, viewpoints_a, distances, elevations = dataset_train.get_random_batch( args.batch_size) images_gt = images_a.cuda() mesh, laplacian_loss, flatten_loss = model(images_gt) renderer.transform.set_eyes_from_angles(distances, elevations, viewpoints_a) images_pred = renderer.render_mesh(mesh) # optimize mesh with silhouette reprojection error and # geometry constraints loss = neg_iou_loss(images_pred[:, 3], images_gt[:, 3]) + \ 0.07 * laplacian_loss + \ 0.004 * flatten_loss Loss_list[i] = loss.item() loop.set_description('Loss: %.4f' % (loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() if i % 100 == 0: image = images_pred.detach().cpu().numpy()[0].transpose( (1, 2, 0)) imagegt = images_gt.detach().cpu().numpy()[0].transpose( (1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) imageio.imsave( os.path.join(args.output_dir, 'deform_%05d.png' % i), (255 * image[..., -1]).astype(np.uint8)) imageio.imsave( os.path.join(args.output_dir, 'gt_%05d.png' % i), (255 * imagegt[..., -1]).astype(np.uint8)) torch.save(model.state_dict(), os.path.join(args.output_dir, 'v6-%d.pth' % args.train_num)) fig = plt.figure() x = range(0, train_num) y = Loss_list plt.plot(x, y) plt.xlabel('Iteration') plt.ylabel('Test loss') plt.show() fig.savefig("loss.jpg") np.set_printoptions(threshold=np.inf) print(Loss_list) # print(Loss_list) # save optimized mesh vals = torch.cat( (dataset_train.get_val(0, 4), dataset_train.get_val(1, 16), dataset_train.get_val(2, 8), dataset_train.get_val(3, 20))).cuda() val_res = model(vals)[0] for i in range(vals.shape[0]): val_res.save_obj(os.path.join(args.output_dir, 'car_%d.obj' % i), idx=i, save_texture=False) img = vals[i].detach().cpu().numpy().transpose((1, 2, 0)) imageio.imsave(os.path.join(args.output_dir, 'car_%d.png' % i), (255 * img).astype(np.uint8)) "Plane dataset"
def main(): parser = argparse.ArgumentParser() parser.add_argument('-i', '--filename-input', type=str, default=os.path.join(data_dir, 'source.npy')) parser.add_argument('-c', '--camera-input', type=str, default=os.path.join(data_dir, 'camera.npy')) parser.add_argument('-t', '--template-mesh', type=str, default=os.path.join(data_dir, 'obj/sphere/sphere_642_s.obj')) parser.add_argument('-o', '--output-dir', type=str, default=os.path.join(data_dir, 'results/output_deform')) parser.add_argument('-b', '--batch-size', type=int, default=120) args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) model = Model(args.template_mesh).cuda() renderer = sr.SoftRenderer(image_size=64, sigma_val=1e-4, aggr_func_rgb='hard', camera_mode='look_at', viewing_angle=15) # read training images and camera poses images = np.load(args.filename_input).astype('float32') / 255. cameras = np.load(args.camera_input).astype('float32') optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99)) camera_distances = torch.from_numpy(cameras[:, 0]) elevations = torch.from_numpy(cameras[:, 1]) viewpoints = torch.from_numpy(cameras[:, 2]) renderer.transform.set_eyes_from_angles(camera_distances, elevations, viewpoints) loop = tqdm.tqdm(list(range(0, 2000))) writer = imageio.get_writer(os.path.join(args.output_dir, 'deform.gif'), mode='I') images_gt = torch.from_numpy(images).cuda() flagSavedGT = False for i in loop: mesh, laplacian_loss, flatten_loss = model(args.batch_size) images_pred = renderer.render_mesh(mesh) # optimize mesh with silhouette reprojection error and # geometry constraints loss = neg_iou_loss(images_pred[:, 3], images_gt[:, 3]) + \ 0.03 * laplacian_loss + \ 0.0003 * flatten_loss loop.set_description('Loss: %.4f' % (loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() if __debug__: for ii in range(args.batch_size): imageio.imsave( os.path.join(args.output_dir, 'Debug/pred_%05d.png' % ii), (255 * images_pred[ii, 3, :, :].detach().cpu().numpy()).astype( np.uint8)) print(images[ii, 3, :, :]) imageio.imsave( os.path.join(args.output_dir, 'Debug/gt_%05d.png' % ii), (255 * images[ii].transpose(1, 2, 0)).astype(np.uint8)) break if i % 10 == 0: images = images_pred.detach().cpu().numpy() globalImg = np.zeros((832, 640), dtype=np.uint8) if not flagSavedGT: globalImgGt = np.zeros((832, 640), dtype=np.uint8) for ii in range(args.batch_size): col = int(ii % 10) row = int(ii / 10) image = images[ii].transpose((1, 2, 0)) globalImg[row * 64:row * 64 + 64, col * 64:col * 64 + 64] = (255 * image[..., -1]).astype(np.uint8) if not flagSavedGT: globalImgGt[row * 64:row * 64 + 64, col * 64:col * 64 + 64] = (255 * images_gt.detach().cpu().numpy()[ ii, 3, :, :]).astype(np.uint8) writer.append_data(globalImg) imageio.imsave( os.path.join(args.output_dir, 'deform_%05d.png' % i), globalImg) if not flagSavedGT: imageio.imsave( os.path.join(args.output_dir, 'groundT_%05d.png' % i), globalImgGt) flagSavedGT = True # save optimized mesh model(1)[0].save_obj(os.path.join(args.output_dir, 'plane.obj'), save_texture=False)
def main(): ########################### # camera settings ########################### camera_distance = 2.732 elevation = 0 azimuth = 0 ########################### # load object ########################### filename_input = os.path.join(data_dir, 'banana.obj') filename_ref = os.path.join(data_dir, 'example3_ref.png') mesh = TriangleMesh.from_obj(filename_input) vertices = mesh.vertices faces = mesh.faces.int() uvs = mesh.uvs face_textures = mesh.face_textures pmax = vertices.max() pmin = vertices.min() pmiddle = (pmax + pmin) / 2 vertices = vertices - pmiddle coef = 8 vertices = vertices * coef vertices = vertices[None, :, :].cuda() faces = faces[None, :, :].cuda() uvs = uvs[None, :, :].cuda() face_textures[None, :, :].cuda() image_gt = torch.from_numpy(imread(filename_ref).astype('float32') / 255.).permute(2, 0, 1)[None, ::].cuda() ########################## #NMR ########################## textures = torch.rand(1, faces.shape[1], 2, 2, 2, 3, dtype=torch.float32) model = Model(textures).cuda() renderer = nr.Renderer(camera_mode='look_at') renderer.eye = nr.get_points_from_angles(camera_distance, elevation, azimuth) renderer.perspective = False renderer.light_intensity_directional = 0.0 renderer.light_intensity_ambient = 1.0 optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing NMR') writer = imageio.get_writer(os.path.join(output_directory_nmr, 'texture.gif'), mode='I') for i in loop: optimizer.zero_grad() new_texture = model() image_pred, _, _ = renderer(vertices, faces, new_texture) loss = torch.sum((image_pred - image_gt)**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f' % (loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) ########################## #Soft Rasterizer ########################## textures = torch.rand(1, faces.shape[1], 2, 3, dtype=torch.float32) model = Model(textures).cuda() renderer = sr.SoftRenderer(image_size=256, sigma_val=3e-5, camera_mode='look_at', perspective=False, light_intensity_directionals=0.0, light_intensity_ambient=1.0) renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing SR') writer = imageio.get_writer(os.path.join(output_directory_sr, 'texture.gif'), mode='I') for i in loop: optimizer.zero_grad() new_texture = model() mesh = sr.Mesh(vertices, faces, new_texture) image_pred = renderer.render_mesh(mesh) loss = torch.sum(((image_pred[:, :3] - image_gt[None, :, :]))**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f' % (loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) ########################### # Dib-Renderer - Vertex Colours ########################### textures = torch.rand(1, vertices.shape[1], 3).cuda() model = Model(textures).cuda() renderer = Dib_Renderer(256, 256, mode='VertexColor') renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing VertexColor') writer = imageio.get_writer(os.path.join(output_directory_dib, 'texture_VertexColor.gif'), mode='I') for i in loop: optimizer.zero_grad() new_texture = model() image_pred, alpha, _ = renderer.forward( points=[vertices, faces[0].long()], colors=[new_texture]) image_pred = torch.cat((image_pred, alpha), dim=3) image_pred = image_pred.permute(0, 3, 1, 2) loss = torch.sum(((image_pred[:, :3] - image_gt[None, :, :]))**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f' % (loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) ########################### # Dib-Renderer - Lambertian ########################### textures = torch.rand(1, 3, 256, 256).cuda() model = Model(textures).cuda() renderer = Dib_Renderer(256, 256, mode='Lambertian') renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing Lambertian') writer = imageio.get_writer(os.path.join(output_directory_dib, 'texture_Lambertian.gif'), mode='I') for i in loop: optimizer.zero_grad() new_texture = model() image_pred, alpha, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), new_texture]) image_pred = torch.cat((image_pred, alpha), dim=3) image_pred = image_pred.permute(0, 3, 1, 2) loss = torch.sum(((image_pred[:, :3] - image_gt[None, :, :]))**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f' % (loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) ########################### # Dib-Renderer - Phong # ########################### textures = torch.rand(1, 3, 256, 256).cuda() model = Model(textures).cuda() renderer = Dib_Renderer(256, 256, mode='Phong') renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99)) loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing Phong') ### Lighting info ### material = np.array([[0.3, 0.3, 0.3], [1.0, 1.0, 1.0], [0.4, 0.4, 0.4]], dtype=np.float32).reshape(-1, 3, 3) material = torch.from_numpy(material).repeat(1, 1, 1).cuda() shininess = np.array([100], dtype=np.float32).reshape(-1, 1) shininess = torch.from_numpy(shininess).repeat(1, 1).cuda() lightdirect = 2 * np.random.rand(1, 3).astype(np.float32) - 1 lightdirect[:, 2] += 2 lightdirect = torch.from_numpy(lightdirect).cuda() writer = imageio.get_writer(os.path.join(output_directory_dib, 'texture_Phong.gif'), mode='I') for i in loop: optimizer.zero_grad() new_texture = model() image_pred, alpha, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), new_texture],\ light= lightdirect, \ material=material, \ shininess=shininess) image_pred = torch.cat((image_pred, alpha), dim=3) image_pred = image_pred.permute(0, 3, 1, 2) loss = torch.sum(((image_pred[:, :3] - image_gt[None, :, :]))**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f' % (loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8)) ########################### # Dib-Renderer - SphericalHarmonics ########################### textures = torch.rand(1, 3, 256, 256).cuda() model = Model(textures).cuda() renderer = Dib_Renderer(256, 256, mode='SphericalHarmonics') renderer.set_look_at_parameters([90 - azimuth], [elevation], [camera_distance]) optimizer = torch.optim.Adam(model.parameters(), 0.01, betas=(0.5, 0.99)) lightparam = np.random.rand(1, 9).astype(np.float32) lightparam[:, 0] += 4 lightparam = torch.from_numpy(lightparam).cuda() loop = tqdm.tqdm(range(2000)) loop.set_description('Optimizing SH') writer = imageio.get_writer(os.path.join(output_directory_dib, 'texture_SH.gif'), mode='I') for i in loop: optimizer.zero_grad() new_texture = model() image_pred, alpha, _ = renderer.forward(points=[vertices, faces[0].long()], \ colors=[uvs, face_textures.long(), new_texture],\ light=lightparam) image_pred = torch.cat((image_pred, alpha), dim=3) image_pred = image_pred.permute(0, 3, 1, 2) loss = torch.sum(((image_pred[:, :3] - image_gt[None, :, :]))**2) loss.backward() optimizer.step() loop.set_description('Loss: %.4f' % (loss.item())) if i % 20 == 0: image = image_pred.detach().cpu().numpy()[0].transpose((1, 2, 0)) writer.append_data((255 * image).astype(np.uint8))
def save_as_gif(verts, faces, output_path, input_img=None, camera_distance=3.0, elevation=30.0, rotation_delta=4, verbose=False): assert len(verts.shape) == 3 assert len(faces.shape) == 3 assert input_img is None or len(input_img.shape) == 4 if torch.cuda.is_available() is not True: return None # soft renderer is only supported under cuda else: if verbose: print("saving as gif...") # downsample, rescale and transpose input_img if input_img is not None: input_img = nn.Upsample((256, 256), mode='bilinear')(input_img) input_img = 255 * input_img if torch.max( input_img) <= 1.0 else input_img input_img = input_img[0, :, :, :].cpu().numpy().transpose((1, 2, 0)) else: input_img = None # prepare for output output_path = os.path.splitext( output_path)[0] + '.gif' # replace extention by .gif os.makedirs(os.path.dirname(output_path), exist_ok=True) # make output dir writer = imageio.get_writer(output_path, mode='I') video = [] # for output if verbose: print("output_path: {}".format(output_path)) # make mesh and renderer mesh = sr.Mesh(verts[0, :, :], faces[0, :, :]) renderer = sr.SoftRenderer(camera_mode='look_at') loop = list(range(0, 360, rotation_delta)) loop = tqdm.tqdm(loop) if verbose else loop for idx, azimuth in enumerate(loop): mesh.reset_() renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth) # render imgs = renderer.render_mesh(mesh) img = imgs.detach().cpu()[0, :, :, :].numpy().transpose( (1, 2, 0)) * 255 img = np.concatenate( (input_img, img[:, :, :3]), axis=1) if input_img is not None else img writer.append_data((img).astype(np.uint8)) video.append( torch.tensor((img).astype(np.uint8).transpose(2, 0, 1)).unsqueeze(0)) writer.close() return torch.cat(video, dim=0).unsqueeze(0) # (#batch, #time, #rgb, #h, #w)