forked from mttu-developers/konabot
124 lines
3.4 KiB
Python
124 lines
3.4 KiB
Python
import numpy as np
|
|
from texture import Texture
|
|
from utilities import ArgParser, Formatter
|
|
|
|
OBJECT_ENTRIES = {}
|
|
|
|
|
|
def make_obj(name: str, aliases: list[str] = []):
|
|
def decorator(cls):
|
|
OBJECT_ENTRIES[name] = cls
|
|
for alias in aliases:
|
|
OBJECT_ENTRIES[alias] = cls
|
|
return cls
|
|
|
|
return decorator
|
|
|
|
|
|
class Transform:
|
|
def __init__(self) -> None:
|
|
self.t: np.ndarray = np.identity(4, dtype=np.float32)
|
|
|
|
@staticmethod
|
|
def normalize(p: np.ndarray) -> np.ndarray:
|
|
return p / np.sqrt(np.dot(p, p))
|
|
|
|
def translate(self, x: float, y: float, z: float):
|
|
mat = np.identity(4, dtype=np.float32)
|
|
mat[0:3, 3] = [x, y, z]
|
|
self.t = mat @ self.t
|
|
return self
|
|
|
|
# scale 会破坏 sdf 的性质，不应该使用
|
|
def scale(self, x: float, y: float, z: float):
|
|
mat = np.identity(4, dtype=np.float32)
|
|
mat[0, 0], mat[1, 1], mat[2, 2] = x, y, z
|
|
self.t = mat @ self.t
|
|
return self
|
|
|
|
def rotate(self, x: float, y: float, z: float):
|
|
cx, sx = np.cos(x), np.sin(x)
|
|
cy, sy = np.cos(y), np.sin(y)
|
|
cz, sz = np.cos(z), np.sin(z)
|
|
|
|
mat = np.identity(4, dtype=np.float32)
|
|
mat[0, 0] = cy * cz
|
|
mat[0, 1] = sx * sy * cz - cx * sz
|
|
mat[0, 2] = cx * sy * cz + sx * sz
|
|
|
|
mat[1, 0] = cy * sz
|
|
mat[1, 1] = sx * sy * sz + cx * cz
|
|
mat[1, 2] = cx * sy * sz - sx * cz
|
|
|
|
mat[2, 0] = -sy
|
|
mat[2, 1] = sx * cy
|
|
mat[2, 2] = cx * cy
|
|
|
|
self.t = mat @ self.t
|
|
return self
|
|
|
|
def lookat(
|
|
self, x: float, y: float, z: float, up: np.ndarray = np.array([0.0, 1.0, 0.0])
|
|
):
|
|
p = self.t[0:3, 3]
|
|
q = np.array([x, y, z])
|
|
zaxis = Transform.normalize(q - p)
|
|
xaxis = Transform.normalize(np.cross(zaxis, up))
|
|
yaxis = Transform.normalize(np.cross(xaxis, zaxis))
|
|
self.t[:3, 0] = xaxis
|
|
self.t[:3, 1] = yaxis
|
|
self.t[:3, 2] = -zaxis
|
|
return self
|
|
|
|
def p_expr(self) -> str:
|
|
inv = np.linalg.inv(self.t) # + 1e-5 * np.identity(4, dtype=np.float32))
|
|
return f"({Formatter.float4(inv)} * float4(p, 1.0)).xyz"
|
|
|
|
|
|
class Object:
|
|
def __init__(self) -> None:
|
|
self.transform: Transform = Transform()
|
|
self.texture: Texture = Texture()
|
|
|
|
def parse_args(self, args: list[str]):
|
|
raise NotImplementedError
|
|
|
|
def sdf_block(self) -> str:
|
|
raise NotImplementedError
|
|
|
|
|
|
@make_obj("cube")
|
|
class Cube(Object):
|
|
def __init__(self, _size: np.ndarray = np.array([1.0, 1.0, 1.0])) -> None:
|
|
super().__init__()
|
|
self.size = _size
|
|
|
|
def parse_args(self, args: list[str]):
|
|
self.size = ArgParser.as_vec3(args)
|
|
|
|
def sdf_block(self) -> str:
|
|
return f"sdCube({self.transform.p_expr()}, float3({self.size[0]}, {self.size[1]}, {self.size[2]}))"
|
|
|
|
|
|
@make_obj("sphere")
|
|
class Sphere(Object):
|
|
def __init__(self, _radius: float = 1.0) -> None:
|
|
super().__init__()
|
|
self.radius = _radius
|
|
|
|
def parse_args(self, args: list[str]):
|
|
self.radius = ArgParser.as_float(args)
|
|
|
|
def sdf_block(self) -> str:
|
|
return f"sdSphere({self.transform.p_expr()}, {self.radius})"
|
|
|
|
|
|
@make_obj("camera", ["cam"])
|
|
class Camera(Object):
|
|
def __init__(self, _focus: float = 1.0) -> None:
|
|
super().__init__()
|
|
self.focus = _focus
|
|
|
|
def parse_args(self, args: list[str]):
|
|
self.focus = ArgParser.as_float(args)
|