heightmap_worker.py

import numpy as np
from PIL import Image, ImageDraw, ImageFilter
import multiprocessing
from canyon_worker import create_canyon_path_worker
from noise import snoise2

def create_heightmap(height_data, params):
    """
    Create RGB heightmap from raw height data with proper canyon application.

    Args:
        height_data (np.ndarray): Raw height data (2D float array, 0-1 range)
        params (dict): Dictionary containing all necessary parameters

    Returns:
        np.ndarray: RGB heightmap array (height, width, 3) with uint8 values
    """
    # Create RGB image with height in red channel
    heightmap = np.zeros((height_data.shape[0], height_data.shape[1], 3), dtype=np.uint8)

    # Get min and max height values
    min_height = params.get("min_height", 1.0)
    max_height = params.get("max_height", 0.5)

    # Adjust height data based on min/max settings
    if min_height != 0.0:
        height_data = height_data + min_height

    if max_height != 1.0:
        current_max = height_data.max()
        if current_max > 0:
            height_data = height_data * (max_height / current_max)

    # Ensure values are within valid range
    height_data = np.clip(height_data, 0.0, 1.0)

    # Get necessary parameters
    size = height_data.shape[0]  # Assuming square heightmap
    seed = params.get("seed", 42)
    canyon_strength = params.get("canyon_strength", 0.0)
    canyon_length = params.get("canyon_length", 0.7)
    canyon_branch_density = params.get("canyon_branch_density", 0.15)
    canyon_count = params.get("canyon_count", 6)
    canyon_seed = params.get("canyon_seed", 42)
    vignette_strength = params.get("vignette_strength", 0.0)
    vignette_radius = params.get("vignette_radius", 0.5)

    # Apply canyon effect
    if canyon_strength > 0:
        height_data = _apply_canyon_effect(height_data, size, seed, canyon_strength,
                                         canyon_length, canyon_branch_density,
                                         canyon_count, canyon_seed)

    # Scale to 0-255 range AFTER canyon application
    height_scaled = (height_data * 255).astype(np.uint8)

    # Apply vignette if strength > 0
    if vignette_strength > 0:
        height_scaled = _apply_vignette(height_scaled, vignette_strength, vignette_radius)

    # Set channels according to Teardown format
    heightmap[:, :, 0] = height_scaled  # Red channel = height

    # Generate grass channel
    if params.get("use_noise_grass", False):
        # Generate grass noise
        grass_noise = _generate_grass_noise(height_data.shape[0], height_data.shape[1], params)
        heightmap[:, :, 1] = grass_noise  # Green channel = grass with noise
    else:
        # Use uniform grass amount
        grass_amount = params.get("grass_amount", 0)
        heightmap[:, :, 1] = grass_amount  # Green channel = grass amount

    # Fill with special value
    special_val = params.get("special_value", 0)
    heightmap[:, :, 2] = special_val

    return heightmap

