Parallel processing for Shapenet Dataset

docker/Dockerfile

@@ -1,10 +1,10 @@
-# Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 
+# Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 # property and proprietary rights in and to this material, related
-# documentation and any modifications thereto. Any use, reproduction, 
-# disclosure or distribution of this material and related documentation 
-# without an express license agreement from NVIDIA CORPORATION or 
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
 # its affiliates is strictly prohibited.
 
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:21.08-py3
@@ -49,3 +49,6 @@ RUN pip install meshzoo ipdb imageio gputil h5py point-cloud-utils imageio image
 
 # HDR image support
 RUN imageio_download_bin freeimage
+
+# Blender Dependancies
+RUN apt-get install -y libxi6 libgconf-2-4 libfontconfig1 libxrender1

render_shapenet_data/render_all.py

@@ -8,6 +8,7 @@
 
 import os
 import argparse
+import time
 
 parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
 parser.add_argument(
@@ -35,10 +36,15 @@
     0.7,
     0.9
 ]
+
 for synset, obj_scale in zip(synset_list, scale_list):
     file_list = sorted(os.listdir(os.path.join(dataset_folder, synset)))
     for idx, file in enumerate(file_list):
+        start_time = time.time()
         render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 >> tmp.out' % (
             blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
         )
         os.system(render_cmd)
+
+        end_time = time.time()
+        print('Time for rendering %d models: %f' % (1, end_time - start_time))

render_shapenet_data/render_parallel.py

@@ -0,0 +1,116 @@
+# 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 argparse
+import subprocess
+import time
+
+parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+parser.add_argument(
+    '--save_folder', type=str, default='./tmp',
+    help='path for saving rendered image')
+parser.add_argument(
+    '--dataset_folder', type=str, default='./tmp',
+    help='path for downloaded 3d dataset folder')
+parser.add_argument(
+    '--blender_root', type=str, default='./tmp',
+    help='path for blender')
+args = parser.parse_args()
+
+save_folder = args.save_folder
+dataset_folder = args.dataset_folder
+blender_root = args.blender_root
+
+synset_list = [
+    # '02958343',  # Car
+    '03001627',  # Chair
+    '03790512'  # Motorbike
+]
+scale_list = [
+    # 0.9,
+    0.7,
+    0.9
+]
+
+for synset, obj_scale in zip(synset_list, scale_list):
+    file_list = sorted(os.listdir(os.path.join(dataset_folder, synset)))
+    idx = 0
+    start_time = time.time()
+    while idx < len(file_list):
+        print("Done with %d/%d" % (idx, len(file_list)))
+
+        stdout = open('stdout.txt', 'w')
+        stderr = open('stderr.txt', 'w')
+        file = file_list[idx]
+        render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 --gpu 0' % (
+            blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
+        )
+        p0 = subprocess.Popen(render_cmd, shell=True, stdout=stdout, stderr=stderr)
+        idx += 1
+
+        file = file_list[idx]
+        render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 --gpu 1' % (
+            blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
+        )
+        p1 = subprocess.Popen(render_cmd, shell=True, stdout=stdout, stderr=stderr)
+        idx += 1
+
+        file = file_list[idx]
+        render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 --gpu 2' % (
+            blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
+        )
+        p2 = subprocess.Popen(render_cmd, shell=True, stdout=stdout, stderr=stderr)
+        idx += 1
+
+        file = file_list[idx]
+        render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 --gpu 3' % (
+            blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
+        )
+        p3 = subprocess.Popen(render_cmd, shell=True, stdout=stdout, stderr=stderr)
+        idx += 1
+
+        file = file_list[idx]
+        render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 --gpu 4' % (
+            blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
+        )
+        p4 = subprocess.Popen(render_cmd, shell=True, stdout=stdout, stderr=stderr)
+        idx += 1
+
+        file = file_list[idx]
+        render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 --gpu 5' % (
+            blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
+        )
+        p5 = subprocess.Popen(render_cmd, shell=True, stdout=stdout, stderr=stderr)
+        idx += 1
+
+        file = file_list[idx]
+        render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 --gpu 6' % (
+            blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
+        )
+        p6 = subprocess.Popen(render_cmd, shell=True, stdout=stdout, stderr=stderr)
+        idx += 1
+
+        file = file_list[idx]
+        render_cmd = '%s -b -P render_shapenet.py -- --output %s %s  --scale %f --views 24 --resolution 1024 --gpu 7' % (
+            blender_root, save_folder, os.path.join(dataset_folder, synset, file, 'model.obj'), obj_scale
+        )
+        p7 = subprocess.Popen(render_cmd, shell=True, stdout=stdout, stderr=stderr)
+        idx += 1
+
+        p0.wait()
+        p1.wait()
+        p2.wait()
+        p3.wait()
+        p4.wait()
+        p5.wait()
+        p6.wait()
+        p7.wait()
+
+    end_time = time.time()
+    print('Time for rendering %d models: %f' % (len(file_list), end_time - start_time))

