diff --git a/hippolyzer/lib/base/mesh.py b/hippolyzer/lib/base/mesh.py index b860081..46f9d1c 100644 --- a/hippolyzer/lib/base/mesh.py +++ b/hippolyzer/lib/base/mesh.py @@ -17,6 +17,7 @@ import recordclass from hippolyzer.lib.base import serialization as se from hippolyzer.lib.base.datatypes import Vector3, Vector2, UUID, TupleCoord from hippolyzer.lib.base.llsd import zip_llsd, unzip_llsd +from hippolyzer.lib.base.serialization import ParseContext LOG = logging.getLogger(__name__) @@ -283,7 +284,7 @@ class VertexWeights(se.SerializableBase): def deserialize(cls, reader: se.Reader, ctx=None): influence_list = [] for _ in range(cls.INFLUENCE_LIMIT): - joint_idx = reader.read(se.U8) + joint_idx = reader.read_bytes(1)[0] if joint_idx == cls.INFLUENCE_TERM: break influence_list.append(VertexWeight(joint_idx, reader.read(cls.INFLUENCE_SER, ctx=ctx))) @@ -321,16 +322,46 @@ class SegmentSerializer: return new_segment +class VecListAdapter(se.Adapter): + def __init__(self, child_spec: se.SERIALIZABLE_TYPE, vec_type: Type): + super().__init__(child_spec) + self.vec_type = vec_type + + def encode(self, val: Any, ctx: Optional[ParseContext]) -> Any: + return val + + def decode(self, val: Any, ctx: Optional[ParseContext], pod: bool = False) -> Any: + new_vals = [] + for elem in val: + new_vals.append(self.vec_type(*elem)) + return new_vals + + +LE_U16: np.dtype = np.dtype(np.uint16).newbyteorder('<') # noqa + + LOD_SEGMENT_SERIALIZER = SegmentSerializer({ # 16-bit indices to the verts making up the tri. Imposes a 16-bit # upper limit on verts in any given material in the mesh. - "TriangleList": se.Collection(None, se.Collection(3, se.U16)), + "TriangleList": se.ExprAdapter( + se.NumPyArray(se.BytesGreedy(), LE_U16, 3), + decode_func=lambda x: x.tolist(), + ), # These are used to interpolate between values in their respective domains # Each position represents a single vert. - "Position": se.Collection(None, se.Vector3U16(0.0, 1.0)), - "TexCoord0": se.Collection(None, se.Vector2U16(0.0, 1.0)), + "Position": VecListAdapter( + se.QuantizedNumPyArray(se.NumPyArray(se.BytesGreedy(), LE_U16, 3), 0.0, 1.0), + Vector3, + ), + "TexCoord0": VecListAdapter( + se.QuantizedNumPyArray(se.NumPyArray(se.BytesGreedy(), LE_U16, 2), 0.0, 1.0), + Vector2, + ), # Normals have a static domain between -1 and 1, so just use that. - "Normal": se.Collection(None, se.Vector3U16(-1.0, 1.0)), + "Normal": VecListAdapter( + se.QuantizedNumPyArray(se.NumPyArray(se.BytesGreedy(), LE_U16, 3), -1.0, 1.0), + Vector3, + ), "Weights": se.Collection(None, VertexWeights) }) diff --git a/hippolyzer/lib/base/serialization.py b/hippolyzer/lib/base/serialization.py index 67cf122..f0c6411 100644 --- a/hippolyzer/lib/base/serialization.py +++ b/hippolyzer/lib/base/serialization.py @@ -10,6 +10,7 @@ from io import SEEK_CUR, SEEK_SET, SEEK_END, RawIOBase, BufferedIOBase from typing import * import lazy_object_proxy +import numpy as np import hippolyzer.lib.base.llsd as llsd import hippolyzer.lib.base.datatypes as dtypes @@ -838,7 +839,7 @@ class QuantizedFloat(QuantizedFloatBase): super().__init__(prim_spec, zero_median=False) self.lower = lower self.upper = upper - # We know the range in `QuantizedFloat` when it's constructed, so we can infer + # We know the range in `QuantizedFloat` when it's constructed, so we can infer # whether or not we should round towards zero in __init__ max_error = (upper - lower) * self.step_mag midpoint = (upper + lower) / 2.0 @@ -1610,7 +1611,9 @@ class BitfieldDataclass(DataclassAdapter): class ExprAdapter(Adapter): - def __init__(self, child_spec: SERIALIZABLE_TYPE, decode_func: Callable, encode_func: Callable): + _ID = lambda x: x + + def __init__(self, child_spec: SERIALIZABLE_TYPE, decode_func: Callable = _ID, encode_func: Callable = _ID): super().__init__(child_spec) self._decode_func = decode_func self._encode_func = encode_func @@ -1659,6 +1662,62 @@ class BinaryLLSD(SerializableBase): writer.write_bytes(llsd.format_binary(val, with_header=False)) +class NumPyArray(Adapter): + """ + An 2-dimensional, dynamic-length array of data from numpy. Greedy. + + Unlike most other serializers, your endianness _must_ be specified in the dtype! + """ + __slots__ = ['dtype', 'elems'] + + def __init__(self, child_spec: Optional[SERIALIZABLE_TYPE], dtype: np.dtype, elems: int): + super().__init__(child_spec) + self.dtype = dtype + self.elems = elems + + def _pick_dtype(self, endian: str) -> np.dtype: + return self.dtype.newbyteorder('>') if endian != "<" else self.dtype + + def decode(self, val: Any, ctx: Optional[ParseContext], pod: bool = False) -> Any: + num_elems = len(val) // self.dtype.itemsize + num_ndims = num_elems // self.elems + buf_array = np.frombuffer(val, dtype=self.dtype, count=num_elems) + return buf_array.reshape((num_ndims, self.elems)) + + def encode(self, val, ctx: Optional[ParseContext]) -> Any: + val: np.ndarray = np.array(val, dtype=self.dtype).flatten() + return val.tobytes() + + +class QuantizedNumPyArray(Adapter): + """Like QuantizedFloat. Only works correctly for unsigned types, no zero midpoint rounding!""" + def __init__(self, child_spec: NumPyArray, lower: float, upper: float): + super().__init__(child_spec) + self.dtype = child_spec.dtype + self.lower = lower + self.upper = upper + self.step_mag = 1.0 / ((2 ** (self.dtype.itemsize * 8)) - 1) + + def encode(self, val: Any, ctx: Optional[ParseContext]) -> Any: + val = np.array(val, dtype=np.float64) + val = np.clip(val, self.lower, self.upper) + delta = self.upper - self.lower + if delta == 0.0: + return np.zeros(val.shape, dtype=self.dtype) + + val -= self.lower + val /= delta + val /= self.step_mag + return np.rint(val).astype(self.dtype) + + def decode(self, val: Any, ctx: Optional[ParseContext], pod: bool = False) -> Any: + val = val.astype(np.float64) + val *= self.step_mag + val *= self.upper - self.lower + val += self.lower + return val + + def subfield_serializer(msg_name, block_name, var_name): def f(orig_cls): global SUBFIELD_SERIALIZERS diff --git a/tests/base/test_serialization.py b/tests/base/test_serialization.py index e663511..f1a3541 100644 --- a/tests/base/test_serialization.py +++ b/tests/base/test_serialization.py @@ -6,6 +6,8 @@ import uuid from io import BytesIO from typing import Optional +import numpy as np + from hippolyzer.lib.base.datatypes import * import hippolyzer.lib.base.serialization as se from hippolyzer.lib.base.llanim import Animation, Joint, RotKeyframe @@ -693,6 +695,46 @@ class NameValueSerializationTests(BaseSerializationTest): deser.to_dict() +class NumPySerializationTests(BaseSerializationTest): + def setUp(self) -> None: + super().setUp() + self.writer.endianness = "<" + + def test_simple(self): + quant_spec = se.Vector3U16(0.0, 1.0) + self.writer.write(quant_spec, Vector3(0, 0.1, 0)) + self.writer.write(quant_spec, Vector3(1, 1, 1)) + + reader = self._get_reader() + np_spec = se.NumPyArray(se.BytesGreedy(), np.dtype(np.uint16), 3) + np_val = reader.read(np_spec) + expected_arr = np.array([[0, 6554, 0], [0xFFFF, 0xFFFF, 0xFFFF]], dtype=np.uint16) + np.testing.assert_array_equal(expected_arr, np_val) + + # Make sure writing the array back works correctly + orig_buf = self.writer.copy_buffer() + self.writer.clear() + self.writer.write(np_spec, expected_arr) + self.assertEqual(orig_buf, self.writer.copy_buffer()) + + def test_quantization(self): + quant_spec = se.Vector3U16(0.0, 1.0) + self.writer.write(quant_spec, Vector3(0, 0.1, 0)) + self.writer.write(quant_spec, Vector3(1, 1, 1)) + + reader = self._get_reader() + np_spec = se.QuantizedNumPyArray(se.NumPyArray(se.BytesGreedy(), np.dtype(np.uint16), 3), 0.0, 1.0) + np_val = reader.read(np_spec) + expected_arr = np.array([[0, 0.1, 0], [1, 1, 1]], dtype=np.float64) + np.testing.assert_array_almost_equal(expected_arr, np_val, decimal=5) + + # Make sure writing the array back works correctly + orig_buf = self.writer.copy_buffer() + self.writer.clear() + self.writer.write(np_spec, expected_arr) + self.assertEqual(orig_buf, self.writer.copy_buffer()) + + class AnimSerializationTests(BaseSerializationTest): SIMPLE_ANIM = b'\x01\x00\x00\x00\x01\x00\x00\x00H\x11\xd1?\x00\x00\x00\x00\x00H\x11\xd1?\x00\x00\x00\x00' \ b'\xcd\xccL>\x9a\x99\x99>\x01\x00\x00\x00\x02\x00\x00\x00mNeck\x00\x01\x00\x00\x00\x03\x00' \