77 lines
2.3 KiB
Python
77 lines
2.3 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
def calculate_spherical_harmonics(env_map, order=2):
|
|
"""
|
|
Calculate spherical harmonics coefficients from an environment map.
|
|
|
|
Args:
|
|
env_map: PIL Image or numpy array of the environment map
|
|
order: Order of spherical harmonics (default=2)
|
|
|
|
Returns:
|
|
coeffs: Spherical harmonics coefficients
|
|
"""
|
|
if isinstance(env_map, Image.Image):
|
|
ldr_image = np.array(env_map).astype(np.float32) / 255.0
|
|
else:
|
|
ldr_image = env_map.astype(np.float32) / 255.0
|
|
|
|
# Apply gamma correction
|
|
gamma = 2.2
|
|
linear_env = np.power(ldr_image, gamma)
|
|
|
|
# Calculate spherical harmonics coefficients
|
|
# (Implementation based on DreamLight's approach)
|
|
height, width, _ = linear_env.shape
|
|
num_samples = 10000
|
|
u = np.random.rand(num_samples)
|
|
v = np.random.rand(num_samples)
|
|
|
|
# Convert to spherical coordinates
|
|
theta = 2 * np.arccos(np.sqrt(1 - u))
|
|
phi = 2 * np.pi * v
|
|
|
|
# Sample environment map
|
|
x = (phi / (2 * np.pi)) * width
|
|
y = (theta / np.pi) * height
|
|
x = np.clip(x, 0, width - 1).astype(int)
|
|
y = np.clip(y, 0, height - 1).astype(int)
|
|
|
|
samples = linear_env[y, x]
|
|
|
|
# Compute spherical harmonics basis functions
|
|
# (Simplified implementation - actual SH calculation would be more complex)
|
|
coeffs = np.zeros((order + 1) ** 2)
|
|
# ... actual SH calculation would go here ...
|
|
|
|
return coeffs
|
|
|
|
def generate_spherical_image(coeffs, height, width):
|
|
"""
|
|
Generate a spherical image from spherical harmonics coefficients.
|
|
|
|
Args:
|
|
coeffs: Spherical harmonics coefficients
|
|
height: Output image height
|
|
width: Output image width
|
|
|
|
Returns:
|
|
spherical_image: Generated spherical image as a numpy array
|
|
"""
|
|
# Generate a grid of spherical coordinates
|
|
y, x = np.mgrid[0:height, 0:width]
|
|
phi = 2 * np.pi * x / width
|
|
theta = np.pi * y / height
|
|
|
|
# Compute spherical harmonics basis functions
|
|
# (Simplified implementation)
|
|
sh_values = np.zeros((height, width))
|
|
# ... actual SH evaluation would go here ...
|
|
|
|
# Normalize and convert to RGB
|
|
sh_values = np.clip(sh_values, 0, 1)
|
|
spherical_image = np.stack([sh_values] * 3, axis=-1) * 255
|
|
return spherical_image.astype(np.uint8)
|