render_shapenet_data/render_shapenet.py

@@ -31,6 +31,9 @@
 parser.add_argument(
     '--resolution', type=int, default=512,
     help='Resolution of the images.')
+parser.add_argument(
+    '--gpu', type=int, default=0,
+    help='gpu.')
 parser.add_argument(
     '--engine', type=str, default='CYCLES',
     help='Blender internal engine for rendering. E.g. CYCLES, BLENDER_EEVEE, ...')
@@ -83,13 +86,12 @@ def enable_cuda_devices():
     # If we have CUDA/OPENCL devices, enable only them, otherwise enable
     # all devices (assumed to be CPU)
     print(cprefs.devices)
-    for device in cprefs.devices:
-        device.use = not accelerated or device.type in acceleratedTypes
+    for idx, device in enumerate(cprefs.devices):
+        device.use = (not accelerated or device.type in acceleratedTypes) and idx == args.gpu
         print('Device enabled ({type}) = {enabled}'.format(type=device.type, enabled=device.use))
 
     return accelerated
 
-
 enable_cuda_devices()
 context.active_object.select_set(True)
 bpy.ops.object.delete()

train_3d.py

@@ -62,7 +62,7 @@ def launch_training(c, desc, outdir, dry_run):
     prev_run_ids = [int(x.group()) for x in prev_run_ids if x is not None]
     cur_run_id = max(prev_run_ids, default=-1) + 1
     if c.inference_vis:
-        c.run_dir = os.path.join(outdir, 'inference')
+        c.run_dir = os.path.join(outdir)
     else:
         c.run_dir = os.path.join(outdir, f'{cur_run_id:05d}-{desc}')
         assert not os.path.exists(c.run_dir)

training/inference_3d.py

