269 lines
10 KiB
Python
269 lines
10 KiB
Python
import bpy
|
||
from pathlib import Path
|
||
|
||
|
||
def find_snow_modifier(terrain_obj):
|
||
"""
|
||
Find the Geometry Nodes modifier that contains snow_depth attribute.
|
||
Returns the modifier or None if not found.
|
||
"""
|
||
for mod in terrain_obj.modifiers:
|
||
if mod.type == 'NODES' and mod.node_group:
|
||
# Check if this modifier's node tree has Store Named Attribute with "snow_depth"
|
||
for node in mod.node_group.nodes:
|
||
if node.type == 'STORE_NAMED_ATTRIBUTE':
|
||
if hasattr(node, 'data_type') and node.name and 'snow' in node.name.lower():
|
||
return mod
|
||
# Check inputs for the name
|
||
for input in node.inputs:
|
||
if input.name == 'Name' and hasattr(input, 'default_value'):
|
||
if input.default_value == 'snow_depth':
|
||
return mod
|
||
# Fallback: check modifier name
|
||
if 'snow' in mod.name.lower():
|
||
return mod
|
||
return None
|
||
|
||
|
||
def bake_snow_depth(terrain_obj, resolution=512, output_path=None, modifier_name=None):
|
||
"""
|
||
Bake snow depth attribute to texture using shader-based Cycles baking.
|
||
Uses the same approach as generate_heightmap.py.
|
||
|
||
Requires:
|
||
- Terrain object with Geometry Nodes modifier that stores 'snow_depth' attribute
|
||
- UV map on terrain mesh
|
||
|
||
Args:
|
||
terrain_obj: Terrain mesh object with snow_depth attribute
|
||
resolution: Texture resolution (square)
|
||
output_path: Path to save EXR file
|
||
modifier_name: Optional specific modifier name to use (e.g., "Snow")
|
||
"""
|
||
|
||
print(f"Baking snow depth for: {terrain_obj.name}")
|
||
print(f"Resolution: {resolution}×{resolution}")
|
||
|
||
# Find the snow geometry nodes modifier
|
||
if modifier_name:
|
||
geo_nodes_modifier = terrain_obj.modifiers.get(modifier_name)
|
||
if not geo_nodes_modifier:
|
||
raise ValueError(f"Modifier '{modifier_name}' not found on {terrain_obj.name}")
|
||
print(f"Using specified modifier: {modifier_name}")
|
||
else:
|
||
geo_nodes_modifier = find_snow_modifier(terrain_obj)
|
||
if not geo_nodes_modifier:
|
||
print("\nAvailable Geometry Nodes modifiers:")
|
||
for mod in terrain_obj.modifiers:
|
||
if mod.type == 'NODES':
|
||
print(f" - {mod.name}")
|
||
raise ValueError(
|
||
f"No Geometry Nodes modifier with 'snow_depth' attribute found on {terrain_obj.name}!\n"
|
||
f"Either add snow accumulation modifier, or specify modifier_name parameter."
|
||
)
|
||
print(f"Found snow modifier: {geo_nodes_modifier.name}")
|
||
|
||
modifier_states = {}
|
||
print(f"\nDisabling modifiers after '{geo_nodes_modifier.name}' for baking...")
|
||
target_mod_index = list(terrain_obj.modifiers).index(geo_nodes_modifier)
|
||
|
||
for i, mod in enumerate(terrain_obj.modifiers):
|
||
modifier_states[mod.name] = {
|
||
'show_viewport': mod.show_viewport,
|
||
'show_render': mod.show_render
|
||
}
|
||
|
||
if i > target_mod_index:
|
||
print(f" Temporarily disabling: {mod.name}")
|
||
mod.show_viewport = False
|
||
mod.show_render = False
|
||
|
||
bpy.context.view_layer.update()
|
||
depsgraph = bpy.context.evaluated_depsgraph_get()
|
||
evaluated_obj = terrain_obj.evaluated_get(depsgraph)
|
||
eval_mesh = evaluated_obj.to_mesh()
|
||
|
||
if 'snow_depth' not in eval_mesh.attributes:
|
||
evaluated_obj.to_mesh_clear()
|
||
for mod_name, state in modifier_states.items():
|
||
mod = terrain_obj.modifiers.get(mod_name)
|
||
if mod:
|
||
mod.show_viewport = state['show_viewport']
|
||
mod.show_render = state['show_render']
|
||
raise ValueError("snow_depth attribute missing from evaluated geometry")
|
||
|
||
print(f"✓ Verified 'snow_depth' attribute exists")
|
||
evaluated_obj.to_mesh_clear()
|
||
|
||
# Ensure object has UV map
|
||
if not terrain_obj.data.uv_layers:
|
||
print("Adding UV map...")
|
||
terrain_obj.data.uv_layers.new(name="UVMap")
|
||
|
||
# Create new image for baking
|
||
bake_image = bpy.data.images.new(
|
||
name="SnowDepth_Bake",
|
||
width=resolution,
|
||
height=resolution,
|
||
alpha=False,
|
||
float_buffer=True,
|
||
is_data=True
|
||
)
|
||
|
||
print(f"Created bake image: {bake_image.name}")
|
||
|
||
original_materials = list(terrain_obj.data.materials)
|
||
print(f"Object has {len(original_materials)} material slot(s): {[mat.name if mat else 'None' for mat in original_materials]}")
|
||
|
||
mat = bpy.data.materials.new(name="SnowDepth_BakeMaterial")
|
||
mat.use_nodes = True
|
||
nodes = mat.node_tree.nodes
|
||
nodes.clear()
|
||
|
||
attr_node = nodes.new(type='ShaderNodeAttribute')
|
||
attr_node.attribute_name = 'snow_depth'
|
||
|
||
emission_node = nodes.new(type='ShaderNodeEmission')
|
||
mat.node_tree.links.new(attr_node.outputs['Fac'], emission_node.inputs['Color'])
|
||
|
||
output_node = nodes.new(type='ShaderNodeOutputMaterial')
|
||
mat.node_tree.links.new(emission_node.outputs['Emission'], output_node.inputs['Surface'])
|
||
|
||
image_node = nodes.new(type='ShaderNodeTexImage')
|
||
image_node.image = bake_image
|
||
image_node.select = True
|
||
nodes.active = image_node
|
||
|
||
terrain_obj.data.materials.clear()
|
||
terrain_obj.data.materials.append(mat)
|
||
print(f"Temporarily replaced all materials with bake material")
|
||
|
||
# Select object and set mode
|
||
bpy.context.view_layer.objects.active = terrain_obj
|
||
terrain_obj.select_set(True)
|
||
|
||
# Ensure we're in object mode
|
||
if bpy.context.object and bpy.context.object.mode != 'OBJECT':
|
||
bpy.ops.object.mode_set(mode='OBJECT')
|
||
|
||
# Setup render settings for baking
|
||
bpy.context.scene.render.engine = 'CYCLES'
|
||
bpy.context.scene.cycles.samples = 1
|
||
bpy.context.scene.cycles.bake_type = 'EMIT'
|
||
|
||
print("Baking with Cycles (EMIT)...")
|
||
bpy.ops.object.bake(type='EMIT', use_clear=True)
|
||
|
||
print("Bake complete!")
|
||
|
||
# Verify bake has data (not all black/zero)
|
||
pixels = list(bake_image.pixels)
|
||
max_value = max(pixels) if pixels else 0.0
|
||
avg_value = sum(pixels) / len(pixels) if pixels else 0.0
|
||
non_zero_count = sum(1 for p in pixels if p > 0.0001)
|
||
print(f"Baked image stats: max={max_value:.4f}, avg={avg_value:.4f}")
|
||
print(f"Non-zero pixels: {non_zero_count} ({non_zero_count / len(pixels) * 100:.1f}%)")
|
||
|
||
if max_value < 0.0001:
|
||
print("\n⚠️ WARNING: Baked image appears to be all black!")
|
||
print(" Possible causes:")
|
||
print(" - 'snow_depth' attribute doesn't exist in the geometry")
|
||
print(" - Geometry Nodes modifier is disabled")
|
||
print(" - Store Named Attribute node is not connected")
|
||
print(" - Wrong modifier selected (try specifying modifier_name)")
|
||
print("\n Continuing anyway, but check your setup...")
|
||
else:
|
||
print(f"✓ Bake contains data (values up to {max_value:.4f}m)")
|
||
|
||
# Save as EXR
|
||
if output_path:
|
||
bake_image.filepath_raw = str(output_path)
|
||
bake_image.file_format = 'OPEN_EXR'
|
||
bake_image.use_half_precision = False
|
||
|
||
scene = bpy.context.scene
|
||
original_color_mode = scene.render.image_settings.color_mode
|
||
original_color_depth = scene.render.image_settings.color_depth
|
||
original_exr_codec = scene.render.image_settings.exr_codec
|
||
|
||
# Use BW mode for single channel (same as heightmap)
|
||
scene.render.image_settings.color_mode = 'BW'
|
||
scene.render.image_settings.color_depth = '32'
|
||
scene.render.image_settings.exr_codec = 'ZIP'
|
||
|
||
print(f"Saving EXR with settings: color_mode=BW, depth=32, codec=ZIP")
|
||
bake_image.save_render(str(output_path), scene=scene)
|
||
|
||
scene.render.image_settings.color_mode = original_color_mode
|
||
scene.render.image_settings.color_depth = original_color_depth
|
||
scene.render.image_settings.exr_codec = original_exr_codec
|
||
|
||
print(f"Saved to: {output_path}")
|
||
print(f"Format: OpenEXR, 32-bit float, ZIP compression")
|
||
print(f"File size: {output_path.stat().st_size / 1024:.1f} KB")
|
||
|
||
bpy.data.images.remove(bake_image)
|
||
bpy.data.materials.remove(mat)
|
||
|
||
terrain_obj.data.materials.clear()
|
||
for original_mat in original_materials:
|
||
terrain_obj.data.materials.append(original_mat)
|
||
print(f"Restored {len(original_materials)} original material(s)")
|
||
|
||
print("\nRestoring modifier states...")
|
||
for mod_name, state in modifier_states.items():
|
||
mod = terrain_obj.modifiers.get(mod_name)
|
||
if mod:
|
||
mod.show_viewport = state['show_viewport']
|
||
mod.show_render = state['show_render']
|
||
print("✓ Modifiers restored")
|
||
|
||
return True
|
||
|
||
|
||
if __name__ == "__main__":
|
||
project_root = Path(bpy.data.filepath).parent.parent
|
||
output_path = project_root / "textures" / "snow_depth.exr"
|
||
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Find terrain object
|
||
terrain_obj = bpy.data.objects.get("TerrainPlane")
|
||
|
||
if not terrain_obj:
|
||
print("'TerrainPlane' not found. Searching for terrain mesh...")
|
||
for obj in bpy.data.objects:
|
||
if obj.type == 'MESH' and ('terrain' in obj.name.lower() or 'plane' in obj.name.lower()):
|
||
terrain_obj = obj
|
||
print(f"Using: {obj.name}")
|
||
break
|
||
|
||
if not terrain_obj:
|
||
raise ValueError("No terrain object found!")
|
||
|
||
# CONFIGURATION: Specify modifier name if you have multiple Geometry Nodes modifiers
|
||
# Leave as None to auto-detect the snow modifier
|
||
# Example: modifier_name = "Snow" or "Snow Accumulation"
|
||
modifier_name = "Snow Accumulation" # Auto-detect by looking for 'snow_depth' attribute
|
||
|
||
bake_snow_depth(
|
||
terrain_obj=terrain_obj,
|
||
resolution=1000,
|
||
output_path=output_path,
|
||
modifier_name=modifier_name # Specify "Snow" if auto-detect fails
|
||
)
|
||
|
||
print("\n" + "="*60)
|
||
print("Snow depth baking complete!")
|
||
print(f"Output: {output_path}")
|
||
print("="*60)
|
||
print("\nNext steps:")
|
||
print("1. Verify snow_depth.exr in textures/ directory")
|
||
print("2. Open in image viewer to check it's not black")
|
||
print("3. Load in game with SnowLayer::load()")
|
||
print("4. Test deformation with player movement")
|
||
print("\nIf bake is black:")
|
||
print("- Check that 'snow_depth' attribute exists in Spreadsheet Editor")
|
||
print("- Verify Geometry Nodes modifier has Store Named Attribute node")
|
||
print("- Try specifying modifier_name='Snow' explicitly in script")
|