custom training

Author: qq456cvb

README.md

@@ -192,7 +192,27 @@ python sunrgbd/eval.py | tee sunrgbd/map.txt
 </details>
 
 # Train on Your Own Object Collections
-Coming soon.
+
+<details>
+<summary><b>Configuration Explained</b></summary>
+
+To train on custom objects, it is necessary to understand some parameters in configuration files.
+- **up_sym**: Whether the objects look like a cylinder from up to bottom (e.g., bottles). This is to ensure the voting target is unambiguous.
+- **right_sym**: Whether the objects look like a cylinder from left to right (e.g., rolls). This is to ensure the voting target is unambiguous.
+- **regress_right**: Whether to predict the right axis. Some symmetric objects only have a up axis well defined (e.g., bowls, bottles), while some do not (e.g., laptops, mugs).
+- **z_right**: Whether the objects are placed such that the right axis is [0, 0, 1] (default: [1, 0, 0]).
+</details>
+<details>
+<summary><b>Voting Statistics Generation</b></summary>
+
+Next, we need to know the ``scale_range`` (used for data augmentation, control possible object scales along the diagonal), ``vote_range`` (the range for center voting targets $\mu,\nu$), and ``scale_mean`` (the average 3D scale, used for scale voting). To generate them, you may refer to ``gen_stats.py``.
+</details>
+
+<details>
+<summary><b>Write Configuration Files and Train</b></summary>
+
+After you prepare the necessary configurations and voting statistics, you can write your own configuration file similar to that in ``config/category``, and then run ``train.py``.
+</details>
 
 # Citation
 ```

gen_stats.py

@@ -0,0 +1,58 @@
+import argparse
+from utils.dataset import generate_target
+from utils.util import typename2shapenetid
+import os
+import open3d as o3d
+import numpy as np
+from tqdm import tqdm
+
+from utils.util import estimate_normals
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--category', default='bowl', help='Category name')
+    parser.add_argument('--shapenet_root', default='/home/neil/disk/ShapeNetCore.v2',  help='ShapeNet root directory')
+    parser.add_argument('--up_sym', action='store_true',  help='If the objects look similar to a cylinder from up ([0, 1, 0]) to bottom')
+    parser.add_argument('--right_sym', action='store_true',  help='If the objects look similar to a cylinder from left to right')
+    parser.add_argument('--z_right', action='store_true',  help='If the objects use [0, 0, 1] as the right axis (default, [1, 0, 0])')
+    args = parser.parse_args()
+    
+    name_path = 'data/shapenet_names/{}.txt'.format(args.category)
+    if os.path.exists(name_path):
+        shapenames = open(name_path).read().splitlines()
+    else:
+        shapenet_id = typename2shapenetid[args.category]
+        shapenames = os.listdir(os.path.join(args.shapenet_root, '{}'.format(shapenet_id)))
+        shapenames = [shapenet_id + '/' + name for name in shapenames]
+    
+    scale_range = [np.inf, -np.inf]
+    vote_range = [0, 0]
+    scale_mean = []
+    for model_name in tqdm(shapenames):
+        shapenet_cls, mesh_name = model_name.split('/')
+        path = os.path.join(args.shapenet_root, f'{shapenet_cls}/{mesh_name}/models/model_normalized.obj')
+        mesh = o3d.io.read_triangle_mesh(path)
+        pc = np.array(mesh.sample_points_uniformly(2048).points)
+        
+        # normalize to center
+        pc -= (np.min(pc, 0) + np.max(pc, 0)) / 2
+        
+        normals = estimate_normals(pc, 60)
+        targets_tr = generate_target(pc, normals, args.up_sym, args.right_sym, args.z_right, 100000)[0]
+        
+        diag_length = np.linalg.norm(np.max(pc, 0) - np.min(pc, 0))
+        
+        scale_range[0] = min(scale_range[0], diag_length)
+        scale_range[1] = max(scale_range[1], diag_length)
+        
+        vote_range[0] = max(vote_range[0], np.max(np.abs(targets_tr[:, 0])))
+        vote_range[1] = max(vote_range[1], np.max(targets_tr[:, 1]))
+        
+        scale_mean.append(np.max(pc, 0))
+    scale_mean = np.mean(scale_mean, 0)
+    
+    print(f'scale_range: {scale_range}')
+    print(f'vote_range: {vote_range}')
+    print(f'scale_mean: {scale_mean}')
+