102 lines
4.0 KiB
Python
102 lines
4.0 KiB
Python
def generate_normal_map_from_alpha(img, strength=1.0, blur_radius=1.0, invert=False):
|
|
"""Generate a normal map from an image's alpha channel using height-based algorithm"""
|
|
try:
|
|
from PIL import Image, ImageFilter
|
|
import numpy as np
|
|
|
|
print(f"[Normal Map] Starting generation with strength={strength}, blur={blur_radius}, invert={invert}")
|
|
|
|
# Extract alpha channel as height map
|
|
if img.mode != 'RGBA':
|
|
img = img.convert('RGBA')
|
|
|
|
# Get alpha channel
|
|
alpha_channel = img.split()[3] # Alpha is the 4th channel
|
|
width, height = alpha_channel.size
|
|
|
|
print(f"[Normal Map] Processing {width}x{height} alpha channel")
|
|
|
|
# Apply blur if specified
|
|
if blur_radius > 0:
|
|
alpha_channel = alpha_channel.filter(ImageFilter.GaussianBlur(radius=blur_radius))
|
|
print(f"[Normal Map] Applied blur with radius {blur_radius}")
|
|
|
|
# Convert to numpy array for gradient calculations
|
|
height_map = np.array(alpha_channel, dtype=np.float32) / 255.0
|
|
|
|
# Apply strength multiplier
|
|
height_map *= strength
|
|
|
|
# Calculate gradients using Sobel operators
|
|
# Sobel X kernel for horizontal gradients
|
|
sobel_x = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float32)
|
|
# Sobel Y kernel for vertical gradients
|
|
sobel_y = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=np.float32)
|
|
|
|
# Pad the height map to handle edges
|
|
padded_height = np.pad(height_map, ((1, 1), (1, 1)), mode='edge')
|
|
|
|
# Calculate gradients
|
|
grad_x = np.zeros_like(height_map)
|
|
grad_y = np.zeros_like(height_map)
|
|
|
|
for i in range(height):
|
|
for j in range(width):
|
|
# Extract 3x3 neighborhood
|
|
neighborhood = padded_height[i:i+3, j:j+3]
|
|
# Apply Sobel operators
|
|
grad_x[i, j] = np.sum(neighborhood * sobel_x)
|
|
grad_y[i, j] = np.sum(neighborhood * sobel_y)
|
|
|
|
print(f"[Normal Map] Calculated gradients")
|
|
|
|
# Convert gradients to normal vectors
|
|
# Normal map RGB values are calculated as:
|
|
# R = (grad_x + 1) * 0.5 -> maps -1,1 to 0,1
|
|
# G = (-grad_y + 1) * 0.5 -> maps -1,1 to 0,1 (Y is flipped for standard normal maps)
|
|
# B = sqrt(1 - grad_x^2 - grad_y^2) -> Z component, always pointing up
|
|
|
|
# Clamp gradients to reasonable range
|
|
grad_x = np.clip(grad_x, -1, 1)
|
|
grad_y = np.clip(grad_y, -1, 1)
|
|
|
|
# Apply invert if specified
|
|
if invert:
|
|
grad_x = -grad_x
|
|
grad_y = -grad_y
|
|
print(f"[Normal Map] Applied inversion")
|
|
|
|
# Calculate normal map channels
|
|
# Red channel: X gradient mapped to 0-1
|
|
normal_r = ((grad_x + 1.0) * 0.5 * 255).astype(np.uint8)
|
|
|
|
# Green channel: Y gradient mapped to 0-1 (flipped)
|
|
normal_g = ((-grad_y + 1.0) * 0.5 * 255).astype(np.uint8)
|
|
|
|
# Blue channel: Z component (pointing up)
|
|
# Calculate Z from X and Y to maintain unit length
|
|
grad_magnitude_sq = grad_x**2 + grad_y**2
|
|
grad_z = np.sqrt(np.maximum(0, 1.0 - grad_magnitude_sq))
|
|
normal_b = (grad_z * 255).astype(np.uint8)
|
|
|
|
print(f"[Normal Map] Calculated normal vectors")
|
|
|
|
# Create the normal map image
|
|
normal_map = Image.new('RGB', (width, height))
|
|
|
|
# Combine channels
|
|
for y in range(height):
|
|
for x in range(width):
|
|
r = int(normal_r[y, x])
|
|
g = int(normal_g[y, x])
|
|
b = int(normal_b[y, x])
|
|
normal_map.putpixel((x, y), (r, g, b))
|
|
|
|
print(f"[Normal Map] Normal map generation completed successfully")
|
|
return normal_map
|
|
|
|
except Exception as e:
|
|
print(f"[Normal Map] Error generating normal map: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return None |