diff --git a/README.md b/README.md index 3a3d8bb..2a7a971 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@

## Update +* add pointclouds rendering code used for paper figure, see `utils/render_mitsuba_pc.py` * When opening an issue, please add @ZENGXH so that I can reponse faster! ## Install diff --git a/utils/_render_mitsuba_cubes.py b/utils/_render_mitsuba_cubes.py new file mode 100644 index 0000000..e377c5e --- /dev/null +++ b/utils/_render_mitsuba_cubes.py @@ -0,0 +1,393 @@ +# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES 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 & AFFILIATES is strictly prohibited. + +import numpy as np +import sys, os, subprocess +import mitsuba as mi +import OpenEXR +import Imath +from PIL import Image +from plyfile import PlyData, PlyElement +import torch +import time +from PIL import Image, ImageChops +random_str = hex(int(time.time() + 12345))[2:] +def trim(im): + bg = Image.new(im.mode, im.size, im.getpixel((0,0))) ##border) + diff = ImageChops.difference(im, bg) + bbox = diff.getbbox() + if bbox: + return im.crop(bbox) + else: + return im + + +PATH_TO_MITSUBA2 = "/home/xzeng/code/mitsuba2/build/dist/mitsuba" ## Codes/mitsuba2/build/dist/mitsuba" # mitsuba exectuable + + +xml_head_segment = \ + """ + + + + + + + + + + + + + + + + + + + + + + + + + + + + +""" + +xml_head = \ + """ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +""" + +# I also use a smaller point size +xml_ball_segment = \ + """ + + + + + + + + + +""" +xml_cube_segment = \ + """ + + + + + + + + + + + + + + +""" + +xml_shape_segment = \ + """ + + + + + + + + + + +""" + ## +# +# +# +# +# +# +# +# +# +# +# +# + + + + +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# +# + +# +# +# +# +# +# +# +# +# +# +# +# +# +# + +xml_tail = \ + """ + + + + + + + + + + + + + + + + + + + + +""" +# +# +# +# +# +# +# + + ## +# +# +# +# +# +# +# +# +# +# + +def colormap(x, y, z): + vec = np.array([x, y, z]) + vec = np.clip(vec, 0.001, 1.0) + norm = np.sqrt(np.sum(vec ** 2)) + vec /= norm + return [vec[0], vec[1], vec[2]] + + +def standardize_bbox(pcl, points_per_object): + pt_indices = np.random.choice(pcl.shape[0], points_per_object, replace=False) + np.random.shuffle(pt_indices) + pcl = pcl[pt_indices] # n by 3 + mins = np.amin(pcl, axis=0) + maxs = np.amax(pcl, axis=0) + center = (mins + maxs) / 2. + scale = np.amax(maxs - mins) + print("Center: {}, Scale: {}".format(center, scale)) + result = ((pcl - center) / scale).astype(np.float32) # [-0.5, 0.5] + return result + + +# only for debugging reasons +def writeply(vertices, ply_file): + sv = np.shape(vertices) + points = [] + for v in range(sv[0]): + vertex = vertices[v] + points.append("%f %f %f\n" % (vertex[0], vertex[1], vertex[2])) + print(np.shape(points)) + file = open(ply_file, "w") + file.write('''ply + format ascii 1.0 + element vertex %d + property float x + property float y + property float z + end_header + %s + ''' % (len(vertices), "".join(points))) + file.close() + + +# as done in https://gist.github.com/drakeguan/6303065 +def ConvertEXRToJPG(exrfile, jpgfile): + File = OpenEXR.InputFile(exrfile) + PixType = Imath.PixelType(Imath.PixelType.FLOAT) + DW = File.header()['dataWindow'] + Size = (DW.max.x - DW.min.x + 1, DW.max.y - DW.min.y + 1) + + rgb = [np.fromstring(File.channel(c, PixType), dtype=np.float32) for c in 'RGB'] + for i in range(3): + print(rgb[i].max(), rgb[i].min()) + rgb[i] = np.where(rgb[i] <= 0.0031308, + (rgb[i] * 12.92) * 255.0, + (1.055 * (rgb[i] ** (1.0 / 2.4)) - 0.055) * 255.0) + + rgb8 = [Image.frombytes("F", Size, c.tostring()).convert("L") for c in rgb] + # rgb8 = [Image.fromarray(c.astype(int)) for c in rgb] + Image.merge("RGB", rgb8).save(jpgfile, "PNG") + ##, quality=95) + img = Image.open(jpgfile) + img = trim(img) + img.save(jpgfile) + + +def render_cubes2png(pcl, filename, colorm=[24,107,239], vs_size=0.5, + sample_count=256, out_width=800, out_height=600, + lookat_1=3, lookat_2=3, lookat_3=3, trim_img=0): + xml_head = xml_head_segment.format( + lookat_1, lookat_2, lookat_3, + sample_count, out_width, out_height) + xml_segments = [xml_head] + print('color: ', colorm) + for i in range(pcl.shape[0]): + color = [c/255.0 for c in colorm] + xml_segments.append(xml_cube_segment.format( + vs_size, vs_size, vs_size, + pcl[i, 0], pcl[i, 1], pcl[i, 2], + *color)) + xml_segments.append(xml_tail) + + xml_content = str.join('', xml_segments) + + xmlFile = '/tmp/tmp_%s.xml'%random_str # ("%s/xml/%s.xml" % (folder, filename)) + exrFile = '/tmp/tmp_%s.exr'%random_str ##("%s/exr/%s.exr" % (folder, filename)) + ## png = ("%s/%s.png" % (folder, filename)) + with open(xmlFile, 'w') as f: + f.write(xml_content) + f.close() + print('save as xmlFile: ', xmlFile) + + if not os.path.exists(os.path.dirname(exrFile)): + os.makedirs(os.path.dirname(exrFile)) + #subprocess.run([PATH_TO_MITSUBA2, '-o', exrFile, xmlFile]) + png = filename + if not os.path.exists(os.path.dirname(png)): + os.makedirs(os.path.dirname(png)) + + #print(['Converting EXR to PNG ...', png]) + #ConvertEXRToJPG(exrFile, png) + scene = mi.load_file(xmlFile) + image = mi.render(scene) ##, spp=256) + mi.util.write_bitmap(png, image) + if trim_img: + img = Image.open(png) + img = trim(img) + img.save(png) + + return png + + +if __name__ == "__main__": + if (len(sys.argv) < 2): + print('filename to npy/ply is not passed as argument. terminated.') + raise ValueError + + pathToFile = sys.argv[1] + + + main(pathToFile) diff --git a/utils/render_mitsuba_pc.py b/utils/render_mitsuba_pc.py new file mode 100644 index 0000000..30588bd --- /dev/null +++ b/utils/render_mitsuba_pc.py @@ -0,0 +1,418 @@ +# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES 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 & AFFILIATES is strictly prohibited. + +import numpy as np +import mitsuba as mi +mi.set_variant("cuda_ad_rgb") +from loguru import logger +import sys, os, subprocess +import copy +import OpenEXR +import Imath +from PIL import Image +## from plyfile import PlyData, PlyElement +import torch +import open3d as o3d +from PIL import Image, ImageChops +import time +random_str = hex(int(time.time() + 12345))[2:] +PATH_TO_MITSUBA2 = "/home/xzeng/code/mitsuba2/build/dist/mitsuba" ## Codes/mitsuba2/build/dist/mitsuba" # mitsuba exectuable + +# replaced by command line arguments +def standardize_bbox_based_on(pcl, eps): + pcl = pcl.numpy()[:, [0,2,1]] + eps = eps.numpy()[:, [0,2,1]] + pcl, center, scale = standardize_bbox(pcl, return_center_scale=1) + eps = (eps - center) / scale if eps is not None else None + offset = - 0.475 - pcl[:,2].min() + eps[:,2] += offset + return torch.from_numpy(eps) + +# PATH_TO_NPY = 'pcl_ex.npy' # the tensor to load +def rotate_pts(pts, r, axis=1, do_transform=0, is_point_flow_data=1, eps=None): + assert(len(pts.shape) == 2), f'require N,3 tensor, get: {pts.shape}' + ## logger.info('rotating pts: {}, get eps: {} ', pts.shape, eps is not None ) + is_tensor = torch.is_tensor(pts) + if not is_tensor: + pts = torch.from_numpy(pts) + if eps is not None and not torch.is_tensor(eps): + eps = torch.from_numpy(eps) + if do_transform: + pcl = pts.cpu().numpy() + eps = eps.cpu().numpy() if eps is not None else None + if not is_point_flow_data: + pcl[:,0] *= -1 + pcl = pcl[:, [2,1,0]] + + if eps is not None: + eps[:,0] *= -1 + eps = eps[:, [2,1,0]] + + pcl, center, scale = standardize_bbox(pcl, return_center_scale=1) + eps = (eps - center) / scale if eps is not None else None + pcl = pcl[:, [2, 0, 1]] + pcl[:,0] *= -1 + pcl[:,2] += 0.0125 + if eps is not None: + eps = eps[:, [2, 0, 1]] + eps[:,0] *= -1 + eps[:,2] += 0.0125 + + offset = - 0.475 - pcl[:,2].min() + pcl[:,2] += offset + if eps is not None: + eps[:,2] += offset + pts = torch.from_numpy(pcl) + eps = torch.from_numpy(eps) if eps is not None else None + + pcl = o3d.geometry.PointCloud() + pcl.points = o3d.utility.Vector3dVector(pts.cpu()) + if axis == 1: + R = pcl.get_rotation_matrix_from_xyz((0, - r * np.pi / 2, 0)) + elif axis == 2: + R = pcl.get_rotation_matrix_from_xyz((0, 0, - r * np.pi / 2)) + elif axis == 0: + R = pcl.get_rotation_matrix_from_xyz((- r * np.pi / 2, 0, 0)) + + #mesh_r = copy.deepcopy(pcl) + #mesh_r.rotate(R, center=(0, 0, 0)) + #pts = np.asarray(mesh_r.points) + + h_center = w_center = 0 + center = np.array([h_center, w_center, 0]).reshape(-1,3) + pts = np.matmul(pts.numpy() - center, R.T) + center + eps = np.matmul(eps.numpy() - center, R.T) + center if eps is not None else eps + + if is_tensor: + pts = torch.from_numpy(pts) + eps = torch.from_numpy(eps) if eps is not None and not torch.is_tensor(eps) else eps + + if eps is not None: + return pts, eps + return pts + +# note that sampler is changed to 'independent' and the ldrfilm is changed to hdrfilm +xml_head_segment = \ + """ + + + + + + + + + + + + + + + + + + + + + + + + + + + + +""" + +# I also use a smaller point size +xml_ball_segment = ['']*10 +xml_ball_segment[0] = \ + """ + + + + + + + + + +""" +xml_ball_segment[1] = \ +""" + + + + + + + + + + +""" + + +xml_ball_segment[2] = \ +""" + + + + + + + + + + +""" + +xml_ball_segment[3] = \ +""" + + + + + + + + + + + + +""" +xml_ball_segment[4] = \ +""" + + + + + + + + + + + + +""" +xml_ball_segment[5] = \ + """ + + + + + + + + + + + + +""" +xml_tail = \ +""" + + + + + + + + + + + + + + + + + + +""" +def trim(im): + bg = Image.new(im.mode, im.size, im.getpixel((0,0))) ##border) + diff = ImageChops.difference(im, bg) + bbox = diff.getbbox() + if bbox: + return im.crop(bbox) + else: + return im + + +def colormap(x, y, z): + if torch.is_tensor(x): + x = x.cpu().numpy() + vec = np.array([x, y, z]) + vec = np.clip(vec, 0.001, 1.0) + norm = np.sqrt(np.sum(vec ** 2)) + vec /= norm + return [vec[0], vec[1], vec[2]] + + +def standardize_bbox(pcl, return_center_scale=0): + #pt_indices = np.random.choice(pcl.shape[0], points_per_object, replace=False) + #np.random.shuffle(pt_indices) + #pcl = pcl[pt_indices] # n by 3 + if torch.is_tensor(pcl): + pcl = pcl.numpy() + mins = np.amin(pcl, axis=0) + maxs = np.amax(pcl, axis=0) + center = (mins + maxs) / 2. + scale = np.amax(maxs - mins) + #print("Center: {}, Scale: {}".format(center, scale)) + result = ((pcl - center) / scale).astype(np.float32) # [-0.5, 0.5] + if return_center_scale: + return result, center, scale + return result + + +# only for debugging reasons +def writeply(vertices, ply_file): + sv = np.shape(vertices) + points = [] + for v in range(sv[0]): + vertex = vertices[v] + points.append("%f %f %f\n" % (vertex[0], vertex[1], vertex[2])) + print(np.shape(points)) + file = open(ply_file, "w") + file.write('''ply + format ascii 1.0 + element vertex %d + property float x + property float y + property float z + end_header + %s + ''' % (len(vertices), "".join(points))) + file.close() + + +# as done in https://gist.github.com/drakeguan/6303065 +def ConvertEXRToJPG(exrfile, jpgfile, trim_img): + File = OpenEXR.InputFile(exrfile) + PixType = Imath.PixelType(Imath.PixelType.FLOAT) + DW = File.header()['dataWindow'] + Size = (DW.max.x - DW.min.x + 1, DW.max.y - DW.min.y + 1) + + rgb = [np.fromstring(File.channel(c, PixType), dtype=np.float32) for c in 'RGB'] + for i in range(3): + rgb[i] = np.where(rgb[i] <= 0.0031308, + (rgb[i] * 12.92) * 255.0, + (1.055 * (rgb[i] ** (1.0 / 2.4)) - 0.055) * 255.0) + + rgb8 = [Image.frombytes("F", Size, c.tostring()).convert("L") for c in rgb] + Image.merge("RGB", rgb8).save(jpgfile, "PNG") ##JPEG", quality=95) + img = Image.open(jpgfile) + if trim_img: + img = trim(img) + img.save(jpgfile) + +def pts2png(input_pts, file_name, colorm=[24,107,239], + skip_if_exists=False, is_color_list=False, + sample_count=256, out_width=1600, out_height=1200, + ball_size=0.025, do_standardize=0, same_computed_loc_color=0, material_id=0, precomputed_color=None, + output_xml_file=None, + use_loc_color=False, lookat_1=3, lookat_2=3, lookat_3=3, do_transform=1, trim_img=0): + """ + Argus: + input_pts: (B,N,3) the points to be render + file_name: list; output image name + """ + assert(len(input_pts.shape) == 3), f'expect: B,N,3; get: {input_pts.shape}' + assert(type(file_name) is list), f'require file_name as list' + xml_head = xml_head_segment.format( + lookat_1, lookat_2, lookat_3, + sample_count, out_width, out_height) + input_pts = input_pts.cpu() + # print('get shape; ', input_pts.shape) + color_list = [] + for pcli in range(0, input_pts.shape[0]): + xmlFile = '/tmp/tmp_%s.xml'%random_str if output_xml_file is None else output_xml_file + # ("%s/xml/%s.xml" % (folder, filename)) + exrFile = '/tmp/tmp_%s.exr'%random_str ##("%s/exr/%s.exr" % (folder, filename)) + png = file_name[pcli] + if skip_if_exists and os.path.exists(png): + print(f'find png: {png}, skip ') + continue + pcl = input_pts[pcli, :, :] + if do_transform: + pcl = standardize_bbox(pcl) + pcl = pcl[:, [2, 0, 1]] + pcl[:, 0] *= -1 + pcl[:, 2] += 0.0125 + + offset = - 0.475 - pcl[:,2].min() + pcl[:,2] += offset + if do_standardize: + pcl = standardize_bbox(pcl) + offset = - 0.475 - pcl[:,2].min() + pcl[:,2] += offset + + xml_segments = [xml_head] + for i in range(pcl.shape[0]): + if precomputed_color is not None: + color = precomputed_color[i] + elif use_loc_color and not same_computed_loc_color: + color = colormap(pcl[i, 0] + 0.5, pcl[i, 1] + 0.5, pcl[i, 2] + 0.5 - 0.0125) + elif use_loc_color and same_computed_loc_color: + if pcli == 0: + color = colormap(pcl[i, 0] + 0.5, pcl[i, 1] + 0.5, pcl[i, 2] + 0.5 - 0.0125) + color_list.append(color) + else: + color = color_list[i] # same color as first shape + elif is_color_list: + color = colorm[pcli] + color = [c/255.0 for c in color] + else: + color = [c/255.0 for c in colorm] + xml_segments.append(xml_ball_segment[material_id].format( + ball_size, + pcl[i, 0], pcl[i, 1], pcl[i, 2], *color)) + ## print('using color: ', color) + xml_segments.append(xml_tail) + + xml_content = str.join('', xml_segments) + + if not os.path.exists(os.path.dirname(xmlFile)): + os.makedirs(os.path.dirname(xmlFile)) + with open(xmlFile, 'w') as f: + f.write(xml_content) + logger.info('[render_mitsuba_pc] write output at: {}', xmlFile) + f.close() + + if not os.path.exists(os.path.dirname(exrFile)): + os.makedirs(os.path.dirname(exrFile)) + if not os.path.exists(os.path.dirname(png)): + os.makedirs(os.path.dirname(png)) + logger.info('*'*20 + f'{png}' +'*'*20) + # mitsuba2 + #subprocess.run([PATH_TO_MITSUBA2, '-o', exrFile, xmlFile]) + #ConvertEXRToJPG(exrFile, png, trim_img) + scene = mi.load_file(xmlFile) + image = mi.render(scene) ##, spp=256) + mi.util.write_bitmap(png, image) + if trim_img: + img = Image.open(png) + img.save(png) + + return png + + +if __name__ == "__main__": + if (len(sys.argv) < 2): + print('filename to npy/ply is not passed as argument. terminated.') + raise ValueError + + pathToFile = sys.argv[1] + + + main(pathToFile) diff --git a/utils/render_voxel_cubes.py b/utils/render_voxel_cubes.py new file mode 100644 index 0000000..9eb496f --- /dev/null +++ b/utils/render_voxel_cubes.py @@ -0,0 +1,128 @@ +# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES 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 & AFFILIATES is strictly prohibited. + +import os +import numpy as np +import open3d as o3d +import sys +import torch +import point_cloud_utils as pcu +from PIL import Image + +sys.path.append('.') +import torch +from utils._render_mitsuba_cubes import render_cubes2png +# from script.paper.vis_mesh.render_mitsuba import reformat_ply +def reformat_ply(input, output, r=0, is_point_flow_data=0, + ascii=False, write_normal=False, fixed_trimesh=1): + m = open3d.io.read_triangle_mesh(input) + pcl = np.asarray(m.vertices) + if not is_point_flow_data: + pcl[:,0] *= -1 + pcl = pcl[:, [2,1,0]] + + pcl = standardize_bbox(pcl) + pcl = pcl[:, [2, 0, 1]] + pcl[:,0] *= -1 + pcl[:,2] += 0.0125 + + offset = - 0.475 - pcl[:,2].min() + pcl[:,2] += offset + m.vertices = open3d.utility.Vector3dVector(pcl) + + R = m.get_rotation_matrix_from_xyz((0, 0, - r * np.pi / 2)) + mesh_r = copy.deepcopy(m) + mesh_r.rotate(R, center=(0, 0, 0)) + ## o3d.visualization.draw_geometries([mesh_r]) + open3d.io.write_triangle_mesh(output, mesh_r, write_ascii=ascii, write_vertex_normals=False) + if fixed_trimesh: + mesh = trimesh.load_mesh(output) ## '../models/featuretype.STL') + trimesh.repair.fix_inversion(mesh) + trimesh.repair.fix_normals(mesh) + mesh.export(output) + logger.info(f'load {input}, write as {output}; ascii={ascii}, write_normal: {write_normal}') + return output + + +def get_vpts(cubes): + import kaolin + v,f = kaolin.ops.conversions.voxelgrids_to_cubic_meshes(cubes) ##voxelgrids_to_trianglemeshes(voxel_volume) + v = [vi.cpu() for vi in v] + f = [fi.cpu() for fi in f] + return v, f + +def create_dir(output_path): + if not os.path.exists(output_path): + os.makedirs(output_path) + +def convert_cube_2_mesh(voxel_size, center_list, output_path, overwrite=0, colorm=[93,64,211], rotate=None, config={}): + scale = 0.5 * (center_list.max() - center_list.min()) + pcl = center_list + mins = np.amin(pcl, axis=0).reshape(1,pcl.shape[-1]) ##np.amin(pcl, axis=1)[:, None, :], axis=0)[None, None, :] + maxs = np.amax(pcl, axis=0).reshape(1,pcl.shape[-1]) ##np.amin(pcl, axis=1)[:, None, :], axis=0)[None, None, :] + + center = ( mins + maxs ) / 2. + scale = np.amax(maxs-mins) + + center_list = (center_list - center) / scale ## (center_list.max() - center_list.min()) + + center_list = center_list[:,[2,0,1]] + center_list[:,0] *= -1 #center_list + offset = - 0.475 - center_list[:,2].min() + center_list[:,2] += offset + if rotate is not None: + pts = torch.from_numpy(center_list) + pcl = o3d.geometry.PointCloud() + pcl.points = o3d.utility.Vector3dVector(pts.cpu()) + R = pcl.get_rotation_matrix_from_xyz((0, 0, - rotate * np.pi / 2)) + center = np.array([0, 0, 0]).reshape(-1,3) + pts = np.matmul(pts.numpy() - center, R.T) + center + center_list = pts + + output_name = output_path + print('center_list: ', center_list.shape) + if not os.path.exists(os.path.dirname(output_name)): + os.makedirs(os.path.dirname(output_name)) + if os.path.exists(output_name) and not overwrite: + print(f'find rendered output: {output_name}, skip') + return + out = render_cubes2png(pcl=center_list, filename=output_name, + vs_size=0.9*voxel_size/scale, colorm=colorm, **config) + print(' save as: ', out) + #img = Image.open(out) + #img.show() + print('output_path: ', output_path) + return out + +def create_unit_cube(): + voxelgrid = torch.tensor([1], device='cuda', dtype=torch.uint8).view(1,1,1,1) + v, f = get_vpts(voxelgrid) + v = v[0] + f = f[0] + v = v - 0.5 # [-0.5, 0.5] + print('v: ', v) + print('f: ', v) + output_name = './script/paper/vis_voxel_exp/unit_cube.ply' + print('save output as: ', output_name) + pcu.save_mesh_vf(output_name, v.numpy(), f.numpy()) + reformat_ply(output_name, output_name, r=0) + +def write_small_cube(): + mesh_compare = './script/paper/vis_voxel_exp/unit_cube2.ply' + m = o3d.io.read_triangle_mesh(mesh_compare) + pcl = np.asarray(m.vertices) * 0.5 + m.vertices = o3d.utility.Vector3dVector(pcl) ##.float()) ## torch.from_numpy(pcl)) + o3d.io.write_triangle_mesh(mesh_compare.replace('.ply', 'w1.ply'), + m, write_ascii=True, write_vertex_normals=True) + +# write_small_cube() +if __name__ == '__main__': + input_path = '/home/xzeng/plots/voxel_exp/voxel_cube/raw' + index = "3 328 91 83 74 73 64 63 54 51 48 45 41 30 22" + convert_cube_2_mesh(input_path, index) +# create_unit_cube()