Add support for rendering samples from objaverse

README.md

@@ -20,6 +20,7 @@ Generate synthetic image datasets by composing multiple 3D objects on a plane an
 
 ### Key features
 - **Asset composition**: place GLB/GLTF/OBJ assets plus optional basic shapes on a ground plane with collision-free grid placement.
+- **Objaverse integration**: directly load and render 3D models from the [Objaverse dataset](https://huggingface.co/datasets/allenai/objaverse) (800K+ objects) via on-the-fly downloads.
 - **Lighting**: sample HDRI environment maps (EXR/HDR) with optional flips, rotations, and exposure scalings.
 - **Camera/motion**: orbiting camera, oscillating camera, orbiting environment light, object rotation, object vertical translation, and dolly zoom.
 - **Outputs**: RGB, normal, depth, albedo, roughness, metallic, and `meta.json` with camera and environment parameters. Optional debug outputs (environment projections, placement grid, `.blend`).
@@ -55,13 +56,20 @@ python blender_datagen_compose.py --config configs/render_orbit_cam.yaml out_dir
 
 You can also switch to different configs for different video rendering.
 
+To render with models from the Objaverse dataset (800K+ objects):
+
+```bash
+python blender_datagen_compose.py --config configs/render_objaverse.yaml
+```
 
 ## Data and assets
 
-- 3D assets: point `base_path` to a directory of assets or a text file listing absolute paths. Supported: `.glb`, `.gltf`, `.obj`.
-- Environment maps: point `envlight` to a directory of `.exr`/`.hdr` files, a single file, or a `.txt` list.
-- Ground plane: set `placement_plane` to a plane `.glb` (e.g., `data/plane_basic/plane.glb`).
-- Basic shapes: set `baseshape_path` to a folder with shape `.glb` files (e.g., `data/basicshapes`).
+- **3D assets**: 
+  - **Local files**: point `base_path` to a directory of assets or a text file listing absolute paths. Supported: `.glb`, `.gltf`, `.obj`.
+  - **Objaverse dataset**: set `use_objaverse: true` and optionally provide `objaverse_csv` path specifying valid instance ids. Models are automatically downloaded from HuggingFace.
+- **Environment maps**: point `envlight` to a directory of `.exr`/`.hdr` files, a single file, or a `.txt` list.
+- **Ground plane**: set `placement_plane` to a plane `.glb` (e.g., `data/plane_basic/plane.glb`).
+- **Basic shapes**: set `baseshape_path` to a folder with shape `.glb` files (e.g., `data/basicshapes`).
 
 
 Edit the config paths to match your rendering assets.
@@ -158,7 +166,11 @@ Environment lighting
 - `num_lighting`: number of lighting variations per scene
 
 Object sampling and placement
-- `base_path`: directory or txt-file list of assets to sample
+- `base_path`: directory or txt-file list of assets to sample (local files)
+- `use_objaverse`: enable Objaverse dataset (default: `false`)
+- `objaverse_csv`: path to CSV with all valid(textured and filtered) Objaverse IDs (default: `data/objaverse_v1_ids.csv`)
+- `objaverse_selection`: path to TXT file which contains the instance ids of interest
+- `objaverse_temp_dir`: custom temp directory for downloads (default: system temp)
 - `glbs_per_scene`: number of main assets per scene
 - `glbs_scale_range`: [min, max] uniform scale
 - `glbs_rotation_range`: [min_deg, max_deg] yaw rotation before placement
@@ -219,7 +231,7 @@ Naming: `<lgt_idx>.<frame_idx>.<pass>.<ext>`
 ## Development
 
 - Main entrypoint: `blender_datagen_compose.py`
-- Utilities: `utils/blender_utils.py`, `utils/render_utils.py`, `utils/image_utils.py`
+- Utilities: `utils/blender_utils.py`, `utils/render_utils.py`, `utils/image_utils.py`, `utils/objaverse_utils.py`
 - Example configs: `configs/*.yaml`
 
 ## Acknowledgments

blender_datagen_compose.py

@@ -27,7 +27,7 @@
 from mathutils import Vector, Matrix, Quaternion
 from pathlib import Path
 from shapely.geometry import Polygon
-from utils import blender_utils, render_utils, image_utils
+from utils import blender_utils, render_utils, image_utils, objaverse_utils
 import matplotlib.pyplot as plt
 from omegaconf import OmegaConf, DictConfig
 from types import SimpleNamespace
@@ -76,7 +76,7 @@ def post_process_rendering(output_dir, feature_fmt='jpg', dump_video=False, vide
     for p in blender_passes:
         img_list = sorted(glob.glob(os.path.join(output_dir, f'{p}.*')))
         if p == 'normal' and len(img_list) > 0:
-            meta_file = json.load(open(os.path.join(output_dir, f'0000.meta.json')))
+            meta_file = json.load(open(os.path.join(output_dir, '0000.meta.json')))
             meta_frames = meta_file['frames']
         for img_path in img_list:
             img_basename = os.path.basename(img_path)
@@ -122,7 +122,7 @@ def post_process_rendering(output_dir, feature_fmt='jpg', dump_video=False, vide
     # Optionally dump videos from RGB frames grouped by lighting index
     if dump_video:
         # Accept any extension for RGB frames (png/jpg)
-        rgb_list = sorted(glob.glob(os.path.join(output_dir, f'*.rgb.*')))
+        rgb_list = sorted(glob.glob(os.path.join(output_dir, '*.rgb.*')))
         # Group frames by lighting index (first token)
         group_dict = {}
         for fp in rgb_list:
@@ -454,7 +454,7 @@ def render_scene(
 class GLTFFileManger:
     def __init__(
         self, files, random_sample=True, rescale=True, 
-        multi_sample_weight=None, check_bbox=False
+        multi_sample_weight=None, check_bbox=False, objaverse_manager=None
     ):
         # just assume files is a list of list of files
         self.files = files
@@ -468,44 +468,76 @@ def __init__(
         self.rescale = rescale
         self.multi_sample_weight = multi_sample_weight
         self.check_bbox = check_bbox
+        self.objaverse_manager = objaverse_manager
+        self.use_objaverse = objaverse_manager is not None
+        self.temp_files_to_cleanup = []
+
         if self.multi_sample_weight is not None:
             assert len(self.multi_sample_weight) == self.num_file_lists
 
     def __len__(self):
+        if self.use_objaverse:
+            return self.objaverse_manager.num_objects
         return self.num_files
 
     def __iter__(self):
         self.idx = 0
         return self
+
+    def cleanup_temp_files(self):
+        """Clean up any temporary files from this iteration"""
+        for file_path in self.temp_files_to_cleanup:
+            if self.objaverse_manager:
+                self.objaverse_manager.cleanup_file(file_path)
+        self.temp_files_to_cleanup = []
 
     def __next__(self): # inifi-gen
         self.idx += 1
         sample_success = False
         obj_mesh = None
+        file = None
+
         while not sample_success:
-            if self.random_sample:
-                idx = np.random.choice(self.num_file_lists, p=self.multi_sample_weight)
-                file = np.random.choice(self.files[idx])
+            # Handle Objaverse sampling
+            if self.use_objaverse:
+                obj_data = self.objaverse_manager.sample_random()
+                if obj_data is None:
+                    logger.warning("Failed to download Objaverse object, retrying...")
+                    continue
+                file = obj_data['file_path']
+                self.temp_files_to_cleanup.append(file)
+                obj_name = f"objaverse_{obj_data['object_id']}"
+            # Handle local file sampling
             else:
-                file = self.files_list[(self.idx-1) % self.num_files]
-                sample_success = True
+                if self.random_sample:
+                    idx = np.random.choice(self.num_file_lists, p=self.multi_sample_weight)
+                    file = np.random.choice(self.files[idx])
+                else:
+                    file = self.files_list[(self.idx-1) % self.num_files]
+                    sample_success = True
+                obj_name = file
+
             try:
                 logger.info(f'Processing: {file}')
                 _obj_mesh = blender_utils.add_object_file(
                     file, with_empty=True, recenter=True, rescale=self.rescale
                 )
-                if self.check_bbox and not check_msh_bbox(obj_mesh):
+                if self.check_bbox and not check_msh_bbox(_obj_mesh):
                     _obj_mesh.clear_objects()
                     del _obj_mesh
                     raise Exception('Invalid bbox')
                 else:
                     obj_mesh = _obj_mesh
                     sample_success = True
             except Exception as e:
+                import traceback
+                traceback.print_exc()
                 logger.info(f'---> Error: {e}, skipping file {file}')
-                self.files[idx].remove(file)
+                if not self.use_objaverse:
+                    self.files[idx].remove(file)
+                # For Objaverse, just continue to next sample
 
-        return {'mesh': obj_mesh, 'name': file}
+        return {'mesh': obj_mesh, 'name': obj_name}
 
 def main():
     # Parse --config and support legacy flags; additional overrides use OmegaConf dotlist (key=val)
@@ -592,7 +624,9 @@ def main():
         # Unsupported/unused but kept for completeness
         'analytical_sky': False,
         'use_objaverse': False,
+        'objaverse_csv': 'data/objaverse_v1_ids.csv',
         'objaverse_selection': None,
+        'objaverse_temp_dir': None,
         'num_frames': 8,
         'sample_shape_texture': False,
     }
@@ -838,18 +872,39 @@ def sample_shape(shapes_files, obj_idx=0):
                     if k in texture:
                         texture_sample_weight[i] = w
         FLAGS.texture_sample_weight = texture_sample_weight / texture_sample_weight.sum()
+    else:
+        num_plane_textures = 0 
+
+    # Initialize Objaverse manager if needed
+    objaverse_manager = None
+    if FLAGS.use_objaverse:
+        logger.info("Initializing Objaverse dataset...")
+        objaverse_csv_path = os.path.abspath(FLAGS.objaverse_csv)
+        if not os.path.exists(objaverse_csv_path):
+            raise FileNotFoundError(f"Objaverse CSV not found: {objaverse_csv_path}")
+
+        # Load selection file if provided
+        selection = None
+        if FLAGS.objaverse_selection is not None:
+            selection = objaverse_utils.load_objaverse_selection(FLAGS.objaverse_selection)
 
-    obj_files = [] 
-    if not FLAGS.use_objaverse and FLAGS.base_path is not None:
+        objaverse_manager = objaverse_utils.ObjaverseManager(
+            csv_path=objaverse_csv_path,
+            temp_dir=FLAGS.objaverse_temp_dir,
+            selection=selection
+        )
+        obj_files = [[]]  # Empty placeholder for ObjaverseManager
+    elif FLAGS.base_path is not None:
+        # Load local files
         base_path_abs = os.path.abspath(FLAGS.base_path)
         if not os.path.isfile(base_path_abs):
-            obj_files.append([os.path.abspath(p) for p in (glob.glob(os.path.join(base_path_abs, "*.glb")) + \
+            obj_files = [[os.path.abspath(p) for p in (glob.glob(os.path.join(base_path_abs, "*.glb")) + \
                 glob.glob(os.path.join(base_path_abs, "*.gltf")) + \
-                glob.glob(os.path.join(base_path_abs, "*.obj")))])
+                glob.glob(os.path.join(base_path_abs, "*.obj")))]]
         else:
-            obj_files.append([os.path.abspath(p) for p in np.loadtxt(base_path_abs, dtype=str).tolist()])
+            obj_files = [[os.path.abspath(p) for p in np.loadtxt(base_path_abs, dtype=str).tolist()]]
     else:
-        raise NotImplementedError("Objaverse is not supported yet")
+        raise ValueError("Either use_objaverse must be True or base_path must be provided")
 
 
     # obj dataloader
@@ -858,9 +913,10 @@ def sample_shape(shapes_files, obj_idx=0):
         random_sample=FLAGS.glbs_random_sample, 
         rescale=FLAGS.glbs_rescale, 
         multi_sample_weight=FLAGS.glbs_multi_sample_weight, 
-        check_bbox=FLAGS.glbs_check_bbox
+        check_bbox=FLAGS.glbs_check_bbox,
+        objaverse_manager=objaverse_manager
     )
-    logger.info(f"Use {len(obj_dataloader)} glbs")
+    logger.info(f"Use {len(obj_dataloader)} objects ({'Objaverse' if FLAGS.use_objaverse else 'local files'})")
 
     # envlight 
     env_src = os.path.abspath(FLAGS.envlight)
@@ -984,18 +1040,25 @@ def sample_shape(shapes_files, obj_idx=0):
 
         if FLAGS.dump_blend:
             # save blender scene
-            bpy.ops.wm.save_as_mainfile(filepath=os.path.join(FLAGS.out_dir, name, f"scene.blend"))
+            bpy.ops.wm.save_as_mainfile(filepath=os.path.join(FLAGS.out_dir, name, "scene.blend"))
         if FLAGS.dump_placement:
             # Plot placement grid with some color map
             fig = plt.figure()
             plt.imshow(placement_grid.T, cmap='tab20', vmin=0, vmax=placement_grid.max().item() + 1)
             # plt.axis('off')
-            plt.savefig(os.path.join(FLAGS.out_dir, name, f"placement.png"))
+            plt.savefig(os.path.join(FLAGS.out_dir, name, "placement.png"))
             plt.close(fig)
         if FLAGS.dump_complete:
             new_complete_file = os.path.join(FLAGS.out_dir, f"COMPLETE_{name}")
             open(new_complete_file, 'w').close()
+
+        # Clean up temporary Objaverse files for this iteration
+        obj_dataloader.cleanup_temp_files()
 
+    # Final cleanup
+    if objaverse_manager is not None:
+        objaverse_manager.cleanup_all()
+
     # clean up and safely exit blender
     blender_utils.clear_scene()
     bpy.ops.wm.quit_blender()

configs/render_objaverse.yaml

@@ -0,0 +1,69 @@
+# Objaverse Rendering Configuration
+# Renders compositional scenes using models from the Objaverse dataset
+
+# Objaverse dataset settings
+use_objaverse: true
+objaverse_csv: 'data/objaverse_v1_ids.csv'
+objaverse_selection: null  # Optional: path to txt file with specific object IDs
+objaverse_temp_dir: null   # Optional: custom temp directory for downloads
+
+# Scene composition
+glbs_per_scene: 2          # Number of Objaverse models per scene
+glbs_random_sample: true
+glbs_rescale: true
+glbs_check_bbox: true
+glbs_scale_range: [1.0, 1.5]
+glbs_rotation_range: [-45, 45]
+glbs_placement_bbox: [-1.0, -1.0, 1.0, 1.0]
+
+# Optional: Add basic shapes
+baseshape_path: 'data/basicshapes'
+shapes_per_scene: 2
+shapes_scale_range: [0.3, 0.8]
+shapes_rotation_range: [-45, 45]
+shapes_placement_bbox: [-1.5, -1.5, 1.5, 1.5]
+
+# Placement plane
+placement_plane: 'data/plane_basic/plane.glb'
+placement_plane_offset: [0, 0, -0.5]
+placement_plane_scale: 10
+
+# Optional: Apply textures to plane
+placement_plane_textures: null  # e.g., 'data/textures'
+
+# Rendering settings
+num_rendering: 10
+num_frames: 8
+num_lighting: 1
+resolution: [256, 256]
+spp: 128
+use_denoise: 'OPTIX'
+transparent_bg: true
+
+# Video mode
+video_mode: 'orbit_cam'  # Options: orbit_cam, oscil_cam, dolly_cam, orbit_lgt, rotat_obj, vtran_obj
+
+# Camera settings
+radius_range: [2.0, 4.0]
+varying_radius: false
+fov_range: [45.0, 45.0]
+cam_phi_range: [0, 360]
+cam_theta_range: [0, 90]
+
+# Environment lighting
+envlight: 'data/envmaps'
+env_scale: 1.0
+random_env_rotation: true
+random_env_flip: true
+random_env_scale: null  # [0.5, 2.0] for random exposure
+
+# Output settings
+out_dir: 'output/objaverse_renders'
+dump_format: 'jpg'
+dump_video: true
+dump_complete: true
+dump_blend: false
+dump_placement: false
+
+# Seed (optional)
+seed: null