@@ -51,9 +51,13 @@ def inference(
     bias_act._init()
     filtered_lrelu._init()
 
+    import random
     device = torch.device('cuda', rank)
     np.random.seed(random_seed * num_gpus + rank)
     torch.manual_seed(random_seed * num_gpus + rank)
+    random.seed(random_seed * num_gpus + rank)
+    torch.cuda.manual_seed(random_seed * num_gpus + rank)
+    torch.cuda.manual_seed_all(random_seed * num_gpus + rank)
     torch.backends.cudnn.enabled = True
     torch.backends.cudnn.benchmark = True  # Improves training speed.
     torch.backends.cuda.matmul.allow_tf32 = True  # Improves numerical accuracy.
@@ -77,23 +81,23 @@ def inference(
         G.load_state_dict(model_state_dict['G'], strict=True)
         G_ema.load_state_dict(model_state_dict['G_ema'], strict=True)
         # D.load_state_dict(model_state_dict['D'], strict=True)
-    grid_size = (5, 5)
+    grid_size = (1, 1)
     n_shape = grid_size[0] * grid_size[1]
     grid_z = torch.randn([n_shape, G.z_dim], device=device).split(1)  # random code for geometry
     grid_tex_z = torch.randn([n_shape, G.z_dim], device=device).split(1)  # random code for texture
     grid_c = torch.ones(n_shape, device=device).split(1)
 
     print('==> generate ')
-    save_visualization(
-        G_ema, grid_z, grid_c, run_dir, 0, grid_size, 0,
-        save_all=False,
-        grid_tex_z=grid_tex_z
-    )
+    # save_visualization(
+    #     G_ema, grid_z, grid_c, run_dir, 0, grid_size, 0,
+    #     save_all=False,
+    #     grid_tex_z=grid_tex_z
+    # )
 
     if inference_to_generate_textured_mesh:
         print('==> generate inference 3d shapes with texture')
         save_textured_mesh_for_inference(
-            G_ema, grid_z, grid_c, run_dir, save_mesh_dir='texture_mesh_for_inference',
+            G_ema, grid_z, grid_c, run_dir, save_mesh_dir="",
             c_to_compute_w_avg=None, grid_tex_z=grid_tex_z)
 
     if inference_save_interpolation:

training/inference_utils.py

@@ -247,7 +247,7 @@ def save_textured_mesh_for_inference(
                     all_uvs.data.cpu().numpy(),
                     mesh_f.data.cpu().numpy(),
                     all_mesh_tex_idx.data.cpu().numpy(),
-                    os.path.join(mesh_dir, '%07d.obj' % (save_mesh_idx))
+                    os.path.join(mesh_dir, 'mesh.obj')
                 )
                 lo, hi = (-1, 1)
                 img = np.asarray(tex_map.permute(1, 2, 0).data.cpu().numpy(), dtype=np.float32)
@@ -260,7 +260,7 @@ def save_textured_mesh_for_inference(
                 img = img * (1 - mask) + dilate_img * mask
                 img = img.clip(0, 255).astype(np.uint8)
                 PIL.Image.fromarray(np.ascontiguousarray(img[::-1, :, :]), 'RGB').save(
-                    os.path.join(mesh_dir, '%07d.png' % (save_mesh_idx)))
+                    os.path.join(mesh_dir, 'albedo.png'))
                 save_mesh_idx += 1
 
 

training/networks_get3d.py

@@ -375,6 +375,25 @@ def extract_3d_shape(
         else:
             mesh_v, mesh_f, sdf, deformation, v_deformed, sdf_reg_loss = self.get_geometry_prediction(ws_geo)
 
+        # Step 1.5: fix the mesh
+        import pymesh
+        mesh_v_processed = []
+        mesh_f_processed = []
+
+        for v, f in zip(mesh_v, mesh_f):
+            mesh = pymesh.form_mesh(vertices=v.cpu().numpy(), faces=f.cpu().numpy())
+            mesh, info = pymesh.remove_isolated_vertices(mesh)
+            mesh, info = pymesh.remove_duplicated_vertices(mesh, 1e-3)
+            print(info)
+
+            tensor_vertices = torch.from_numpy(mesh.vertices).float().to(self.device)
+            tensor_faces = torch.from_numpy(mesh.faces).long().to(self.device)
+            mesh_v_processed.append(tensor_vertices)
+            mesh_f_processed.append(tensor_faces)
+
+        mesh_v = mesh_v_processed
+        mesh_f = mesh_f_processed
+
         # Step 2: use x-atlas to get uv mapping for the mesh
         from training.extract_texture_map import xatlas_uvmap
         all_uvs = []

training/utils/utils_3d.py

@@ -37,7 +37,7 @@ def savemeshtes2(pointnp_px3, tcoords_px2, facenp_fx3, facetex_fx3, fname):
     fid.write('Ks 0.4 0.4 0.4\n')
     fid.write('Ns 10\n')
     fid.write('illum 2\n')
-    fid.write('map_Kd %s.png\n' % na)
+    fid.write('map_Kd albedo.png\n')
     fid.close()
     ####