Merge pull request #222 from princeton-vl/floating_indoors

Floating indoors

docs/CHANGELOG.md

@@ -78,4 +78,5 @@ v1.5.1
 
 v1.6.0
 - Add geometric tile pattern materials
-- Tune window parameters and materials
+- Tune window parameters and materials
+- Add floating object placement generator and example command

docs/HelloRoom.md

@@ -120,6 +120,15 @@ These settings are intended for debugging or for generating tailored datasets. I
 
 If you are using the commands from [Creating large datasets](#creating-large-datasets) you will instead add these configs as `--overrides` to the end of your command, rather than `-p`
 
+### Generate Rooms with Floating Objects
+
+To enable floating objects in a room, add the override `compose_indoors.floating_objs_enabled=True`. For example:
+```
+python -m infinigen_examples.generate_indoors --seed 0 --task coarse --output_folder outputs/indoors/coarse -g fast_solve.gin singleroom.gin -p compose_indoors.floating_objs_enabled=True compose_indoors.terrain_enabled=False restrict_solving.restrict_parent_rooms=\[\"DiningRoom\"\]
+
+```
+By default, between 15 and 25 objects are generated and have their size normalized to fit within a 0.5m cube. The number of objects can be configured with `compose_indoors.num_floating` and normalization can be disabled with `compose_indoors.norm_floating=False`. Collisions/intersections between floating objects and existing solved objects are off by default and can be enabled with `compose_indoors.enable_collision_floating=True` and `compose_indoors.enable_collision_solved=True`.
+
 ## Run unit tests
 ```
 pytest tests/ --disable-warnings

infinigen/assets/placement/floating_objects.py

@@ -0,0 +1,164 @@
+import logging
+
+import bmesh
+import bpy
+import numpy as np
+from mathutils.bvhtree import BVHTree
+
+from infinigen.core.placement.factory import AssetFactory
+from infinigen.core.util import blender as butil
+from infinigen.core.util.random import random_general as rg
+
+logger = logging.getLogger(__name__)
+
+
+def create_bvh_tree_from_object(obj):
+    bm = bmesh.new()
+    bm.from_mesh(obj.data)
+    bm.transform(obj.matrix_world)
+    bvh = BVHTree.FromBMesh(bm)
+    bm.free()
+    return bvh
+
+
+def check_bvh_intersection(bvh1, bvh2):
+    if isinstance(bvh2, list):
+        return any([check_bvh_intersection(bvh1, bvh) for bvh in bvh2])
+    else:
+        return bvh1.overlap(bvh2)
+
+
+def raycast_sample(min_dist, sensor_coords, pix_it, camera, bvhtree):
+    cam_location = camera.matrix_world.to_translation()
+
+    for _ in range(1500):
+        x = pix_it[np.random.randint(0, pix_it.shape[0])][0]
+        y = pix_it[np.random.randint(0, pix_it.shape[0])][1]
+
+        direction = (sensor_coords[y, x] - camera.matrix_world.translation).normalized()
+
+        location, normal, index, distance = bvhtree.ray_cast(cam_location, direction)
+
+        if location:
+            if distance <= min_dist:
+                continue
+            random_distance = np.random.uniform(min_dist, distance)
+            sampled_point = cam_location + direction.normalized() * random_distance
+            return sampled_point
+
+    logger.info("Couldnt find far enough away pixel to raycast to")
+    return None
+
+
+def bbox_sample(bbox):
+    raise NotImplementedError
+
+
+class FloatingObjectPlacement:
+    def __init__(
+        self,
+        generators: list[AssetFactory],
+        camera: bpy.types.Object,
+        background_objs: bpy.types.Object | list[bpy.types.Object],
+        collision_objs: bpy.types.Object | list[bpy.types.Object],
+        bbox=None,
+    ):
+        self.generators = generators
+        self.camera = camera
+
+        if not isinstance(background_objs, list):
+            background_objs = [background_objs]
+        self.background_objs = background_objs
+
+        if not isinstance(collision_objs, list):
+            collision_objs = [collision_objs]
+        self.collision_objs = collision_objs
+
+        self.bbox = bbox
+
+    def place_objs(
+        self,
+        num_objs,
+        min_dist=1,
+        sample_retries=200,
+        raycast=True,
+        normalize=False,
+        collision_placed=False,
+        collision_existing=False,
+    ):
+        background_copied = butil.join_objects(
+            [butil.copy(obj) for obj in self.background_objs]
+        )
+        room_bvh = create_bvh_tree_from_object(background_copied)
+        butil.delete(background_copied)
+
+        if len(self.collision_objs) > 0:
+            objs_copied = butil.join_objects(
+                [butil.copy(obj) for obj in self.collision_objs]
+            )
+            existing_obj_bvh = create_bvh_tree_from_object(objs_copied)
+            butil.delete(objs_copied)
+        else:
+            existing_obj_bvh = None
+
+        placed_obj_bvhs = []
+
+        from infinigen.core.placement.camera import get_sensor_coords
+
+        sensor_coords, pix_it = get_sensor_coords(self.camera, sparse=False)
+        num_place = rg(num_objs)
+        for i in range(num_place):
+            fac = np.random.choice(self.generators)(np.random.randint(1e7))
+            asset = fac.spawn_asset(0)
+            fac.finalize_assets([asset])
+            max_dim = max(asset.dimensions.x, asset.dimensions.y, asset.dimensions.z)
+
+            if normalize:
+                if max_dim != 0:
+                    normalize_scale = 0.5 / max(
+                        asset.dimensions.x, asset.dimensions.y, asset.dimensions.z
+                    )
+                else:
+                    normalize_scale = 1
+
+                asset.scale = (normalize_scale, normalize_scale, normalize_scale)
+
+            for j in range(sample_retries):
+                if raycast:
+                    point = raycast_sample(
+                        min_dist, sensor_coords, pix_it, self.camera, room_bvh
+                    )
+                else:
+                    point = bbox_sample()
+
+                if point is None:
+                    continue
+
+                asset.rotation_mode = "XYZ"
+                asset.rotation_euler = np.random.uniform(-np.pi, np.pi, 3)
+                asset.location = point
+
+                bpy.context.view_layer.update()  # i can redo this later without view updates if necessary, but currently it doesn't incur significant overhead
+                bvh = create_bvh_tree_from_object(asset)
+
+                if (
+                    check_bvh_intersection(bvh, room_bvh)
+                    or (
+                        not collision_existing
+                        and existing_obj_bvh is not None
+                        and check_bvh_intersection(bvh, existing_obj_bvh)
+                    )
+                    or (
+                        not collision_placed
+                        and check_bvh_intersection(bvh, placed_obj_bvhs)
+                    )
+                ):
+                    logger.debug(f"Sample {j} of asset {i} rejected, resampling...")
+                    if i == sample_retries - 1:
+                        butil.delete(asset)
+                else:
+                    logger.info(
+                        f"{self.__class__.__name__} placing object {i}/{num_place}, {asset.name=}"
+                    )
+                    placed_obj_bvhs.append(bvh)
+                    break

infinigen_examples/configs_indoor/base_indoors.gin

@@ -60,4 +60,10 @@ compose_indoors.hide_other_rooms_enabled = False
 compose_indoors.fancy_clouds_chance = 0.5
 compose_indoors.grass_chance = 0.5
 compose_indoors.rocks_chance = 0.5
-compose_indoors.near_distance = 20
+compose_indoors.near_distance = 20
+
+compose_indoors.floating_objs_enabled = False
+compose_indoors.num_floating = ('discrete_uniform', 15, 25)
+compose_indoors.norm_floating_size = True
+compose_indoors.enable_collision_floating = False
+compose_indoors.enable_collision_solved = False

infinigen_examples/generate_indoors.py

@@ -20,9 +20,11 @@
 import gin
 import numpy as np
 
+from infinigen import repo_root
 from infinigen.assets import lighting
 from infinigen.assets.materials import invisible_to_camera
 from infinigen.assets.objects.wall_decorations.skirting_board import make_skirting_board
+from infinigen.assets.placement.floating_objects import FloatingObjectPlacement
 from infinigen.assets.utils.decorate import read_co
 from infinigen.core import execute_tasks, init, placement, surface, tagging
 from infinigen.core import tags as t
@@ -52,6 +54,10 @@
     place_cam_overhead,
     restrict_solving,
 )
+from infinigen_examples.util.test_utils import (
+    import_item,
+    load_txt_list,
+)
 
 from . import generate_nature  # noqa F401 # needed for nature gin configs to load
 
@@ -134,6 +140,7 @@ def add_coarse_terrain():
     terrain, terrain_mesh = p.run_stage(
         "terrain", add_coarse_terrain, use_chance=False, default=(None, None)
     )
+
     p.run_stage("sky_lighting", lighting.sky_lighting.add_lighting, use_chance=False)
 
     consgraph = home_constraints()
@@ -307,6 +314,32 @@ def solve_small():
         "populate_assets", populate.populate_state_placeholders, state, use_chance=False
     )
 
+    def place_floating():
+        pholder_rooms = butil.get_collection("placeholders:room_meshes")
+        pholder_cutters = butil.get_collection("placeholders:portal_cutters")
+        pholder_objs = butil.get_collection("placeholders")
+
+        obj_fac_names = load_txt_list(
+            repo_root() / "tests" / "assets" / "list_indoor_meshes.txt"
+        )
+        facs = [import_item(path) for path in obj_fac_names]
+
+        placer = FloatingObjectPlacement(
+            generators=facs,
+            camera=cam_util.get_camera(0, 0),
+            background_objs=list(pholder_cutters.objects) + list(pholder_rooms.objects),
+            collision_objs=list(pholder_objs.objects),
+        )
+
+        placer.place_objs(
+            num_objs=overrides.get("num_floating", 20),
+            normalize=overrides.get("norm_floating_size", True),
+            collision_placed=overrides.get("enable_collision_floating", False),
+            collision_existing=overrides.get("enable_collision_solved", False),
+        )
+
+    p.run_stage("floating_objs", place_floating, use_chance=False, default=state)
+
     door_filter = r.Domain({t.Semantics.Door}, [(cl.AnyRelation(), stages["rooms"])])
     window_filter = r.Domain(
         {t.Semantics.Window}, [(cl.AnyRelation(), stages["rooms"])]