def _apply_canyon_effect(height_data, size, seed, canyon_strength, canyon_length,
                        canyon_branch_density, canyon_count, canyon_seed):
    """Apply canyon effect to height data."""
    # Create canyon mask image
    canyon_img = Image.new('L', (size, size), 0)
    draw = ImageDraw.Draw(canyon_img)

    # Fixed random generator with canyon-specific seed
    fixed_rng = np.random.RandomState(canyon_seed)

    # Center point for canyons
    cx, cy = size / 2, size / 2

    # Number of canyons per edge
    canyons_per_edge = canyon_count

    # Generate edge points
    edge_points = []

    # Generate evenly distributed edge points with consistent randomness
    def generate_edge_points(edge_type, count):
        points = []
        spacing = size / count

        if edge_type == "top":
            y = 0
            for i in range(count):
                offset = fixed_rng.uniform(0.2, 0.8)
                x = int((i + offset) * spacing)
                points.append((x, y))
        elif edge_type == "bottom":
            y = size - 1
            for i in range(count):
                offset = fixed_rng.uniform(0.2, 0.8)
                x = int((i + offset) * spacing)
                points.append((x, y))
        elif edge_type == "left":
            x = 0
            for i in range(count):
                offset = fixed_rng.uniform(0.2, 0.8)
                y = int((i + offset) * spacing)
                points.append((x, y))
        elif edge_type == "right":
            x = size - 1
            for i in range(count):
                offset = fixed_rng.uniform(0.2, 0.8)
                y = int((i + offset) * spacing)
                points.append((x, y))

        return points

    # Get points from all four edges
    edge_points.extend(generate_edge_points("top", canyons_per_edge))
    edge_points.extend(generate_edge_points("bottom", canyons_per_edge))
    edge_points.extend(generate_edge_points("left", canyons_per_edge))
    edge_points.extend(generate_edge_points("right", canyons_per_edge))

    # Create arguments for parallel processing
    worker_args = [(start_x, start_y, cx, cy, canyon_length, canyon_branch_density, canyon_seed, size)
                  for start_x, start_y in edge_points]

    # Use multiprocessing to generate canyon paths
    try:
        ctx = multiprocessing.get_context('spawn')
        num_workers = min(ctx.cpu_count(), 4)  # Use at most 4 workers

        with ctx.Pool(processes=num_workers) as pool:
            canyon_paths_results = pool.map(create_canyon_path_worker, worker_args)

        # Filter out None results
        all_canyon_paths = [path for path in canyon_paths_results if path is not None]
    except Exception as e:
        # Fallback to single-process if multiprocessing fails
        print(f"Multiprocessing error for canyons: {e}. Using single-process mode.")
        all_canyon_paths = []

        # Process each starting point sequentially as fallback
        for start_x, start_y in edge_points:
            result = create_canyon_path_worker((start_x, start_y, cx, cy, canyon_length,
                                               canyon_branch_density, canyon_seed, size))
            if result is not None:
                all_canyon_paths.append(result)

    # Apply the actual canyon_length parameter to draw only portions of each path
    for canyon in all_canyon_paths:
        # Get the main path
        path_points = canyon['main_path']

        # Calculate the actual path length based on canyon_length parameter
        length_variation = fixed_rng.uniform(-0.05, 0.05)
        length_ratio = canyon_length + length_variation
        length_ratio = max(0.2, min(0.95, length_ratio))

        # Get the number of points to use based on length_ratio
        points_to_use = max(2, int(len(path_points) * length_ratio))
        used_path = path_points[:points_to_use]

        # Smooth the main path
        if len(used_path) > 3:
            smoothed_path = []
            window_size = 3

            for i in range(len(used_path)):
                if i < window_size // 2 or i >= len(used_path) - window_size // 2:
                    smoothed_path.append(used_path[i])
                else:
                    x_avg = sum(p[0] for p in used_path[i-window_size//2:i+window_size//2+1]) / window_size
                    y_avg = sum(p[1] for p in used_path[i-window_size//2:i+window_size//2+1]) / window_size
                    smoothed_path.append((int(x_avg), int(y_avg)))

            # Draw the main path with width scaled by canyon_strength
            for j in range(len(smoothed_path) - 1):
                progress = j / (len(smoothed_path) - 1)
                width = max(2, int((1.0 - progress * 0.7) * canyon_strength * 15))
                intensity = int(255 * (1.0 - progress * 0.2))
                draw.line([smoothed_path[j], smoothed_path[j+1]],
                         fill=intensity,
                         width=width)

        # Draw branches - only those that connect to the used portion of the main path
        for branch_points in canyon['branches']:
            # Get the starting point of the branch
            branch_start = branch_points[0]

            # Check if this branch connects to the used portion of the path
            is_connected = False
            for p in used_path:
                if abs(p[0] - branch_start[0]) <= 1 and abs(p[1] - branch_start[1]) <= 1:
                    is_connected = True
                    break

            # Only draw branches that connect to the used path
            if is_connected:
                # Calculate how much of the branch to use based on main path length ratio
                branch_length_ratio = length_ratio * 1.2  # Branches can be a bit longer
                branch_length_ratio = min(1.0, branch_length_ratio)  # Cap at 100%

                # Get the points to use
                branch_points_to_use = max(2, int(len(branch_points) * branch_length_ratio))
                used_branch = branch_points[:branch_points_to_use]

                # Smooth the branch
                if len(used_branch) > 3:
                    smoothed_branch = []
                    window_size = 3

                    for i in range(len(used_branch)):
                        if i < window_size // 2 or i >= len(used_branch) - window_size // 2:
                            smoothed_branch.append(used_branch[i])
                        else:
                            x_avg = sum(p[0] for p in used_branch[i-window_size//2:i+window_size//2+1]) / window_size
                            y_avg = sum(p[1] for p in used_branch[i-window_size//2:i+window_size//2+1]) / window_size
                            smoothed_branch.append((int(x_avg), int(y_avg)))

                    # Draw the branch
                    for k in range(len(smoothed_branch) - 1):
                        prog = k / (len(smoothed_branch) - 1)
                        branch_width = max(1, int((1.0 - prog * 0.7) * canyon_strength * 5))
                        intensity = int(220 * (1.0 - prog * 0.3))
                        draw.line([smoothed_branch[k], smoothed_branch[k+1]],
                                 fill=intensity,
                                 width=branch_width)

    # Apply Gaussian blur scaled with canyon_strength
    blur_radius = max(1.0, min(3.0, size / 256 * canyon_strength * 2))
    canyon_img = canyon_img.filter(ImageFilter.GaussianBlur(radius=blur_radius))

    # Convert to numpy array
    canyon_mask = np.array(canyon_img) / 255.0

    # Apply to heightmap with canyon_strength
    height_data = height_data * (1.0 - canyon_mask * min(1.0, canyon_strength * 1.2))

    return height_data

def _generate_grass_noise(height, width, params):
    """Generate grass noise using exactly the same coordinate system as the main noise."""
    # Grass-specific parameters
    octaves = params.get("grass_noise_octaves", 4)
    persistence = params.get("grass_noise_persistence", 0.5)
    grass_scale = params.get("grass_noise_scale", 50.0)
    density = params.get("noise_grass_density", 0.5)
    grass_amount = params.get("grass_amount", 0)
    seed = params.get("seed", 42) + 1000  # Offset seed to differentiate grass

    # Scaling and resolution matching
    base_scale = params.get("scale", 150.0)
    preview_res = params.get("noise_size", 512)  # Size used in preview
    export_res = width  # Export width — assume square

    zoom_factor = export_res / preview_res
    adjusted_scale = base_scale * zoom_factor * 0.05  # Tame zoom with 0.05 multiplier

    # Use current focus point
    focus_x = max(0.0, min(1.0, params.get("focus_x", 0.5)))
    focus_y = max(0.0, min(1.0, params.get("focus_y", 0.5)))

    # Coordinate offset matching main noise
    world_offset_x = (focus_x - 0.5) * width
    world_offset_y = (focus_y - 0.5) * height

    base_offset_x = 1000.0
    base_offset_y = 1000.0

    # Output grass map
    grass_map = np.zeros((height, width), dtype=np.uint8)

    for y in range(height):
        for x in range(width):
            # Match coordinate sampling with base texture
            sample_x = base_offset_x + (x - world_offset_x) / adjusted_scale
            sample_y = base_offset_y + (y - world_offset_y) / adjusted_scale

            # Apply grass-specific scaling
            grass_sample_x = sample_x / grass_scale
            grass_sample_y = sample_y / grass_scale

            noise_value = snoise2(
                grass_sample_x,
                grass_sample_y,
                octaves=octaves,
                persistence=persistence,
                lacunarity=2.0,
                base=seed
            )

            normalized_noise = (noise_value + 1) / 2

            if normalized_noise < density:
                grass_map[y, x] = grass_amount
            else:
                grass_map[y, x] = 0

    return grass_map

def _apply_vignette(image, strength, radius):
    """Apply vignette effect to the image."""
    # Support both grayscale and color images
    if image.ndim == 2:
        # Grayscale image
        height, width = image.shape
        channels = 1
        is_gray = True
    else:
        # Color image
        height, width, channels = image.shape
        is_gray = False

    # Create vignette mask
    y, x = np.ogrid[:height, :width]
    center_x, center_y = width / 2, height / 2
    # Calculate distance from center
    distance = np.sqrt((x - center_x)**2 + (y - center_y)**2)
    max_distance = np.sqrt(center_x**2 + center_y**2)
    # Create vignette mask with adjustable radius
    vignette = 1 - np.clip((distance / (max_distance * radius)), 0, 1) * strength
    vignette = np.clip(vignette, 0, 1)
    # Apply vignette mask
    if is_gray:
        # For grayscale images, directly multiply and return 2D array
        return (image * vignette).astype(np.uint8)
    else:
        # For color images, apply per-channel and return image
        for i in range(channels):
            image[:, :, i] = (image[:, :, i] * vignette).astype(np.uint8)
        return image