Update RigNet/quick_start.py

RigNet/quick_start.py

@@ -110,18 +110,26 @@ def create_single_data(mesh_filename):
             # Create voxel grid
             voxel_grid = mesh_tri.voxelized(pitch=pitch)
             
-            # Ensure it's exactly 88x88x88 by padding/cropping
+            # Get current voxel matrix
             vox_matrix = voxel_grid.matrix
             current_shape = vox_matrix.shape
             
-            # 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]
+            print(f"  Original voxel shape: {current_shape}")
+            
+            # Resize to exactly 88x88x88 by padding/cropping each dimension
+            target_shape = (88, 88, 88)
+            resized = np.zeros(target_shape, dtype=bool)
+            
+            # Calculate how much to copy in each dimension
+            x_size = min(current_shape[0], target_shape[0])
+            y_size = min(current_shape[1], target_shape[1])
+            z_size = min(current_shape[2], target_shape[2])
+            
+            # Copy the overlapping region
+            resized[:x_size, :y_size, :z_size] = vox_matrix[:x_size, :y_size, :z_size]
+            
+            vox_matrix = resized
+            print(f"  Resized voxel shape: {vox_matrix.shape}")
             
             # Create binvox-compatible object with ALL required attributes
             class Voxels:
@@ -130,14 +138,14 @@ def create_single_data(mesh_filename):
                     self.dims = dims
                     self.translate = translate
                     self.scale = scale
-                    self.axis_order = axis_order  # Required by binvox_rw
+                    self.axis_order = axis_order
             
             vox_obj = Voxels(
                 data=vox_matrix,
                 dims=[88, 88, 88],
                 translate=[0.0, 0.0, 0.0],
                 scale=1.0,
-                axis_order='xyz'  # Standard axis order
+                axis_order='xyz'
             )
             
             # Save as binvox format for caching
@@ -148,6 +156,8 @@ def create_single_data(mesh_filename):
             
         except Exception as e:
             print(f"  ERROR: Trimesh voxelization failed: {e}")
+            import traceback
+            traceback.print_exc()
             raise Exception(f"Voxelization failed: {e}")
     
     # Load voxel data
@@ -157,6 +167,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