135 lines
4.3 KiB
Python
135 lines
4.3 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
class Ball2Envmap:
|
|
"""
|
|
DiffusionLight Ball2Envmap class
|
|
|
|
"""
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"chromeball": ("IMAGE",),
|
|
"envmap_height": ("INT", {"default": 256, "min": 1, "max": 8192, "step": 1, "label": "Image Size"}),
|
|
"anti_aliasing": (["1", "2", "4", "8", "16"], {"label": "MSAA Anti-Aliasing Scale", "default": "4"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
FUNCTION = "convert"
|
|
|
|
def convert(self, chromeball, anti_aliasing="4", envmap_height=256):
|
|
"""
|
|
Convert an environment map to a ball2envmap format.
|
|
|
|
Args:
|
|
chromeball (IMAGE): The input environment map image. #Tensor of image format shape (range 0-1) shape [1, H, W, 3]
|
|
|
|
Returns:
|
|
tuple: A tuple containing the converted image.
|
|
"""
|
|
# Assuming envmap is already in the correct format
|
|
msaa_scale = int(anti_aliasing)
|
|
envmap = ball2envmap(chromeball, msaa_scale, envmap_height)
|
|
return (envmap, )
|
|
|
|
|
|
# HELPER FUNCTION
|
|
def ball2envmap(chromeball, msaa_scale, envmap_height):
|
|
"""
|
|
Helper function to convert a chromeball image to an environment map.
|
|
|
|
Args:
|
|
chromeball (torch.Tensor): The input chromeball image tensor.
|
|
|
|
Returns:
|
|
torch.Tensor: The converted environment map tensor.
|
|
"""
|
|
# Assuming the input is already in the correct format
|
|
|
|
I = np.array([1,0, 0]) # incoming vector, pointing to the camera
|
|
|
|
# compute normal map that create from reflect vector
|
|
env_grid = create_envmap_grid(envmap_height * msaa_scale)
|
|
reflect_vec = get_cartesian_from_spherical(env_grid[...,1], env_grid[...,0])
|
|
normal = get_normal_vector(I[None,None], reflect_vec)
|
|
|
|
# turn from normal map to position to lookup [Range: 0,1]
|
|
pos = (normal + 1.0) / 2
|
|
pos = 1.0 - pos
|
|
pos = pos[...,1:]
|
|
|
|
env_map = None
|
|
|
|
# using pytorch method for bilinear interpolation
|
|
with torch.no_grad():
|
|
# convert position to pytorch grid look up
|
|
grid = torch.from_numpy(pos)[None].float()
|
|
grid = grid * 2 - 1 # convert to range [-1,1]
|
|
|
|
# convert ball to support pytorch
|
|
ball_image = chromeball.permute(0,3,1,2) # [1,3,H,W]
|
|
envmap = torch.nn.functional.grid_sample(ball_image, grid, mode='bilinear', padding_mode='border', align_corners=True)
|
|
envmap = torch.nn.functional.interpolate(envmap, size=(envmap_height, envmap_height * 2), mode='bilinear', align_corners=False)
|
|
envmap = envmap.permute(0,2,3,1) # [1,H,W,3]
|
|
|
|
return envmap
|
|
|
|
def create_envmap_grid(size: int):
|
|
"""
|
|
BLENDER CONVENSION
|
|
Create the grid of environment map that contain the position in sperical coordinate
|
|
Top left is (0,0) and bottom right is (pi/2, 2pi)
|
|
"""
|
|
|
|
theta = torch.linspace(0, np.pi * 2, size * 2)
|
|
phi = torch.linspace(0, np.pi, size)
|
|
|
|
#use indexing 'xy' torch match vision's homework 3
|
|
theta, phi = torch.meshgrid(theta, phi ,indexing='xy')
|
|
theta_phi = torch.cat([theta[..., None], phi[..., None]], dim=-1)
|
|
theta_phi = theta_phi.numpy()
|
|
|
|
|
|
|
|
return theta_phi
|
|
|
|
def get_cartesian_from_spherical(theta: np.array, phi: np.array, r = 1.0):
|
|
"""
|
|
BLENDER CONVENSION
|
|
theta: vertical angle
|
|
phi: horizontal angle
|
|
r: radius
|
|
"""
|
|
x = r * np.sin(theta) * np.cos(phi)
|
|
y = r * np.sin(theta) * np.sin(phi)
|
|
z = r * np.cos(theta)
|
|
return np.concatenate([x[...,None],y[...,None],z[...,None]], axis=-1)
|
|
|
|
def get_normal_vector(incoming_vector: np.ndarray, reflect_vector: np.ndarray):
|
|
"""
|
|
BLENDER CONVENSION
|
|
incoming_vector: the vector from the point to the camera
|
|
reflect_vector: the vector from the point to the light source
|
|
"""
|
|
#N = 2(R ⋅ I)R - I
|
|
N = (incoming_vector + reflect_vector) / np.linalg.norm(incoming_vector + reflect_vector, axis=-1, keepdims=True)
|
|
return N
|
|
|
|
def get_cartesian_from_spherical(theta: np.array, phi: np.array, r = 1.0):
|
|
"""
|
|
BLENDER CONVENSION
|
|
theta: vertical angle
|
|
phi: horizontal angle
|
|
r: radius
|
|
"""
|
|
x = r * np.sin(theta) * np.cos(phi)
|
|
y = r * np.sin(theta) * np.sin(phi)
|
|
z = r * np.cos(theta)
|
|
return np.concatenate([x[...,None],y[...,None],z[...,None]], axis=-1)
|