diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 15cad92d3..83a1cee07 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -78,4 +78,5 @@ v1.5.1 v1.6.0 - Add geometric tile pattern materials -- Tune window parameters and materials \ No newline at end of file +- Tune window parameters and materials +- Add floating object placement generator and example command \ No newline at end of file diff --git a/docs/HelloRoom.md b/docs/HelloRoom.md index aaeea6cfa..fb4cadffd 100644 --- a/docs/HelloRoom.md +++ b/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 diff --git a/infinigen/OcMesher b/infinigen/OcMesher index d3d1441ab..2cdcbacbe 160000 --- a/infinigen/OcMesher +++ b/infinigen/OcMesher @@ -1 +1 @@ -Subproject commit d3d1441ab57c48db3ec40c621fc3d0c323579e8a +Subproject commit 2cdcbacbe62ef79dc6031e0131f916266b7372e3 diff --git a/infinigen/assets/placement/floating_objects.py b/infinigen/assets/placement/floating_objects.py new file mode 100644 index 000000000..21e47f1ed --- /dev/null +++ b/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 diff --git a/infinigen_examples/configs_indoor/base_indoors.gin b/infinigen_examples/configs_indoor/base_indoors.gin index 4b69c72ec..f0f07059a 100644 --- a/infinigen_examples/configs_indoor/base_indoors.gin +++ b/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 \ No newline at end of file +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 \ No newline at end of file diff --git a/infinigen_examples/generate_indoors.py b/infinigen_examples/generate_indoors.py index a557f4232..cc9ef165d 100644 --- a/infinigen_examples/generate_indoors.py +++ b/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"])]