Files
konabot/konabot/plugins/marchtoy/obj.py
2026-04-27 23:31:38 +08:00

279 lines
8.5 KiB
Python

import numpy as np
from konabot.plugins.marchtoy.texture import Texture
from konabot.plugins.marchtoy.utilities import ArgParser, Formatter
from nonebot import logger
from typing import Optional
import re
OBJECT_ENTRIES = {}
def make_obj(*name: str):
def decorator(cls):
# OBJECT_ENTRIES[name] = cls
for alias in [*name]:
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.linalg.norm(p) + 1e-8)
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 的性质 梯度大小会变 导致 overshoot 等问题
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):
x = x / 180 * np.pi
y = y / 180 * np.pi
z = z / 180 * np.pi
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.to_vec4(inv)} * vec4(p, 1.0)).xyz"
class Object:
def __init__(self) -> None:
self.transform: Transform = Transform()
self.texture: Texture = Texture()
self.round_corner: float = 0.0
def parse_args(self, args: list[str]):
raise NotImplementedError
def get_transformed(self, p: np.ndarray) -> np.ndarray:
if p.shape != (3,):
raise Exception(f"{p} is not a vec3")
p = p.copy()
inv = np.linalg.inv(self.transform.t)
return (inv @ np.array((*p, 1)))[:3]
def sdf_block_glsl(self) -> str:
raise NotImplementedError
@make_obj("cube", "box")
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: str):
self.size = ArgParser.as_vec3(args)
def sdf_block_glsl(self) -> str:
return f"sdCube({self.transform.p_expr()}, vec3({self.size[0]}, {self.size[1]}, {self.size[2]}))"
@make_obj("sphere", "ball")
class Sphere(Object):
def __init__(self, _radius: float = 1.0) -> None:
super().__init__()
self.radius = _radius
def parse_args(self, args: str):
self.radius = ArgParser.as_float(args)
def sdf_block_glsl(self) -> str:
return f"sdSphere({self.transform.p_expr()}, {self.radius})"
@make_obj("cylinder", "cyl")
class Cylinder(Object):
def __init__(self, _radius: float = 1.0, _height: float = 1.0) -> None:
super().__init__()
self.radius = _radius
self.height = _height
def parse_args(self, args: str):
param = ArgParser.as_vec2(args)
self.radius = param[0]
self.height = param[1]
def sdf_block_glsl(self) -> str:
return (
f"sdCappedCylinder({self.transform.p_expr()}, {self.radius}, {self.height})"
)
@make_obj("torus")
class Torus(Object):
def __init__(self, _r1: float = 1.0, _r2: float = 0.4) -> None:
super().__init__()
self.r1 = _r1
self.r2 = _r2
def parse_args(self, args: list[str]):
param = ArgParser.as_vec2(args)
self.r1 = param[0]
self.r2 = param[1]
def sdf_block_glsl(self) -> str:
return f"sdTorus({self.transform.p_expr()}, vec2({self.r1}, {self.r2}))"
@make_obj("capsule", "pill")
class Capsule(Object):
def __init__(self, _h: float = 1.0, _r: float = 0.25) -> None:
super().__init__()
self.h = _h
self.r = _r
def parse_args(self, args: list[str]):
param = ArgParser.as_vec2(args)
self.h = param[0]
self.r = param[1]
def sdf_block_glsl(self) -> str:
return f"sdVerticalCapsule({self.transform.p_expr()}, {self.h}, {self.r})"
@make_obj("smoothed", "mix")
class Smoothed(Object):
def __init__(self, _k: float = 12.00):
super().__init__()
self.block_a: str = "INF"
self.block_b: str = "INF"
self.k: float = _k
def parse_args(self, args: list[str]):
from konabot.plugins.marchtoy.scene import Scene
try:
if not len(args) >= 2:
raise Exception("expecting at least 2 args")
scene_a = Scene(args[0])
scene_b = Scene(args[1])
self.block_a = scene_a.canvas_objs[0][1]
self.block_b = scene_b.canvas_objs[0][1]
if len(args) > 2:
self.k = ArgParser.as_float(args[2], 0.25)
except Exception as e:
raise Exception(f"cannot build smoothed object over {args}: {e}")
def sdf_block_glsl(self):
return f"smin({self.block_a}, {self.block_b}, {self.k})"
@make_obj("intersect", "bool")
class BoolIntersect(Object):
def __init__(self):
super().__init__()
def parse_args(self, args: list[str]):
from konabot.plugins.marchtoy.scene import Scene
try:
if not len(args) >= 2:
raise Exception("expecting at least 2 args")
scene_a = Scene(args[0])
scene_b = Scene(args[1])
self.block_a = scene_a.canvas_objs[0][1]
self.block_b = scene_b.canvas_objs[0][1]
if len(args) > 2:
self.k = ArgParser.as_float(args[2], 0.25)
except Exception as e:
raise Exception(f"cannot build bool object over {args}: {e}")
def sdf_block_glsl(self):
return f"max({self.block_a}, {self.block_b})"
@make_obj("substract", "minus")
class BoolSubstract(Object):
def __init__(self):
super().__init__()
def parse_args(self, args: list[str]):
from konabot.plugins.marchtoy.scene import Scene
try:
if not len(args) >= 2:
raise Exception("expecting at least 2 args")
scene_a = Scene(args[0])
scene_b = Scene(args[1])
self.block_a = scene_a.canvas_objs[0][1]
self.block_b = scene_b.canvas_objs[0][1]
if len(args) > 2:
self.k = ArgParser.as_float(args[2], 0.25)
except Exception as e:
raise Exception(f"cannot build bool object over {args}: {e}")
def sdf_block_glsl(self):
return f"max({self.block_a}, -{self.block_b})"
@make_obj("add", "addition")
class BoolAddition(Object):
def __init__(self):
super().__init__()
def parse_args(self, args: list[str]):
from konabot.plugins.marchtoy.scene import Scene
try:
if not len(args) >= 2:
raise Exception("expecting at least 2 args")
scene_a = Scene(args[0])
scene_b = Scene(args[1])
self.block_a = scene_a.canvas_objs[0][1]
self.block_b = scene_b.canvas_objs[0][1]
except Exception as e:
raise Exception(f"cannot build bool object over {args}: {e}")
def sdf_block_glsl(self):
return f"min({self.block_a}, {self.block_b})"
@make_obj("camera", "cam")
class Camera(Object):
def __init__(self, _focus: float = 2.0) -> None:
super().__init__()
self.focus = _focus
# self.transform.translate(8.0, 8.0, 8.0).lookat(0.0, 0.0, 0.0)
def parse_args(self, args: list[str]):
self.focus = ArgParser.as_float(args)