konabot/konabot/plugins/marchtoy/obj.py

import numpy as np
from konabot.plugins.marchtoy.texture import Texture
from konabot.plugins.marchtoy.utilities import ArgParser, Formatter

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.sqrt(np.dot(p, 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):
        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)} * vec4(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", "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: list[str]):
        self.size = ArgParser.as_vec3(args)

    def sdf_block(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: list[str]):
        self.radius = ArgParser.as_float(args)

    def sdf_block(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: list[str]):
        param = ArgParser.as_vec2(args)
        self.radius = param[0]
        self.height = param[1]

    def sdf_block(self) -> str:
        return f"sdCappedCylinder({self.transform.p_expr()}, {self.radius}, {self.height})"



@make_obj("torus")
class Torus(Object):
    def __init__(self, _r1: float = 0.5, _r2: float = 1.0) -> 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(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(self) -> str:
        return f"sdVerticalCapsule({self.transform.p_expr()}, {self._h}, {self._r})"



@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)