fix bug for training

utils/data_load.py

@@ -45,6 +45,15 @@ class DataLoadUtil:
         mesh.apply_transform(world_object_pose)
         return mesh
     
+    @staticmethod
+    def get_bbox_diag(model_dir, object_name):
+        model_path = os.path.join(model_dir, object_name, "mesh.obj")
+        mesh = trimesh.load(model_path)
+        bbox = mesh.bounding_box.extents
+        diagonal_length = np.linalg.norm(bbox)
+        return diagonal_length
+    
+    
     @staticmethod
     def save_mesh_at(model_dir, output_dir, object_name, scene_name, world_object_pose):
         mesh = DataLoadUtil.load_mesh_at(model_dir, object_name, world_object_pose)
@@ -192,6 +201,24 @@ class DataLoadUtil:
             "points_world": target_points_world,
             "points_camera": target_points_camera
         }
+    
+    @staticmethod
+    def get_point_cloud(depth, cam_intrinsic, cam_extrinsic):
+        h, w = depth.shape
+        i, j = np.meshgrid(np.arange(w), np.arange(h), indexing='xy')
+        
+        z = depth
+        x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
+        y = (j - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
+        
+        points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
+        points_camera_aug = np.concatenate([points_camera, np.ones((points_camera.shape[0], 1))], axis=-1)
+        
+        points_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
+        return {
+            "points_world": points_world,
+            "points_camera": points_camera
+        }
         
     @staticmethod
     def get_target_point_cloud_world_from_path(path, binocular=False, random_downsample_N=65536, voxel_size = 0.005, target_mask_label=(0,255,0,255)):