SCMayS/hydata
089
1import os2import sys3import glob4import numpy as np5import trimesh6import matplotlib.pyplot as plt7from PIL import Image8import argparse9from pathlib import Path10import traceback11 12# Enable verbose debugging13print("Starting GLB visualization script...")14print(f"Python version: {sys.version}")15print(f"Current directory: {os.getcwd()}")16 17# Set OpenGL platform to EGL before importing pyrender18os.environ['PYOPENGL_PLATFORM'] = 'egl'19print(f"Set PYOPENGL_PLATFORM to: {os.environ.get('PYOPENGL_PLATFORM')}")20 21try:22 import pyrender23 print(f"Successfully imported pyrender {pyrender.__version__}")24except Exception as e:25 print(f"Error importing pyrender: {e}")26 traceback.print_exc()27 sys.exit(1)28 29def create_look_at_matrix(eye, target, up=None):30 """Create a 'look at' transformation matrix for camera positioning31 32 Args:33 eye: Camera position34 target: Point the camera is looking at35 up: Up direction (defaults to [0,0,1] if not provided)36 37 Returns:38 4x4 transformation matrix39 """40 if up is None:41 up = np.array([0.0, 0.0, 1.0]) # Default up direction42 43 # Calculate forward direction (z)44 forward = np.array(target) - np.array(eye)45 forward = forward / np.linalg.norm(forward)46 47 # Calculate right direction (x)48 right = np.cross(forward, up)49 if np.linalg.norm(right) < 1e-6:50 # If forward and up are parallel, choose a different up vector51 alternate_up = np.array([0.0, 1.0, 0.0]) if np.allclose(up, [0,0,1]) else np.array([0.0, 0.0, 1.0])52 right = np.cross(forward, alternate_up)53 right = right / np.linalg.norm(right)54 55 # Recalculate the orthogonal up vector (y)56 corrected_up = np.cross(right, forward)57 corrected_up = corrected_up / np.linalg.norm(corrected_up)58 59 # Create rotation matrix - each column is one of the basis vectors60 rotation = np.eye(4)61 rotation[:3, 0] = right62 rotation[:3, 1] = corrected_up63 rotation[:3, 2] = -forward # Negate because in OpenGL, camera looks down -z64 65 # Create translation matrix66 translation = np.eye(4)67 translation[:3, 3] = -np.array(eye) 68 69 # Combine rotation and translation70 result = np.matmul(rotation, translation)71 72 # Instead, just set the position directly and keep the calculated rotation73 result[:3, 3] = np.array(eye)74 75 return result76 77class GLBRenderer:78 """Class to render 3D meshes from GLB files using PyRender with EGL backend"""79 80 def __init__(self, output_dir="glb_visualizations", size=(512, 512), verbose=True):81 self.output_dir = output_dir82 self.size = size83 self.verbose = verbose84 os.makedirs(output_dir, exist_ok=True)85 86 # Default camera parameters87 self.camera_distance_factor = 1.588 self.camera_elevation = 3089 90 if verbose:91 print(f"GLB Renderer initialized with output directory: {os.path.abspath(output_dir)}")92 print(f"Render size: {size[0]}x{size[1]}")93 print(f"Using PYOPENGL_PLATFORM: {os.environ.get('PYOPENGL_PLATFORM', 'default')}")94 95 def render_model(self, glb_path, angles=None, contact_sheet=True):96 """Render multiple views of a GLB model and optionally create a contact sheet"""97 if angles is None:98 angles = [0, 45, 90, 135, 180, 225, 270, 315]99 100 model_name = Path(glb_path).stem101 model_output_dir = os.path.join(self.output_dir, model_name)102 os.makedirs(model_output_dir, exist_ok=True)103 104 if self.verbose:105 print(f"Rendering GLB model: {model_name}")106 print(f"Output directory: {model_output_dir}")107 108 # Load the model109 try:110 print(f"Loading GLB from: {glb_path}")111 model = trimesh.load(glb_path)112 113 # Get model bounds and center for camera positioning114 if isinstance(model, trimesh.Scene):115 print(f"Loaded GLB as a Scene with {len(model.geometry)} geometries")116 117 # Get overall scene bounds118 bounds = np.zeros((2, 3))119 center = np.zeros(3)120 mesh_size = 1.0 # Default size if we can't compute bounds121 122 # Try to calculate bounds from all geometries, regardless of transform123 if len(model.geometry) > 0:124 all_vertices = []125 126 # First collect all vertices from all geometries127 for name, geom in model.geometry.items():128 if hasattr(geom, 'vertices') and len(geom.vertices) > 0:129 # Try to get transform for this geometry130 try:131 transform = model.graph.get(name)[0]132 transformed_verts = trimesh.transformations.transform_points(geom.vertices, transform)133 all_vertices.append(transformed_verts)134 except (ValueError, KeyError, IndexError) as e:135 print(f"Warning: Could not get transform for {name}, using identity. Error: {e}")136 # Use identity transform as fallback137 all_vertices.append(geom.vertices)138 139 # If we have any vertices, compute bounds140 if all_vertices:141 # Combine all vertices142 combined_vertices = np.vstack(all_vertices)143 bounds[0] = np.min(combined_vertices, axis=0)144 bounds[1] = np.max(combined_vertices, axis=0)145 center = (bounds[0] + bounds[1]) / 2146 mesh_size = np.max(bounds[1] - bounds[0])147 else:148 print("Warning: No usable vertices found in geometry, using default bounds")149 bounds[0] = np.array([-1.0, -1.0, -1.0])150 bounds[1] = np.array([1.0, 1.0, 1.0])151 center = np.zeros(3)152 mesh_size = 2.0153 else:154 # Empty scene, use default bounds155 print("Warning: Empty scene, using default bounds")156 bounds[0] = np.array([-1.0, -1.0, -1.0])157 bounds[1] = np.array([1.0, 1.0, 1.0])158 center = np.zeros(3)159 mesh_size = 2.0160 else:161 print(f"Loaded GLB as a single Mesh with {len(model.vertices)} vertices and {len(model.faces)} faces")162 bounds = model.bounds163 center = (bounds[0] + bounds[1]) / 2164 mesh_size = np.max(bounds[1] - bounds[0])165 166 if self.verbose:167 print(f"Model bounds: {bounds}")168 print(f"Model center: {center}")169 print(f"Model size: {mesh_size}")170 171 except Exception as e:172 print(f"Error loading GLB {glb_path}: {e}")173 traceback.print_exc()174 return None175 176 # Render each angle177 render_paths = []178 for angle in angles:179 output_path = os.path.join(model_output_dir, f"angle_{angle:03d}.png")180 try:181 self._render_view(model, angle, output_path, center, mesh_size)182 render_paths.append(output_path)183 except Exception as e:184 print(f"Error rendering angle {angle}: {e}")185 traceback.print_exc()186 # Create a placeholder image187 img = Image.new('RGB', self.size, color=(200, 200, 200))188 img.save(output_path)189 render_paths.append(output_path)190 191 # Create contact sheet if requested192 if contact_sheet and render_paths:193 sheet_path = os.path.join(self.output_dir, f"{model_name}_contact_sheet.png")194 try:195 self._create_contact_sheet(render_paths, sheet_path)196 if self.verbose:197 print(f"Created contact sheet: {sheet_path}")198 except Exception as e:199 print(f"Error creating contact sheet: {e}")200 traceback.print_exc()201 202 return render_paths203 204 def _render_view(self, model, angle, output_path, center, mesh_size):205 """Render a single view of the GLB model"""206 print(f"Rendering view at angle {angle}°...")207 208 # Create a scene with stronger ambient light to prevent completely dark renders209 scene = pyrender.Scene(bg_color=[0.9, 0.9, 0.9, 1.0], ambient_light=[0.7, 0.7, 0.7])210 211 # Add model to the scene212 if isinstance(model, trimesh.Scene):213 # For a scene, add each mesh with its transform214 for name, geom in model.geometry.items():215 if not isinstance(geom, trimesh.Trimesh):216 # Skip non-mesh geometries217 continue218 219 try:220 # Get the transform for this geometry221 try:222 transform = model.graph.get(name)[0]223 except (ValueError, KeyError, IndexError):224 # If no valid transform path exists, use identity transform225 print(f"No valid transform path for {name}, using identity transform")226 transform = np.eye(4)227 228 # Convert trimesh to pyrender mesh229 mesh_pyrender = pyrender.Mesh.from_trimesh(geom, smooth=False)230 231 # Add to scene with the geometry's transform232 scene.add(mesh_pyrender, pose=transform)233 except Exception as e:234 print(f"Error adding mesh {name} to scene: {e}")235 continue236 else:237 # For a single mesh, add it directly238 try:239 mesh_pyrender = pyrender.Mesh.from_trimesh(model, smooth=False)240 scene.add(mesh_pyrender, pose=np.eye(4))241 except Exception as e:242 print(f"Error adding mesh to scene: {e}")243 raise244 245 # Calculate camera distance based on mesh size246 camera_distance = mesh_size * self.camera_distance_factor247 camera_distance = max(camera_distance, 1.0) # Ensure minimum distance248 249 print(f"Mesh size: {mesh_size}, Camera distance: {camera_distance}")250 251 # Position camera based on angle and elevation252 angle_rad = np.radians(angle)253 elevation_rad = np.radians(self.camera_elevation)254 255 # Calculate camera position256 x = camera_distance * np.cos(elevation_rad) * np.sin(angle_rad)257 y = camera_distance * np.cos(elevation_rad) * np.cos(angle_rad)258 z = camera_distance * np.sin(elevation_rad)259 eye = np.array([x, y, z]) + center260 261 # Create camera pose matrix using our custom function262 camera_pose = create_look_at_matrix(eye, center)263 264 # Create and add camera to scene265 print("Adding camera to scene...")266 camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0)267 scene.add(camera, pose=camera_pose)268 269 # Add multiple lights from different directions270 print("Adding lights to scene...")271 272 # 1. Main light from camera direction273 light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=4.0)274 scene.add(light, pose=camera_pose)275 276 # 2. Add several point lights around the object277 for light_angle in [0, 90, 180, 270]: # Lights at cardinal directions278 light_angle_rad = np.radians(light_angle)279 lx = camera_distance * 0.8 * np.sin(light_angle_rad)280 ly = camera_distance * 0.8 * np.cos(light_angle_rad)281 lz = camera_distance * 0.8 # Position lights above the object282 283 light_pose = np.eye(4)284 light_pose[:3, 3] = np.array([lx, ly, lz]) + center285 286 point_light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=2.0)287 scene.add(point_light, pose=light_pose)288 289 # Add a bright light from above290 top_light_pose = np.eye(4)291 top_light_pose[:3, 3] = center + np.array([0, 0, camera_distance])292 top_light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=3.0)293 scene.add(top_light, pose=top_light_pose)294 295 # Render296 print("Rendering scene...")297 r = pyrender.OffscreenRenderer(self.size[0], self.size[1])298 color, depth = r.render(scene)299 r.delete()300 301 # Save image302 print(f"Saving rendered image to: {output_path}")303 img = Image.fromarray(color)304 img.save(output_path)305 306 return output_path307 308 def _create_contact_sheet(self, image_paths, output_path, cols=4):309 """Create a contact sheet from multiple images"""310 if not image_paths:311 return None312 313 print(f"Creating contact sheet from {len(image_paths)} images...")314 315 # Determine rows and columns316 n_images = len(image_paths)317 rows = (n_images + cols - 1) // cols318 319 # Open first image to get dimensions320 with Image.open(image_paths[0]) as img:321 img_width, img_height = img.size322 323 # Create contact sheet324 sheet_width = cols * img_width325 sheet_height = rows * img_height326 contact_sheet = Image.new('RGB', (sheet_width, sheet_height), (255, 255, 255))327 328 # Add each image to contact sheet329 for i, img_path in enumerate(image_paths):330 if os.path.exists(img_path):331 try:332 img = Image.open(img_path)333 row = i // cols334 col = i % cols335 x = col * img_width336 y = row * img_height337 contact_sheet.paste(img, (x, y))338 except Exception as e:339 print(f"Error adding image {img_path} to contact sheet: {e}")340 341 # Save contact sheet342 contact_sheet.save(output_path)343 print(f"Contact sheet saved to: {output_path}")344 return output_path345 346def main():347 print("Entering main function...")348 349 parser = argparse.ArgumentParser(description="Render GLB files using PyRender with EGL")350 parser.add_argument('--input-dir', type=str, default="/mnt/data/yma71/vrg/0330/Hunyuan3D-2/objaverse_downloads",351 help="Directory containing GLB files")352 parser.add_argument('--output-dir', type=str, default="/mnt/data/yma71/vrg/0330/Hunyuan3D-2/glb_visualizations",353 help="Directory to save rendered images")354 parser.add_argument('--angles', type=str, default="0,45,90,135,180,225,270,315",355 help="Comma-separated list of camera angles")356 parser.add_argument('--size', type=str, default="512,512",357 help="Render size as width,height (e.g. '512,512')")358 parser.add_argument('--no-contact-sheet', action='store_true',359 help="Disable contact sheet generation")360 parser.add_argument('--elevation', type=float, default=30,361 help="Camera elevation angle in degrees")362 parser.add_argument('--distance', type=float, default=2.5,363 help="Camera distance factor (relative to model size)")364 parser.add_argument('--verbose', action='store_true',365 help="Print verbose output")366 parser.add_argument('--model', type=str, default=None,367 help="Render a specific GLB file instead of all models in input-dir")368 369 args = parser.parse_args()370 print(f"Parsed arguments: {args}")371 372 # Parse render size373 size = tuple(map(int, args.size.split(',')))374 375 # Parse angles376 angles = list(map(int, args.angles.split(',')))377 378 # Create renderer379 print("Creating GLB renderer...")380 renderer = GLBRenderer(381 output_dir=args.output_dir, 382 size=size, 383 verbose=args.verbose or True # Force verbose for debugging384 )385 386 # Set camera parameters387 renderer.camera_elevation = args.elevation388 renderer.camera_distance_factor = args.distance389 390 # Handle single model case391 if args.model:392 if not os.path.exists(args.model):393 print(f"GLB file not found: {args.model}")394 return395 396 print(f"Processing single GLB model: {args.model}")397 renderer.render_model(398 args.model,399 angles=angles,400 contact_sheet=not args.no_contact_sheet401 )402 return403 404 # Find all GLB files405 print(f"Looking for GLB files in: {args.input_dir}")406 glb_files = sorted(glob.glob(os.path.join(args.input_dir, "*.glb")))407 408 if not glb_files:409 print(f"No GLB files found in {args.input_dir}")410 return411 412 print(f"Found {len(glb_files)} GLB files to render: {[os.path.basename(f) for f in glb_files]}")413 414 # Process each file415 for i, glb_path in enumerate(glb_files):416 print(f"\n[{i+1}/{len(glb_files)}] Processing: {os.path.basename(glb_path)}")417 try:418 renderer.render_model(419 glb_path, 420 angles=angles, 421 contact_sheet=not args.no_contact_sheet422 )423 except Exception as e:424 print(f"Error processing {os.path.basename(glb_path)}: {e}")425 traceback.print_exc()426 427 print(f"\nVisualization complete! Results saved to {args.output_dir}")428 429if __name__ == "__main__":430 try:431 main()432 except Exception as e:433 print(f"Unhandled exception in main: {e}")434 traceback.print_exc()