SCMayS/hydata
095
1import os2import sys3import glob4import numpy as np5import trimesh6from PIL import Image7import argparse8from pathlib import Path9import traceback10import random11from tqdm import tqdm12import itertools13 14# Set OpenGL platform to EGL before importing pyrender15os.environ['PYOPENGL_PLATFORM'] = 'egl'16print(f"Set PYOPENGL_PLATFORM to: {os.environ.get('PYOPENGL_PLATFORM')}")17 18try:19 import pyrender20 print(f"Successfully imported pyrender {pyrender.__version__}")21except Exception as e:22 print(f"Error importing pyrender: {e}")23 traceback.print_exc()24 sys.exit(1)25 26# Import core rendering functionality from visualize_glb_models27from visualize_glb_models import create_look_at_matrix, GLBRenderer28 29class RotationDatasetGenerator(GLBRenderer):30 """Generate dataset of GLB models with paired rotated views for rotation prediction training"""31 32 def __init__(self, output_dir="rotation_dataset", size=(512, 512), verbose=True):33 super().__init__(output_dir=output_dir, size=size, verbose=verbose)34 # Ensure outputs directory exists35 os.makedirs(output_dir, exist_ok=True)36 37 # Valid rotation angles (divisible by 30, between -180 and 180)38 self.rotation_angles = [-150, -120, -90, -60, -30, 30, 60, 90, 120, 150, 180]39 # Valid rotation axes40 self.rotation_axes = ['x', 'y', 'z']41 42 # Dictionary to track figure numbers for each glb file43 self.glb_to_fignum = {}44 self.current_fignum = 045 46 def extract_euler_angles(self, matrix):47 """Extract Euler angles (in degrees) from a transformation matrix."""48 # Extract rotation matrix (top-left 3x3)49 rotation_matrix = matrix[:3, :3]50 51 # Convert to Euler angles (in radians)52 euler_angles = trimesh.transformations.euler_from_matrix(rotation_matrix, 'sxyz')53 54 # Convert to degrees55 euler_degrees = np.degrees(euler_angles)56 57 # Return as a tuple of x, y, z angles58 return tuple(round(angle, 2) for angle in euler_degrees)59 60 def generate_specific_orientation(self, x_angle=0.0, y_angle=0.0, z_angle=0.0):61 """Generate a specific orientation using the given Euler angles (in degrees)"""62 # Convert angles to radians63 x_rad = np.radians(x_angle)64 y_rad = np.radians(y_angle)65 z_rad = np.radians(z_angle)66 67 # Create rotation matrices for each axis68 x_rot = trimesh.transformations.rotation_matrix(x_rad, [1, 0, 0])69 y_rot = trimesh.transformations.rotation_matrix(y_rad, [0, 1, 0])70 z_rot = trimesh.transformations.rotation_matrix(z_rad, [0, 0, 1])71 72 # Combine rotations (order matters: first z, then y, then x)73 rotation_matrix = trimesh.transformations.concatenate_matrices(x_rot, y_rot, z_rot)74 75 return rotation_matrix76 77 def render_separate_views(self, model, base_transform, initial_z_rotation, rotation_axis, rotation_angle, 78 output_folder, center, mesh_size, show_axes=False):79 """Render initial view and rotated view as separate files in a folder"""80 # Create scenes for both views81 scene1 = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0], ambient_light=[0.7, 0.7, 0.7])82 scene2 = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0], ambient_light=[0.7, 0.7, 0.7])83 #scene1 = pyrender.Scene(bg_color=[0.9, 0.9, 0.9, 1.0], ambient_light=[0.7, 0.7, 0.7])84 #scene2 = pyrender.Scene(bg_color=[0.9, 0.9, 0.9, 1.0], ambient_light=[0.7, 0.7, 0.7])85 86 # Apply initial z-rotation after base transform87 initial_z_rotation_matrix = self.create_rotation_matrix('z', initial_z_rotation)88 initial_transform = trimesh.transformations.concatenate_matrices(initial_z_rotation_matrix, base_transform)89 90 # Apply the initial transformation to scene191 self._add_model_to_scene(model, scene1, initial_transform)92 93 # Calculate the combined transformation (base + initial z-rotation + main rotation)94 main_rotation_matrix = self.create_rotation_matrix(rotation_axis, rotation_angle)95 combined_transform = trimesh.transformations.concatenate_matrices(main_rotation_matrix, initial_transform)96 97 self._add_model_to_scene(model, scene2, combined_transform)98 99 # Setup camera and lights for both scenes100 camera_distance = mesh_size * self.camera_distance_factor101 camera_distance = max(camera_distance, 1.0)102 103 # Standard front view camera (negative z-axis)104 eye = np.array([0, 0, camera_distance]) + center105 target = center106 up = [0, 1, 0] # Standard "up" direction107 camera_pose = create_look_at_matrix(eye, target, up)108 109 # Add camera and lights to both scenes110 camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0)111 scene1.add(camera, pose=camera_pose)112 scene2.add(camera, pose=camera_pose)113 114 # Add light to scenes115 self._add_lights_to_scene(scene1, camera_pose, center, camera_distance)116 self._add_lights_to_scene(scene2, camera_pose, center, camera_distance)117 118 # Add coordinate axes if requested119 if show_axes:120 # Scale axes size based on mesh size121 axis_length = mesh_size * 0.4 # Slightly smaller to avoid overwhelming the object122 123 # Place axes at the center of the object124 axis_pose1 = np.eye(4)125 axis_pose1[:3, 3] = center # Center position for first scene126 127 axis_pose2 = np.eye(4)128 axis_pose2[:3, 3] = center # Center position for second scene129 130 # Add the coordinate axes to both scenes131 self._add_coordinate_axes(scene1, pose=axis_pose1, axis_length=axis_length)132 self._add_coordinate_axes(scene2, pose=axis_pose2, axis_length=axis_length)133 134 # Render both scenes135 #r = pyrender.OffscreenRenderer(self.size[0], self.size[1], point_size=1.0, antialias_samples=4)136 r = pyrender.OffscreenRenderer(self.size[0], self.size[1], point_size=1.0)137 138 # Create the output folder if it doesn't exist139 os.makedirs(output_folder, exist_ok=True)140 141 # Get the base filename from the folder142 folder_name = os.path.basename(output_folder)143 144 # Render and save the initial orientation image with the new naming pattern145 image1, _ = r.render(scene1)146 ini_path = os.path.join(output_folder, f"{folder_name}_ini.png")147 Image.fromarray(image1).save(ini_path)148 149 # Render and save the rotated image with the new naming pattern150 image2, _ = r.render(scene2)151 rot_path = os.path.join(output_folder, f"{folder_name}_rot.png")152 Image.fromarray(image2).save(rot_path)153 154 r.delete()155 156 return output_folder157 158 def _add_coordinate_axes(self, scene, pose=None, axis_length=1.0):159 """Add coordinate axes visualization to scene (x:red, y:green, z:blue)"""160 if pose is None:161 pose = np.eye(4)162 163 origin = pose[:3, 3]164 165 # Create meshes for each axis166 for i, (axis, color) in enumerate(zip(['x', 'y', 'z'], 167 [[1.0, 0.0, 0.0], 168 [0.0, 1.0, 0.0], 169 [0.0, 0.0, 1.0]])):170 # Create direction vector based on the pose171 direction = np.zeros(3)172 direction[i] = 1.0173 direction = pose[:3, :3] @ direction # Transform direction by rotation part of pose174 175 # Create cylinder for the shaft176 cylinder_height = axis_length - 0.2 # Leave space for cone177 cylinder_radius = axis_length * 0.02 # Thin line178 179 # Create a trimesh cylinder180 cylinder = trimesh.creation.cylinder(181 radius=cylinder_radius,182 height=cylinder_height,183 sections=8184 )185 186 # Rotate cylinder to point in the right direction187 z_axis = np.array([0, 0, 1])188 rotation_axis = np.cross(z_axis, direction)189 rotation_angle = np.arccos(np.dot(z_axis, direction) / (np.linalg.norm(z_axis) * np.linalg.norm(direction)))190 191 if np.linalg.norm(rotation_axis) > 1e-6: # Avoid zero division192 rotation_matrix = trimesh.transformations.rotation_matrix(rotation_angle, rotation_axis)193 cylinder.apply_transform(rotation_matrix)194 195 # Position cylinder196 translation = origin + direction * cylinder_height / 2197 translation_matrix = trimesh.transformations.translation_matrix(translation)198 cylinder.apply_transform(translation_matrix)199 200 # Create a cone for the arrow tip201 cone_height = 0.2202 cone_radius = cylinder_radius * 2.5203 cone = trimesh.creation.cone(radius=cone_radius, height=cone_height, sections=8)204 205 # Rotate cone to point in the right direction206 if np.linalg.norm(rotation_axis) > 1e-6:207 cone.apply_transform(rotation_matrix)208 209 # Position cone at the end of the cylinder210 tip_translation = origin + direction * axis_length - direction * cone_height / 2211 tip_translation_matrix = trimesh.transformations.translation_matrix(tip_translation)212 cone.apply_transform(tip_translation_matrix)213 214 # Combine shaft and arrow into one mesh215 axis_mesh = trimesh.util.concatenate([cylinder, cone])216 217 # Apply material color218 material = pyrender.MetallicRoughnessMaterial(219 baseColorFactor=color + [1.0], # RGBA220 metallicFactor=0.0,221 roughnessFactor=0.5222 )223 224 # Add to scene225 mesh = pyrender.Mesh.from_trimesh(axis_mesh, material=material)226 scene.add(mesh)227 228 def create_rotation_matrix(self, axis, angle_degrees):229 """Create a rotation matrix about the camera-aligned axis by the given angle in degrees"""230 angle_rad = np.radians(angle_degrees)231 232 if axis == 'x':233 # Horizontal axis (left-right in the image)234 return trimesh.transformations.rotation_matrix(angle_rad, [1, 0, 0])235 elif axis == 'y': 236 # Vertical axis (up-down in the image)237 return trimesh.transformations.rotation_matrix(angle_rad, [0, 1, 0])238 elif axis == 'z':239 # Depth axis (in-out of the image)240 return trimesh.transformations.rotation_matrix(angle_rad, [0, 0, 1])241 else:242 raise ValueError(f"Invalid rotation axis: {axis}. Must be 'x', 'y', or 'z'")243 244 def render_paired_views(self, model, base_transform, initial_z_rotation, rotation_axis, rotation_angle, 245 output_path, center, mesh_size, show_axes=False):246 """Render initial view and rotated view side by side using PIL instead of matplotlib"""247 # Create scenes for both views248 scene1 = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0], ambient_light=[0.7, 0.7, 0.7])249 scene2 = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0], ambient_light=[0.7, 0.7, 0.7])250 251 # Apply initial z-rotation after base transform252 initial_z_rotation_matrix = self.create_rotation_matrix('z', initial_z_rotation)253 initial_transform = trimesh.transformations.concatenate_matrices(initial_z_rotation_matrix, base_transform)254 255 # Add model with initial transformation to scene1256 self._add_model_to_scene(model, scene1, initial_transform)257 258 # Calculate the combined transformation (base + initial z-rotation + main rotation)259 main_rotation_matrix = self.create_rotation_matrix(rotation_axis, rotation_angle)260 combined_transform = trimesh.transformations.concatenate_matrices(main_rotation_matrix, initial_transform)261 262 # Add model with combined transformation to scene2263 self._add_model_to_scene(model, scene2, combined_transform)264 265 # Setup camera and lights for both scenes266 camera_distance = mesh_size * self.camera_distance_factor267 camera_distance = max(camera_distance, 1.0)268 269 # Standard front view camera (negative z-axis)270 eye = np.array([0, 0, camera_distance]) + center271 target = center272 up = [0, 1, 0] # Standard "up" direction273 camera_pose = create_look_at_matrix(eye, target, up)274 275 # Add camera and lights to both scenes276 camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0)277 scene1.add(camera, pose=camera_pose)278 scene2.add(camera, pose=camera_pose)279 280 # Add light to scenes281 self._add_lights_to_scene(scene1, camera_pose, center, camera_distance)282 self._add_lights_to_scene(scene2, camera_pose, center, camera_distance)283 284 # Add coordinate axes if requested285 if show_axes:286 # Scale axes size based on mesh size287 axis_length = mesh_size * 0.4 # Slightly smaller to avoid overwhelming the object288 289 # Place axes at the center of the object290 axis_pose1 = np.eye(4)291 axis_pose1[:3, 3] = center # Center position for first scene292 293 axis_pose2 = np.eye(4)294 axis_pose2[:3, 3] = center # Center position for second scene295 296 # Apply the appropriate transform to the second scene's axes297 # to maintain consistency with the object's rotation298 axis_pose2 = trimesh.transformations.concatenate_matrices(main_rotation_matrix, axis_pose2)299 300 self._add_coordinate_axes(scene1, pose=axis_pose1, axis_length=axis_length)301 self._add_coordinate_axes(scene2, pose=axis_pose2, axis_length=axis_length)302 303 # Render both scenes304 #r = pyrender.OffscreenRenderer(self.size[0], self.size[1], point_size=1.0, antialias_samples=4)305 r = pyrender.OffscreenRenderer(self.size[0], self.size[1], point_size=1.0)306 307 image1, _ = r.render(scene1)308 image2, _ = r.render(scene2)309 310 r.delete()311 312 # Combine images side by side using PIL313 pil_img1 = Image.fromarray(image1)314 pil_img2 = Image.fromarray(image2)315 316 # Create a new image with twice the width (no gap)317 combined_width = pil_img1.width + pil_img2.width318 combined_height = max(pil_img1.height, pil_img2.height)319 combined_img = Image.new('RGB', (combined_width, combined_height))320 321 # Paste the two images side by side with no gap322 combined_img.paste(pil_img1, (0, 0))323 combined_img.paste(pil_img2, (pil_img1.width, 0))324 325 # Save the combined image326 combined_img.save(output_path)327 328 return output_path329 330 def _add_model_to_scene(self, model, scene, transform=None):331 """Add a model to a scene with the given transformation"""332 if transform is None:333 transform = np.eye(4)334 335 if isinstance(model, trimesh.Scene):336 # For a scene, add each mesh with its transform337 for name, geom in model.geometry.items():338 if not isinstance(geom, trimesh.Trimesh):339 continue340 341 try:342 # Get the geometry's transform343 try:344 geom_transform = model.graph.get(name)[0]345 except (ValueError, KeyError, IndexError):346 geom_transform = np.eye(4)347 348 # Combine with the provided transform349 combined_transform = trimesh.transformations.concatenate_matrices(transform, geom_transform)350 351 # Add to scene - CHANGED smooth=False to smooth=True352 mesh_pyrender = pyrender.Mesh.from_trimesh(geom, smooth=True)353 scene.add(mesh_pyrender, pose=combined_transform)354 except Exception as e:355 if self.verbose:356 print(f"Info: Skipped mesh {name}: {str(e)[:100]}")357 continue358 else:359 # For a single mesh, add it directly with the transform360 try:361 # CHANGED smooth=False to smooth=True362 mesh_pyrender = pyrender.Mesh.from_trimesh(model, smooth=True)363 scene.add(mesh_pyrender, pose=transform)364 except Exception as e:365 if self.verbose:366 print(f"Error adding mesh to scene: {e}")367 raise368 369 def _add_lights_to_scene(self, scene, camera_pose, center, camera_distance):370 """Add more balanced lighting to a scene"""371 # 1. Main light from camera direction (reduced intensity)372 main_light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=2.0) # Reduced from 4.0373 scene.add(main_light, pose=camera_pose)374 375 # 2. Add point lights around the object (fewer lights, reduced intensity)376 for light_angle in [0, 180]: # Reduced from 4 lights to 2377 light_angle_rad = np.radians(light_angle)378 lx = camera_distance * 0.8 * np.cos(light_angle_rad)379 ly = camera_distance * 0.3380 lz = camera_distance * 0.8 * np.sin(light_angle_rad)381 382 light_pose = np.eye(4)383 light_pose[:3, 3] = np.array([lx, ly, lz]) + center384 385 point_light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=1.0) # Reduced from 2.0386 scene.add(point_light, pose=light_pose)387 388 # 3. Add a soft light from above (reduced intensity)389 top_light_pose = np.eye(4)390 top_light_pose[:3, 3] = center + np.array([0, camera_distance, 0])391 top_light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1.5) # Reduced from 3.0392 scene.add(top_light, pose=top_light_pose)393 394 def generate_dataset(self, glb_paths, rotation_axes=None, rotation_angles=None, 395 x_angles=None, y_angles=None, z_angles=None, initial_z_rotations=None,396 show_axes=False, generation_mode='combined'):397 """Generate a dataset from multiple GLB files with specific orientations and rotations"""398 dataset_info = []399 400 # Use default rotation axes if not specified401 if rotation_axes is None:402 rotation_axes = self.rotation_axes403 404 # Use default rotation angles if not specified405 if rotation_angles is None:406 rotation_angles = self.rotation_angles407 408 # Define the orientation angles to use409 # Use custom x_angles if provided, otherwise use default410 if x_angles is None:411 x_angles = [0.0] # Default is fixed at 0412 # Use custom y_angles if provided, otherwise use default413 if y_angles is None:414 y_angles = [0, 45, 90, 135, 180, 225, 270, 315]415 # Use custom z_angles if provided, otherwise use default416 if z_angles is None:417 z_angles = [0, 180]418 # Use custom initial_z_rotations if provided, otherwise use default419 if initial_z_rotations is None:420 initial_z_rotations = [0.0] # Default is no initial z rotation421 422 # Calculate all orientation combinations423 orientation_configs = list(itertools.product(x_angles, y_angles, z_angles))424 425 # Calculate all rotation configurations (combinations of axis and angle)426 rotation_configs = list(itertools.product(rotation_axes, rotation_angles))427 428 # Calculate total number of samples429 total_samples = len(glb_paths) * len(orientation_configs) * len(initial_z_rotations) * len(rotation_configs)430 431 with tqdm(total=total_samples, desc="Generating dataset") as pbar:432 for i, glb_path in enumerate(glb_paths):433 if self.verbose:434 print(f"\nProcessing model {i+1}/{len(glb_paths)}: {Path(glb_path).name}")435 436 # Load the GLB model once per model to save time437 try:438 # In your mesh loading code439 model = trimesh.load(glb_path, process=False, force='mesh')440 441 # Get model bounds and center for camera positioning442 if isinstance(model, trimesh.Scene):443 center, mesh_size = self._get_scene_bounds(model)444 else:445 bounds = model.bounds446 center = (bounds[0] + bounds[1]) / 2447 mesh_size = np.max(bounds[1] - bounds[0])448 449 # Get the GLB file name450 mesh_name = Path(glb_path).stem451 452 # Check if this GLB file already has a figure number assigned453 if mesh_name not in self.glb_to_fignum:454 # Assign a new figure number455 self.glb_to_fignum[mesh_name] = self.current_fignum456 self.current_fignum += 1457 458 # Get the assigned figure number459 fig_num = self.glb_to_fignum[mesh_name]460 461 # Generate a sample for each orientation, initial z rotation, and main rotation combination462 sample_id = 0463 for x_angle, y_angle, z_angle in orientation_configs:464 # Generate the specific orientation465 base_transform = self.generate_specific_orientation(x_angle, y_angle, z_angle)466 467 # Extract orientation angles for CSV output468 orientation_angles = (x_angle, y_angle, z_angle)469 470 for initial_z_rotation in initial_z_rotations:471 for rotation_axis, rotation_angle in rotation_configs:472 # Create output filename with ground truth info - USING THE ORIGINAL NAMING CONVENTION473 # Remove the initial_z_rotation from the filename to match the original format474 base_name = f"fig{fig_num:04d}_{sample_id:03d}_{rotation_axis}_{rotation_angle}"475 476 if generation_mode == 'separate':477 # Create a folder for the separate images478 output_folder = os.path.join(self.output_dir, base_name)479 480 # Render the separate views481 self.render_separate_views(482 model=model,483 base_transform=base_transform,484 initial_z_rotation=initial_z_rotation,485 rotation_axis=rotation_axis,486 rotation_angle=rotation_angle,487 output_folder=output_folder,488 center=center,489 mesh_size=mesh_size,490 show_axes=show_axes491 )492 493 output_path = output_folder494 else:495 # Combined mode (default)496 output_path = os.path.join(self.output_dir, f"{base_name}.png")497 498 # Render the paired views499 self.render_paired_views(500 model=model,501 base_transform=base_transform,502 initial_z_rotation=initial_z_rotation,503 rotation_axis=rotation_axis,504 rotation_angle=rotation_angle,505 output_path=output_path,506 center=center,507 mesh_size=mesh_size,508 show_axes=show_axes509 )510 511 # Add to dataset info512 dataset_info.append({513 "file_path": output_path,514 "model": mesh_name,515 "sample_id": sample_id,516 "initial_z_rotation": initial_z_rotation,517 "rotation_axis": rotation_axis,518 "rotation_angle": rotation_angle,519 "orientation_x": orientation_angles[0],520 "orientation_y": orientation_angles[1],521 "orientation_z": orientation_angles[2],522 "generation_mode": generation_mode523 })524 525 sample_id += 1526 pbar.update(1)527 528 except Exception as e:529 if self.verbose:530 print(f"Error processing {glb_path}: {e}")531 traceback.print_exc()532 # Update progress bar for skipped configurations533 skipped_count = len(orientation_configs) * len(initial_z_rotations) * len(rotation_configs)534 pbar.update(skipped_count)535 536 # Create a CSV file with dataset information537 self.save_dataset_info(dataset_info)538 539 return dataset_info540 541 def _get_scene_bounds(self, scene):542 """Get the bounds and size of a scene"""543 # Get overall scene bounds544 bounds = np.zeros((2, 3))545 center = np.zeros(3)546 mesh_size = 1.0 # Default size if we can't compute bounds547 548 # Try to calculate bounds from all geometries549 if len(scene.geometry) > 0:550 all_vertices = []551 552 # First collect all vertices from all geometries553 for name, geom in scene.geometry.items():554 if hasattr(geom, 'vertices') and len(geom.vertices) > 0:555 # Try to get transform for this geometry556 try:557 transform = scene.graph.get(name)[0]558 transformed_verts = trimesh.transformations.transform_points(geom.vertices, transform)559 all_vertices.append(transformed_verts)560 except (ValueError, KeyError, IndexError) as e:561 if self.verbose:562 print(f"Warning: Could not get transform for {name}, using identity")563 # Use identity transform as fallback564 all_vertices.append(geom.vertices)565 566 # If we have any vertices, compute bounds567 if all_vertices:568 combined_vertices = np.vstack(all_vertices)569 bounds[0] = np.min(combined_vertices, axis=0)570 bounds[1] = np.max(combined_vertices, axis=0)571 center = (bounds[0] + bounds[1]) / 2572 mesh_size = np.max(bounds[1] - bounds[0])573 else:574 if self.verbose:575 print("Warning: No usable vertices found in geometry, using default bounds")576 577 return center, mesh_size578 579 def save_dataset_info(self, dataset_info):580 """Save dataset information to a CSV file"""581 import csv582 583 csv_path = os.path.join(self.output_dir, "dataset_info.csv")584 585 with open(csv_path, 'w', newline='') as csvfile:586 # Add initial_z_rotation to fieldnames587 fieldnames = ['file_path', 'model', 'sample_id', 'initial_z_rotation', 'rotation_axis', 'rotation_angle', 588 'orientation_x', 'orientation_y', 'orientation_z', 'generation_mode']589 writer = csv.DictWriter(csvfile, fieldnames=fieldnames)590 591 writer.writeheader()592 for item in dataset_info:593 writer.writerow(item)594 595 if self.verbose:596 print(f"Dataset information saved to {csv_path}")597 598def main():599 parser = argparse.ArgumentParser(description="Generate a dataset of 3D model rotations")600 parser.add_argument('--input-dir', type=str, default="/mnt/data/yma71/vrg/0330/Hunyuan3D-2/objaverse_downloads",601 help="Directory containing GLB files")602 parser.add_argument('--output-dir', type=str, default="/mnt/data/yma71/vrg/0330/Hunyuan3D-2/rotation_dataset",603 help="Directory to save generated dataset")604 parser.add_argument('--size', type=str, default="1024,1024",605 help="Render size as width,height (e.g. '1024,1024')")606 parser.add_argument('--random-seed', type=int, default=42,607 help="Random seed for reproducibility")608 parser.add_argument('--rotation-axes', type=str, default=None,609 help="Comma-separated list of rotation axes (e.g. 'x,y,z')")610 parser.add_argument('--rotation-angles', type=str, default=None,611 help="Comma-separated list of rotation angles (e.g. '30,60,90')")612 parser.add_argument('--model', type=str, default=None,613 help="Use a specific GLB file instead of all in input-dir")614 parser.add_argument('--show-axes', action='store_true',615 help="Show coordinate axes in the renderings")616 parser.add_argument('--generation-mode', type=str, choices=['combined', 'separate'], default='combined',617 help="Mode for generating images (combined = one image with both views, separate = folder with two images)")618 parser.add_argument('--verbose', action='store_true',619 help="Print verbose output")620 parser.add_argument('--x-angles', type=str, default=None,621 help="Comma-separated list of x angles (e.g. '0,45,90')")622 parser.add_argument('--y-angles', type=str, default=None,623 help="Comma-separated list of y angles (e.g. '0,45,90,135,180,225,270,315')")624 parser.add_argument('--z-angles', type=str, default=None,625 help="Comma-separated list of z angles (e.g. '0,180')")626 # Add new argument for initial z rotations627 parser.add_argument('--initial-z-rotations', type=str, default=None,628 help="Comma-separated list of initial z rotations (e.g. '0,90,180,270')")629 630 args = parser.parse_args()631 632 # Set random seed for reproducibility (still useful for some operations)633 random.seed(args.random_seed)634 np.random.seed(args.random_seed)635 636 # Parse render size637 size = tuple(map(int, args.size.split(',')))638 639 # Create dataset generator640 generator = RotationDatasetGenerator(641 output_dir=args.output_dir,642 size=size,643 verbose=args.verbose644 )645 646 # Get GLB files to process647 if args.model:648 if not os.path.exists(args.model):649 print(f"GLB file not found: {args.model}")650 return651 glb_files = [args.model]652 else:653 # Search recursively for GLB files in input_dir and all subdirectories654 glb_files = []655 for root, dirs, files in os.walk(args.input_dir):656 for file in files:657 if file.lower().endswith('.glb'):658 glb_files.append(os.path.join(root, file))659 glb_files = sorted(glb_files)660 661 if not glb_files:662 print(f"No GLB files found in {args.input_dir} or its subdirectories")663 return664 665 print(f"Found {len(glb_files)} GLB files to process")666 if args.verbose:667 # Print first few files to verify correct discovery668 for i, file in enumerate(glb_files[:5]):669 print(f" {i+1}. {file}")670 if len(glb_files) > 5:671 print(f" ... and {len(glb_files)-5} more files")672 673 # Parse rotation axes if provided674 rotation_axes = None675 if args.rotation_axes:676 rotation_axes = args.rotation_axes.split(',')677 # Validate rotation axes678 for axis in rotation_axes:679 if axis not in ['x', 'y', 'z']:680 print(f"Invalid rotation axis: {axis}. Must be one of 'x', 'y', 'z'")681 return682 683 # Parse rotation angles if provided684 rotation_angles = None685 if args.rotation_angles:686 try:687 rotation_angles = [int(angle) for angle in args.rotation_angles.split(',')]688 689 # Make sure angles are within the valid range690 for angle in rotation_angles:691 if not (-180 <= angle <= 180):692 print(f"All rotation angles must be between -180 and 180 degrees. Got {angle}")693 return694 695 except ValueError as e:696 print(f"Invalid rotation angles format. Use comma-separated integers: {e}")697 return698 699 # Parse x angles if provided700 x_angles = None701 if args.x_angles:702 try:703 x_angles = [float(angle) for angle in args.x_angles.split(',')]704 except ValueError as e:705 print(f"Invalid x angles format. Use comma-separated numbers: {e}")706 return707 708 # Parse y angles if provided709 y_angles = None710 if args.y_angles:711 try:712 y_angles = [float(angle) for angle in args.y_angles.split(',')]713 except ValueError as e:714 print(f"Invalid y angles format. Use comma-separated numbers: {e}")715 return716 717 # Parse z angles if provided718 z_angles = None719 if args.z_angles:720 try:721 z_angles = [float(angle) for angle in args.z_angles.split(',')]722 except ValueError as e:723 print(f"Invalid z angles format. Use comma-separated numbers: {e}")724 return725 726 # Parse initial z rotations if provided727 initial_z_rotations = None728 if args.initial_z_rotations:729 try:730 initial_z_rotations = [float(angle) for angle in args.initial_z_rotations.split(',')]731 except ValueError as e:732 print(f"Invalid initial z rotations format. Use comma-separated numbers: {e}")733 return734 735 # Update print statements to include all angles736 print("Using initial orientations:")737 if x_angles:738 print(f" x = {x_angles}")739 else:740 print(" x = [0.0]")741 if y_angles:742 print(f" y = {y_angles}")743 else:744 print(" y = [0, 45, 90, 135, 180, 225, 270, 315]")745 if z_angles:746 print(f" z = {z_angles}")747 else:748 print(" z = [0, 180]")749 750 # Print initial z rotations751 if initial_z_rotations:752 print(f"Using initial z rotations: {initial_z_rotations}")753 else:754 print("Using default initial z rotation: [0.0]")755 756 # Update calculation of total combinations757 x_count = len(x_angles) if x_angles else 1 # Default is [0.0]758 y_count = len(y_angles) if y_angles else 8 # Default is [0, 45, 90, 135, 180, 225, 270, 315]759 z_count = len(z_angles) if z_angles else 2 # Default is [0, 180]760 initial_z_count = len(initial_z_rotations) if initial_z_rotations else 1 # Default is [0.0]761 762 if rotation_axes:763 print(f"Using rotation axes: {rotation_axes}")764 else:765 print(f"Using default rotation axes: {generator.rotation_axes}")766 767 if rotation_angles:768 print(f"Using rotation angles: {rotation_angles}")769 else:770 print(f"Using default rotation angles: {generator.rotation_angles}")771 772 if args.show_axes:773 print("Coordinate axes will be shown in the renderings")774 775 rot_axis_count = len(rotation_axes) if rotation_axes else len(generator.rotation_axes)776 rot_angle_count = len(rotation_angles) if rotation_angles else len(generator.rotation_angles)777 778 total_orientations = x_count * y_count * z_count779 total_rotations = rot_axis_count * rot_angle_count780 total_combinations_per_model = total_orientations * initial_z_count * total_rotations781 782 print(f"Total orientations: {total_orientations}")783 print(f"Total initial z rotations per orientation: {initial_z_count}")784 print(f"Total rotations per initial state: {total_rotations}")785 print(f"Total combinations per model: {total_combinations_per_model}")786 print(f"Total samples to generate: {total_combinations_per_model * len(glb_files)}")787 788 # Generate the dataset with new parameters789 dataset_info = generator.generate_dataset(790 glb_files,791 rotation_axes=rotation_axes,792 rotation_angles=rotation_angles,793 x_angles=x_angles,794 y_angles=y_angles,795 z_angles=z_angles,796 initial_z_rotations=initial_z_rotations,797 show_axes=args.show_axes,798 generation_mode=args.generation_mode799 )800 801 print(f"\nDataset generation complete!")802 print(f"Generated {len(dataset_info)} samples")803 print(f"Results saved to {args.output_dir}")804 805if __name__ == "__main__":806 try:807 main()808 except Exception as e:809 print(f"Unhandled exception in main: {e}")810 traceback.print_exc()