Skip to content

Commit

Permalink
Bugfix v1.2.5 - Terrain bugfix for multi-task command; Terrain.bounds…
Browse files Browse the repository at this point in the history
… and Terrain.populated_bounds parameter

* fix reinitializing terrain

* add bounds and populated_bounds for terrain
  • Loading branch information
mazeyu authored and araistrick committed Apr 3, 2024
1 parent 5132903 commit 18be26c
Show file tree
Hide file tree
Showing 12 changed files with 78 additions and 63 deletions.
2 changes: 1 addition & 1 deletion infinigen/OcMesher
2 changes: 1 addition & 1 deletion infinigen/__init__.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
import logging

__version__ = "1.2.4"
__version__ = "1.2.5"
2 changes: 0 additions & 2 deletions infinigen/terrain/assets/landtiles/custom.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,6 @@ def multi_mountains_asset(
coverage=params["coverage"],
slope_freq=params["slope_freq"],
slope_height=params["slope_height"],
is_asset=True,
)
heightmap = mountains.get_heightmap(X, Y)
mountains.cleanup()
Expand Down Expand Up @@ -176,7 +175,6 @@ def coast_asset(
coverage=params1["coverage"],
slope_freq=params1["slope_freq"],
slope_height=params1["slope_height"],
is_asset=True,
)
heightmap = mountains.get_heightmap(X, Y)
mountains.cleanup()
Expand Down
41 changes: 24 additions & 17 deletions infinigen/terrain/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,8 @@ def get_surface_type(surface, degrade_sdf_to_displacement=True):


class OcMesher(UntexturedOcMesher):
def __init__(self, cameras, **kwargs):
UntexturedOcMesher.__init__(self, get_caminfo(cameras)[0], **kwargs)
def __init__(self, cameras, bounds, **kwargs):
UntexturedOcMesher.__init__(self, get_caminfo(cameras)[0], bounds, **kwargs)

def __call__(self, kernels):
sdf_kernels = [(lambda x, k0=k: k0(x)[Vars.SDF]) for k in kernels]
Expand All @@ -59,8 +59,8 @@ def __call__(self, kernels):
return mesh

class CollectiveOcMesher(UntexturedOcMesher):
def __init__(self, cameras, **kwargs):
UntexturedOcMesher.__init__(self, get_caminfo(cameras)[0], **kwargs)
def __init__(self, cameras, bounds, **kwargs):
UntexturedOcMesher.__init__(self, get_caminfo(cameras)[0], bounds, **kwargs)

def __call__(self, kernels):
sdf_kernels = [lambda x: np.stack([k(x)[Vars.SDF] for k in kernels], -1).min(axis=-1)]
Expand All @@ -74,6 +74,7 @@ def __call__(self, kernels):

