Update RigNet/quick_start.py

RigNet/quick_start.py

@@ -90,64 +90,61 @@ def create_single_data(mesh_filename):
     # batch
     batch = torch.zeros(len(v), dtype=torch.long)
 
-    # ========== VOXEL: Use trimesh instead of binvox (NO X11 REQUIRED) ==========
+    # voxel - Use trimesh instead of binvox (no X11 required)
     binvox_file = mesh_filename.replace('_remesh.obj', '_normalized.binvox')
     
     if not os.path.exists(binvox_file):
-        print(f"  voxelizing mesh with trimesh (no binvox)...")
+        print(f"  voxelizing mesh (trimesh-based, no binvox)...")
         
         try:
             # Load mesh with trimesh
             mesh_tri = trimesh.load(normalized_obj)
             
-            # Voxelize to 88x88x88 grid
-            # Use exact=True for accurate inside/outside determination
-            voxel_grid = mesh_tri.voxelized(pitch=1.0/88, method='subdivide')
+            # Voxelize: create a 88x88x88 grid
+            # Calculate pitch to fit mesh in 88^3 grid
+            bounds = mesh_tri.bounds
+            dims = bounds[1] - bounds[0]
+            max_dim = max(dims)
+            pitch = max_dim / 88.0
             
-            # Get voxel matrix and ensure it's 88x88x88
-            vox_matrix = voxel_grid.matrix.astype(bool)
-            current_shape = vox_matrix.shape
-            
-            # Create a centered 88x88x88 array
-            final_vox = np.zeros((88, 88, 88), dtype=bool)
-            
-            # Calculate centering offsets
-            offset_x = max(0, (88 - current_shape[0]) // 2)
-            offset_y = max(0, (88 - current_shape[1]) // 2)
-            offset_z = max(0, (88 - current_shape[2]) // 2)
+            # Create voxel grid
+            voxel_grid = mesh_tri.voxelized(pitch=pitch)
             
-            # Copy data to centered position
-            end_x = min(offset_x + current_shape[0], 88)
-            end_y = min(offset_y + current_shape[1], 88)
-            end_z = min(offset_z + current_shape[2], 88)
-            
-            src_end_x = end_x - offset_x
-            src_end_y = end_y - offset_y
-            src_end_z = end_z - offset_z
+            # Ensure it's exactly 88x88x88 by padding/cropping
+            vox_matrix = voxel_grid.matrix
+            current_shape = vox_matrix.shape
             
-            final_vox[offset_x:end_x, offset_y:end_y, offset_z:end_z] = \
-                vox_matrix[:src_end_x, :src_end_y, :src_end_z]
+            # Pad to 88x88x88 if smaller
+            if any(s < 88 for s in current_shape):
+                padded = np.zeros((88, 88, 88), dtype=bool)
+                padded[:current_shape[0], :current_shape[1], :current_shape[2]] = vox_matrix
+                vox_matrix = padded
+            # Crop to 88x88x88 if larger
+            elif any(s > 88 for s in current_shape):
+                vox_matrix = vox_matrix[:88, :88, :88]
             
-            # Create binvox-compatible object
-            class VoxelData:
-                def __init__(self, data, dims, translate, scale):
+            # Create binvox-compatible object with ALL required attributes
+            class Voxels:
+                def __init__(self, data, dims, translate, scale, axis_order):
                     self.data = data
                     self.dims = dims
                     self.translate = translate
                     self.scale = scale
+                    self.axis_order = axis_order  # Required by binvox_rw
             
-            vox_obj = VoxelData(
-                data=final_vox,
+            vox_obj = Voxels(
+                data=vox_matrix,
                 dims=[88, 88, 88],
                 translate=[0.0, 0.0, 0.0],
-                scale=1.0
+                scale=1.0,
+                axis_order='xyz'  # Standard axis order
             )
             
             # Save as binvox format for caching
             with open(binvox_file, 'wb') as f:
                 binvox_rw.write(vox_obj, f)
             
-            print(f"  ✓ Voxelization complete: {binvox_file} (trimesh method)")
+            print(f"  ✓ Voxelization complete: {binvox_file}")
             
         except Exception as e:
             print(f"  ERROR: Trimesh voxelization failed: {e}")
@@ -159,6 +156,7 @@ def create_single_data(mesh_filename):
 
     data = Data(x=v[:, 3:6], pos=v[:, 0:3], tpl_edge_index=tpl_e, geo_edge_index=geo_e, batch=batch)
     return data, vox, surface_geodesic, translation_normalize, scale_normalize
+
 # def create_single_data(mesh_filename):
 #     """
 #     create input data for the network. The data is wrapped by Data structure in pytorch-geometric library