@gin.configurable
class Terrain:
instance = None
def __init__(
self,
seed,
Expand All @@ -85,18 +86,23 @@ def __init__(
device="cpu",
main_terrain=TerrainNames.OpaqueTerrain,
under_water=False,
min_distance=1
min_distance=1,
populated_bounds=(-75, 75, -75, 75, -25, 55),
bounds=(-500, 500, -500, 500, -500, 500),
):
self.seed = seed
self.device = device
self.surface_registry = surface_registry
self.main_terrain = main_terrain
self.under_water = under_water
self.min_distance = min_distance
self.populated_bounds = populated_bounds
self.bounds = bounds

if Task.Coarse not in task and Task.FineTerrain not in task:
if Terrain.instance is not None:
self.__dict__ = Terrain.instance.__dict__.copy()
return

with Timer('Create terrain'):
if asset_folder is None:
if not ASSET_ENV_VAR in os.environ:
Expand All @@ -121,6 +127,7 @@ def __init__(
self.elements_list = list(self.elements.values())
logger.info(f"Terrain elements: {[x.__class__.name for x in self.elements_list]}")
transfer_scene_info(self, scene_infos)
Terrain.instance = self

def __del__(self):
self.cleanup()
Expand All @@ -145,10 +152,10 @@ def export(self,
if opaque_elements != []:
attributes_dict[TerrainNames.OpaqueTerrain] = set()
if dynamic:
if spherical: mesher = OpaqueSphericalMesher(cameras=cameras)
else: mesher = OcMesher(cameras=cameras)
if spherical: mesher = OpaqueSphericalMesher(cameras, self.bounds)
else: mesher = OcMesher(cameras, self.bounds)
else:
mesher = UniformMesher()
mesher = UniformMesher(self.populated_bounds)
with Timer(f"meshing {TerrainNames.OpaqueTerrain}"):
mesh = mesher([element for element in opaque_elements])
meshes_dict[TerrainNames.OpaqueTerrain] = mesh
Expand All @@ -163,9 +170,9 @@ def export(self,
if element.__class__.name == ElementNames.Atmosphere:
special_args["pixels_per_cube"] = 100
special_args["inv_scale"] = 1
if spherical: mesher = TransparentSphericalMesher(cameras=cameras, **special_args)
else: mesher = OcMesher(cameras=cameras, simplify_occluded=False, **special_args)
else: mesher = UniformMesher(enclosed=True)
if spherical: mesher = TransparentSphericalMesher(cameras, self.bounds, **special_args)
else: mesher = OcMesher(cameras, self.bounds, simplify_occluded=False, **special_args)
else: mesher = UniformMesher(self.populated_bounds, enclosed=True)
with Timer(f"meshing {element.__class__.name}"):
mesh = mesher([element])
meshes_dict[element.__class__.name] = mesh
Expand All @@ -176,10 +183,10 @@ def export(self,
if collective_transparent_elements != []:
attributes_dict[TerrainNames.CollectiveTransparentTerrain] = set()
if dynamic:
if spherical: mesher = TransparentSphericalMesher(cameras=cameras)
else: mesher = CollectiveOcMesher(cameras=cameras, simplify_occluded=False)
if spherical: mesher = TransparentSphericalMesher(cameras, self.bounds)
else: mesher = CollectiveOcMesher(cameras, self.bounds, simplify_occluded=False)
else:
mesher = UniformMesher()
mesher = UniformMesher(self.populated_bounds)
with Timer(f"meshing {TerrainNames.CollectiveTransparentTerrain}"):
mesh = mesher([element for element in collective_transparent_elements])
meshes_dict[TerrainNames.CollectiveTransparentTerrain] = mesh
Expand Down Expand Up @@ -211,7 +218,7 @@ def export(self,
if len(attributes_dict[mesh_name]) == 1:
meshes_dict[mesh_name].vertex_attributes.pop(list(attributes_dict[mesh_name])[0])
else:
self.bounding_box = np.array(mesher.dimensions)[::2], np.array(mesher.dimensions)[1::2]
self.bounding_box = np.array(self.populated_bounds)[::2], np.array(self.populated_bounds)[1::2]

return meshes_dict, attributes_dict

Expand Down
12 changes: 1 addition & 11 deletions infinigen/terrain/elements/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,7 @@
@gin.configurable
class Element:
called_time = {}
def __init__(self, lib_name, material, transparency, bounding_shape=None, is_asset=False):
self.bounding_shape = bounding_shape
self.is_asset = is_asset
def __init__(self, lib_name, material, transparency):
if lib_name in Element.called_time:
lib_name_X = f"{lib_name}_{Element.called_time[lib_name]}"
print(f"{lib_name} already loaded, loading {lib_name_X} instead")
Expand Down Expand Up @@ -79,14 +77,6 @@ def __call__(self, positions, sdf_only=False):
else:
auxs.append(None)
self.call(N, ASFLOAT(AC(positions.astype(np.float32))), ASFLOAT(sdf), *[POINTER(c_float)() if x is None else ASFLOAT(x) for x in auxs])
bounding_shape = self.bounding_shape
if not self.is_asset and bounding_shape is not None:
if bounding_shape[0] == "cube":
x_min, x_max, y_min, y_max, z_min, z_max = bounding_shape[1:]
out_bound = (positions[:, 0] < x_min) | (positions[:, 0] > x_max) \
| (positions[:, 1] < y_min) | (positions[:, 1] > y_max) \
| (positions[:, 2] < z_min) | (positions[:, 2] > z_max)
sdf[out_bound] = 1e6
ret = {}
ret[Vars.SDF] = sdf

Expand Down
3 changes: 1 addition & 2 deletions infinigen/terrain/elements/mountains.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ def __init__(
slope_octaves=9,
material=Materials.MountainCollection,
transparency=Transparency.Opaque,
is_asset=False,
):
self.device = device
min_freq = rg(min_freq)
Expand All @@ -58,5 +57,5 @@ def __init__(
slope_freq, slope_octaves, slope_height,
], dtype=np.float32))

Element.__init__(self, "mountains", material, transparency, is_asset=is_asset)
Element.__init__(self, "mountains", material, transparency)
self.tag = ElementTag.Terrain
40 changes: 26 additions & 14 deletions infinigen/terrain/mesher/spherical_mesher.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,17 @@
magnifier = 1e6

@gin.configurable
def kernel_caller(kernels, XYZ, r_bound_max=None, cam_loc=None):
def kernel_caller(kernels, XYZ, bounds=None):
sdfs = []
for kernel in kernels:
ret = kernel(XYZ, sdf_only=1)
sdf = ret[Vars.SDF]
if r_bound_max is not None:
R = np.linalg.norm(XYZ - cam_loc, axis=-1)
mask = R / r_bound_max - 1
sdf = np.maximum(sdf, mask * magnifier)
if bounds is not None:
out_bound = np.zeros(len(XYZ), dtype=bool)
for i in range(3):
out_bound |= XYZ[:, i] <= bounds[i*2]
out_bound |= XYZ[:, i] >= bounds[i*2+1]
sdf[out_bound] = 1e6 # because of skimage mc only provides coords, which is has precision limit
sdfs.append(sdf)
ret = np.stack(sdfs, -1)
return ret
Expand All @@ -32,28 +34,38 @@ def kernel_caller(kernels, XYZ, r_bound_max=None, cam_loc=None):
class SphericalMesher:
def __init__(self,
cameras,
bounds,
r_min=1,
r_max=1e3,
complete_depth_test=True,
):
_, self.cam_pose, self.fov, self.H, self.W, _ = get_caminfo(cameras)
full_info, self.cam_pose, self.fov, self.H, self.W, _ = get_caminfo(cameras)
cams = full_info[0]
assert self.fov[0] < np.pi / 2 and self.fov[1] < np.pi / 2, "the algorithm does not support larger-than-90-degree fov yet"
self.r_min = r_min
self.r_max = r_max
self.complete_depth_test = complete_depth_test
self.bounds = bounds
self.r_max = 0
for cam in cams:
for i in range(2):
for j in range(2):
for k in range(2):
r_max = np.linalg.norm(np.array([self.bounds[i], self.bounds[2+j], self.bounds[4+k]]) - cam[:3, 3])
self.r_max = max(self.r_max, r_max)
self.r_max *= 1.1

@gin.configurable
class OpaqueSphericalMesher(SphericalMesher):
def __init__(self,
cameras,
bounds,
base_90d_resolution=None,
pixels_per_cube=1.84,
test_downscale=5,
upscale1=2,
upscale2=4,
r_lengthen=1,
):
SphericalMesher.__init__(self, cameras)
SphericalMesher.__init__(self, cameras, bounds)
inview_upscale_coarse = upscale1
inview_upscale_fine = upscale1 * upscale2
outview_upscale = 1
Expand Down Expand Up @@ -84,7 +96,7 @@ def __init__(self,
test_downscale=test_downscale,
complete_depth_test=self.complete_depth_test,
)
self.frontview_mesher.kernel_caller = kernel_caller
self.frontview_mesher.kernel_caller = lambda k, xyz: kernel_caller(k, xyz, self.bounds)
self.background_mesher = CubeSphericalMesher(
self.cam_pose,
self.r_min, self.r_max,
Expand All @@ -94,7 +106,7 @@ def __init__(self,
H_fov=rounded_fov[0], W_fov=rounded_fov[1],
N0=N0, N1=N1,
)
self.background_mesher.kernel_caller = kernel_caller
self.background_mesher.kernel_caller = lambda k, xyz: kernel_caller(k, xyz, self.bounds)

def __call__(self, kernels):
with Timer("OpaqueSphericalMesher: frontview_mesher"):
Expand All @@ -110,14 +122,15 @@ def __call__(self, kernels):
class TransparentSphericalMesher(SphericalMesher):
def __init__(self,
cameras,
bounds,
base_90d_resolution=None,
pixels_per_cube=1.84,
test_downscale=5,
inv_scale=8,
r_lengthen=3,
camera_annotation_frames=None,
):
SphericalMesher.__init__(self, cameras)
SphericalMesher.__init__(self, cameras, bounds)
self.cameras = cameras
self.camera_annotation_frames = camera_annotation_frames
assert bool(base_90d_resolution is None) ^ bool(pixels_per_cube is None)
Expand All @@ -144,8 +157,7 @@ def __init__(self,
N0=N0, N1=N1,
complete_depth_test=self.complete_depth_test,
)
r_bound_max = np.exp(np.log(self.r_min) + (np.log(self.r_max) - np.log(self.r_min)) * (base_R - 0.5) / base_R)
self.mesher.kernel_caller = lambda k, xyz: kernel_caller(k, xyz, r_bound_max, self.cam_pose[:3, 3])
self.mesher.kernel_caller = lambda k, xyz: kernel_caller(k, xyz, self.bounds)


def __call__(self, kernels):
Expand Down
22 changes: 11 additions & 11 deletions infinigen/terrain/mesher/uniform_mesher.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ def __init__(self, desc, verbose=False):
@gin.configurable
class UniformMesher:
def __init__(self,
dimensions=(-75, 75, -75, 75, -25, 55),
bounds,
subdivisions=(64, -1, -1), # -1 means automatic
upscale=3,
enclosed=False,
Expand All @@ -40,29 +40,29 @@ def __init__(self,
):
self.enclosed = enclosed
self.upscale = upscale
self.dimensions = dimensions
# Lx, Ly, Lz = dimensions[1] - dimensions[0], dimensions[3] - dimensions[2], dimensions[5] - dimensions[4]
self.bounds = bounds

assert(np.sum(subdivisions == -1) in [0, 2])
for i, s in enumerate(subdivisions):
if s != -1:
coarse_voxel_size = (dimensions[i * 2 + 1] - dimensions[i * 2]) / s
coarse_voxel_size = (bounds[i * 2 + 1] - bounds[i * 2]) / s

if subdivisions[0] != -1:
self.x_N = subdivisions[0]
else:
self.x_N = int((dimensions[1] - dimensions[0]) / coarse_voxel_size)
self.x_N = int((bounds[1] - bounds[0]) / coarse_voxel_size)
if subdivisions[1] != -1:
self.y_N = subdivisions[1]
else:
self.y_N = int((dimensions[3] - dimensions[2]) / coarse_voxel_size)
self.y_N = int((bounds[3] - bounds[2]) / coarse_voxel_size)
if subdivisions[2] != -1:
self.z_N = subdivisions[2]
else:
self.z_N = int((dimensions[5] - dimensions[4]) / coarse_voxel_size)
self.z_N = int((bounds[5] - bounds[4]) / coarse_voxel_size)

self.x_min, self.x_max = dimensions[0], dimensions[1]
self.y_min, self.y_max = dimensions[2], dimensions[3]
self.z_min, self.z_max = dimensions[4], dimensions[5]
self.x_min, self.x_max = bounds[0], bounds[1]
self.y_min, self.y_max = bounds[2], bounds[3]
self.z_min, self.z_max = bounds[4], bounds[5]
self.closing_margin = coarse_voxel_size / upscale / 2
self.verbose = verbose
self.bisection_iters = bisection_iters
Expand Down Expand Up @@ -94,7 +94,7 @@ def kernel_caller(self, kernels, XYZ):
out_bound = (XYZ[:, 0] < self.x_min + self.closing_margin) | (XYZ[:, 0] > self.x_max - self.closing_margin) \
| (XYZ[:, 1] < self.y_min + self.closing_margin) | (XYZ[:, 1] > self.y_max - self.closing_margin) \
| (XYZ[:, 2] < self.z_min + self.closing_margin) | (XYZ[:, 2] > self.z_max - self.closing_margin)
sdf[out_bound] = 1e11
sdf[out_bound] = 1e6
sdfs.append(sdf)
return np.stack(sdfs, -1)

Expand Down
11 changes: 10 additions & 1 deletion infinigen/tools/datarelease_toolkit.py
Original file line number Diff line number Diff line change
Expand Up @@ -324,8 +324,17 @@ def reorganize_old_framesfolder(frames_old):

frames_dest = frames_old.parent/"frames"

excludes = [
"version.txt",
"polycounts.txt",
"version.txt",
"MaskTag.json",
"scene.blend",
"pipeline_coarse.csv",
]

for img_path in frames_old.iterdir():
if img_path.is_dir():
if img_path.is_dir() or img_path.name in excludes:
continue
dtype, *_ = img_path.name.split('_')
idxs = parse_suffix(img_path.name)
Expand Down
2 changes: 1 addition & 1 deletion infinigen_examples/configs/scene_types/cave.gin
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ Waterbody.height = -5
full/render_image.min_samples = 500

# camera selection config
UniformMesher.dimensions = (-25, 25, -25, 25, -25, 0)
Terrain.populated_bounds = (-25, 25, -25, 25, -25, 0)
keep_cam_pose_proposal.terrain_coverage_range = (1, 1)
camera_selection_preprocessing.terrain_tags_ratio = {"cave": (0.3, 1), ("closeup", 4): (0, 0.3)}

Expand Down
2 changes: 1 addition & 1 deletion infinigen_examples/configs/scene_types/cliff.gin
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ WarpedRocks.slope_shift=-15
Ground.height = -15

# camera selection config
UniformMesher.dimensions = (-25, 25, -25, 25, -15, 35)
Terrain.populated_bounds = (-25, 25, -25, 25, -15, 35)
keep_cam_pose_proposal.terrain_coverage_range = (0, 0.8)
camera_selection_preprocessing.terrain_tags_ratio = {("altitude", 10, 1e9): (0.01, 1)}

Expand Down
2 changes: 1 addition & 1 deletion infinigen_examples/configs/scene_types/river.gin
Original file line number Diff line number Diff line change
Expand Up @@ -37,5 +37,5 @@ scene.ground_chance = 0
scene.warped_rocks_chance = 0

# camera selection config
UniformMesher.dimensions = (-25, 25, -25, 25, -15, 35)
Terrain.populated_bounds = (-25, 25, -25, 25, -15, 35)
camera_selection_preprocessing.terrain_tags_ratio = {("altitude", -1e9, 0.75): (0.1, 1), "liquid": (0.05, 1)}

0 comments on commit 18be26c

Please sign in to comment.