kinfer 0.3.2__cp312-cp312-macosx_11_0_arm64.whl → 0.4.0__cp312-cp312-macosx_11_0_arm64.whl

kinfer/serialize/numpy.py

@@ -1,426 +0,0 @@
-"""Defines a serializer for Numpy arrays."""
-
-from typing import cast
-
-import numpy as np
-
-from kinfer import proto as P
-from kinfer.serialize.base import (
-    AudioFrameSerializer,
-    CameraFrameSerializer,
-    ImuSerializer,
-    JointCommandsSerializer,
-    JointPositionsSerializer,
-    JointTorquesSerializer,
-    JointVelocitiesSerializer,
-    MultiSerializer,
-    Serializer,
-    StateTensorSerializer,
-    TimestampSerializer,
-    VectorCommandSerializer,
-)
-from kinfer.serialize.utils import (
-    as_float,
-    check_names_match,
-    convert_angular_position,
-    convert_angular_velocity,
-    convert_torque,
-    dtype_num_bytes,
-    dtype_range,
-    numpy_dtype,
-    parse_bytes,
-)
-
-
-class NumpyBaseSerializer:
-    def __init__(self: "NumpyBaseSerializer", dtype: np.dtype | None = None) -> None:
-        self.dtype = dtype or np.float32
-
-
-class NumpyJointPositionsSerializer(NumpyBaseSerializer, JointPositionsSerializer[np.ndarray]):
-    def serialize_joint_positions(
-        self: "NumpyJointPositionsSerializer",
-        schema: P.JointPositionsSchema,
-        value: P.JointPositionsValue,
-    ) -> np.ndarray:
-        value_map = {v.joint_name: v for v in value.values}
-        check_names_match("schema", schema.joint_names, "value", list(value_map.keys()))
-        array = np.array(
-            [
-                convert_angular_position(value_map[name].value, value_map[name].unit, schema.unit)
-                for name in schema.joint_names
-            ],
-            dtype=self.dtype,
-        )
-
-        return array
-
-    def deserialize_joint_positions(
-        self: "NumpyJointPositionsSerializer",
-        schema: P.JointPositionsSchema,
-        value: np.ndarray,
-    ) -> P.JointPositionsValue:
-        if value.shape != (len(schema.joint_names),):
-            raise ValueError(
-                f"Shape of array must match number of joint names: {value.shape} != {len(schema.joint_names)}"
-            )
-        value_list = cast(list[float], value.astype(float).tolist())
-        return P.JointPositionsValue(
-            values=[
-                P.JointPositionValue(
-                    joint_name=name,
-                    value=float(value_list[i]),
-                    unit=schema.unit,
-                )
-                for i, name in enumerate(schema.joint_names)
-            ]
-        )
-
-
-class NumpyJointVelocitiesSerializer(NumpyBaseSerializer, JointVelocitiesSerializer[np.ndarray]):
-    def serialize_joint_velocities(
-        self: "NumpyJointVelocitiesSerializer",
-        schema: P.JointVelocitiesSchema,
-        value: P.JointVelocitiesValue,
-    ) -> np.ndarray:
-        value_map = {v.joint_name: v for v in value.values}
-        check_names_match("schema", schema.joint_names, "value", list(value_map.keys()))
-        array = np.array(
-            [
-                convert_angular_velocity(value_map[name].value, value_map[name].unit, schema.unit)
-                for name in schema.joint_names
-            ],
-            dtype=self.dtype,
-        )
-        return array
-
-    def deserialize_joint_velocities(
-        self: "NumpyJointVelocitiesSerializer",
-        schema: P.JointVelocitiesSchema,
-        value: np.ndarray,
-    ) -> P.JointVelocitiesValue:
-        if value.shape != (len(schema.joint_names),):
-            raise ValueError(
-                f"Shape of array must match number of joint names: {value.shape} != {len(schema.joint_names)}"
-            )
-        value_list = cast(list[float], value.astype(float).tolist())
-        return P.JointVelocitiesValue(
-            values=[
-                P.JointVelocityValue(joint_name=name, value=value_list[i], unit=schema.unit)
-                for i, name in enumerate(schema.joint_names)
-            ]
-        )
-
-
-class NumpyJointTorquesSerializer(NumpyBaseSerializer, JointTorquesSerializer[np.ndarray]):
-    def serialize_joint_torques(
-        self: "NumpyJointTorquesSerializer",
-        schema: P.JointTorquesSchema,
-        value: P.JointTorquesValue,
-    ) -> np.ndarray:
-        value_map = {v.joint_name: v for v in value.values}
-        check_names_match("schema", schema.joint_names, "value", list(value_map.keys()))
-        array = np.array(
-            [convert_torque(value_map[name].value, value_map[name].unit, schema.unit) for name in schema.joint_names],
-            dtype=self.dtype,
-        )
-        return array
-
-    def deserialize_joint_torques(
-        self: "NumpyJointTorquesSerializer",
-        schema: P.JointTorquesSchema,
-        value: np.ndarray,
-    ) -> P.JointTorquesValue:
-        if value.shape != (len(schema.joint_names),):
-            raise ValueError(
-                f"Shape of array must match number of joint names: {value.shape} != {len(schema.joint_names)}"
-            )
-        value_list = cast(list[float], value.astype(float).tolist())
-        return P.JointTorquesValue(
-            values=[
-                P.JointTorqueValue(joint_name=name, value=float(value_list[i]), unit=schema.unit)
-                for i, name in enumerate(schema.joint_names)
-            ]
-        )
-
-
-class NumpyJointCommandsSerializer(NumpyBaseSerializer, JointCommandsSerializer[np.ndarray]):
-    def _convert_value_to_array(
-        self: "NumpyJointCommandsSerializer",
-        value: P.JointCommandValue,
-        schema: P.JointCommandsSchema,
-    ) -> np.ndarray:
-        return np.array(
-            [
-                convert_torque(value.torque, value.torque_unit, schema.torque_unit),
-                convert_angular_velocity(value.velocity, value.velocity_unit, schema.velocity_unit),
-                convert_angular_position(value.position, value.position_unit, schema.position_unit),
-                value.kp,
-                value.kd,
-            ],
-            dtype=self.dtype,
-        )
-
-    def _convert_array_to_value(
-        self: "NumpyJointCommandsSerializer",
-        values: list[float],
-        schema: P.JointCommandsSchema,
-        name: str,
-    ) -> P.JointCommandValue:
-        if len(values) != 5:
-            raise ValueError(f"Shape of array must match number of joint commands: {len(values)} != 5")
-        return P.JointCommandValue(
-            joint_name=name,
-            torque=values[0],
-            velocity=values[1],
-            position=values[2],
-            kp=values[3],
-            kd=values[4],
-            torque_unit=schema.torque_unit,
-            velocity_unit=schema.velocity_unit,
-            position_unit=schema.position_unit,
-        )
-
-    def serialize_joint_commands(
-        self: "NumpyJointCommandsSerializer",
-        schema: P.JointCommandsSchema,
-        value: P.JointCommandsValue,
-    ) -> np.ndarray:
-        value_map = {v.joint_name: v for v in value.values}
-        check_names_match("schema", schema.joint_names, "value", list(value_map.keys()))
-        array = np.stack(
-            [self._convert_value_to_array(value_map[name], schema) for name in schema.joint_names],
-            axis=0,
-        )
-        return array
-
-    def deserialize_joint_commands(
-        self: "NumpyJointCommandsSerializer",
-        schema: P.JointCommandsSchema,
-        value: np.ndarray,
-    ) -> P.JointCommandsValue:
-        if value.shape != (len(schema.joint_names), 5):
-            raise ValueError(
-                "Shape of array must match number of joint names and commands: "
-                f"{value.shape} != ({len(schema.joint_names)}, 5)"
-            )
-        value_list = cast(list[list[float]], value.astype(float).tolist())
-        return P.JointCommandsValue(
-            values=[
-                self._convert_array_to_value(value_list[i], schema, name) for i, name in enumerate(schema.joint_names)
-            ]
-        )
-
-
-class NumpyCameraFrameSerializer(NumpyBaseSerializer, CameraFrameSerializer[np.ndarray]):
-    def serialize_camera_frame(
-        self: "NumpyCameraFrameSerializer",
-        schema: P.CameraFrameSchema,
-        value: P.CameraFrameValue,
-    ) -> np.ndarray:
-        np_arr = parse_bytes(value.data, P.DType.UINT8)
-        array = np_arr.astype(self.dtype) / 255.0
-        if array.size != schema.channels * schema.height * schema.width:
-            raise ValueError(
-                "Length of data must match number of channels, height, and width: "
-                f"{array.size} != {schema.channels} * {schema.height} * {schema.width}"
-            )
-        array = array.reshape(schema.channels, schema.height, schema.width)
-        return array
-
-    def deserialize_camera_frame(
-        self: "NumpyCameraFrameSerializer",
-        schema: P.CameraFrameSchema,
-        value: np.ndarray,
-    ) -> P.CameraFrameValue:
-        np_arr = (value * 255.0).flatten().astype(np.uint8)
-        return P.CameraFrameValue(data=np_arr.tobytes())
-
-
-class NumpyAudioFrameSerializer(NumpyBaseSerializer, AudioFrameSerializer[np.ndarray]):
-    def serialize_audio_frame(
-        self: "NumpyAudioFrameSerializer",
-        schema: P.AudioFrameSchema,
-        value: P.AudioFrameValue,
-    ) -> np.ndarray:
-        value_bytes = value.data
-        if len(value_bytes) != schema.channels * schema.sample_rate * dtype_num_bytes(schema.dtype):
-            raise ValueError(
-                "Length of data must match number of channels, sample rate, and dtype: "
-                f"{len(value_bytes)} != {schema.channels} * {schema.sample_rate} * {dtype_num_bytes(schema.dtype)}"
-            )
-        _, max_value = dtype_range(schema.dtype)
-        np_arr = parse_bytes(value_bytes, schema.dtype)
-        array = np_arr.astype(self.dtype)
-        array = array.reshape(schema.channels, -1)
-        array = array / max_value
-        return array
-
-    def deserialize_audio_frame(
-        self: "NumpyAudioFrameSerializer",
-        schema: P.AudioFrameSchema,
-        value: np.ndarray,
-    ) -> P.AudioFrameValue:
-        _, max_value = dtype_range(schema.dtype)
-        np_arr = (value * max_value).flatten().astype(numpy_dtype(schema.dtype))
-        return P.AudioFrameValue(data=np_arr.tobytes())
-
-
-class NumpyImuSerializer(NumpyBaseSerializer, ImuSerializer[np.ndarray]):
-    def serialize_imu(
-        self: "NumpyImuSerializer",
-        schema: P.ImuSchema,
-        value: P.ImuValue,
-    ) -> np.ndarray:
-        vectors = []
-        if schema.use_accelerometer:
-            vectors.append(
-                np.array(
-                    [value.linear_acceleration.x, value.linear_acceleration.y, value.linear_acceleration.z],
-                    dtype=self.dtype,
-                )
-            )
-        if schema.use_gyroscope:
-            vectors.append(
-                np.array(
-                    [value.angular_velocity.x, value.angular_velocity.y, value.angular_velocity.z],
-                    dtype=self.dtype,
-                )
-            )
-        if schema.use_magnetometer:
-            vectors.append(
-                np.array(
-                    [value.magnetic_field.x, value.magnetic_field.y, value.magnetic_field.z],
-                    dtype=self.dtype,
-                )
-            )
-        if not vectors:
-            raise ValueError("IMU has nothing to serialize")
-        return np.stack(vectors, axis=0)
-
-    def deserialize_imu(
-        self: "NumpyImuSerializer",
-        schema: P.ImuSchema,
-        value: np.ndarray,
-    ) -> P.ImuValue:
-        num_vectors = sum([schema.use_accelerometer, schema.use_gyroscope, schema.use_magnetometer])
-        if value.shape != (num_vectors, 3):
-            raise ValueError(
-                f"Shape of array must match number of vectors and components: {value.shape} != ({num_vectors}, 3)"
-            )
-        vectors = cast(list[list[float]], value.astype(float).tolist())
-        imu_value = P.ImuValue()
-        if schema.use_accelerometer:
-            x, y, z = vectors.pop(0)
-            imu_value.linear_acceleration.x = as_float(x)
-            imu_value.linear_acceleration.y = as_float(y)
-            imu_value.linear_acceleration.z = as_float(z)
-        if schema.use_gyroscope:
-            x, y, z = vectors.pop(0)
-            imu_value.angular_velocity.x = as_float(x)
-            imu_value.angular_velocity.y = as_float(y)
-            imu_value.angular_velocity.z = as_float(z)
-        if schema.use_magnetometer:
-            x, y, z = vectors.pop(0)
-            imu_value.magnetic_field.x = as_float(x)
-            imu_value.magnetic_field.y = as_float(y)
-            imu_value.magnetic_field.z = as_float(z)
-        return imu_value
-
-
-class NumpyTimestampSerializer(NumpyBaseSerializer, TimestampSerializer[np.ndarray]):
-    def serialize_timestamp(
-        self: "NumpyTimestampSerializer",
-        schema: P.TimestampSchema,
-        value: P.TimestampValue,
-    ) -> np.ndarray:
-        elapsed_seconds = value.seconds - schema.start_seconds
-        elapsed_nanos = value.nanos - schema.start_nanos
-        if elapsed_nanos < 0:
-            elapsed_seconds -= 1
-            elapsed_nanos += 1_000_000_000
-        total_elapsed_seconds = elapsed_seconds + elapsed_nanos / 1_000_000_000
-        return np.array([total_elapsed_seconds], dtype=self.dtype)
-
-    def deserialize_timestamp(
-        self: "NumpyTimestampSerializer",
-        schema: P.TimestampSchema,
-        value: np.ndarray,
-    ) -> P.TimestampValue:
-        total_elapsed_seconds = float(value.item())
-        elapsed_seconds = int(total_elapsed_seconds)
-        elapsed_nanos = int((total_elapsed_seconds - elapsed_seconds) * 1_000_000_000)
-        return P.TimestampValue(seconds=elapsed_seconds, nanos=elapsed_nanos)
-
-
-class NumpyVectorCommandSerializer(NumpyBaseSerializer, VectorCommandSerializer[np.ndarray]):
-    def serialize_vector_command(
-        self: "NumpyVectorCommandSerializer",
-        schema: P.VectorCommandSchema,
-        value: P.VectorCommandValue,
-    ) -> np.ndarray:
-        return np.array(value.values, dtype=self.dtype)
-
-    def deserialize_vector_command(
-        self: "NumpyVectorCommandSerializer",
-        schema: P.VectorCommandSchema,
-        value: np.ndarray,
-    ) -> P.VectorCommandValue:
-        if value.shape != (schema.dimensions,):
-            raise ValueError(f"Shape of array must match number of dimensions: {value.shape} != {schema.dimensions}")
-        values = cast(list[float], value.astype(float).tolist())
-        return P.VectorCommandValue(values=values)
-
-
-class NumpyStateTensorSerializer(NumpyBaseSerializer, StateTensorSerializer[np.ndarray]):
-    def serialize_state_tensor(
-        self: "NumpyStateTensorSerializer",
-        schema: P.StateTensorSchema,
-        value: P.StateTensorValue,
-    ) -> np.ndarray:
-        value_bytes = value.data
-        if len(value_bytes) != np.prod(schema.shape) * dtype_num_bytes(schema.dtype):
-            raise ValueError(
-                "Length of data must match number of elements: "
-                f"{len(value_bytes)} != {np.prod(schema.shape)} * {dtype_num_bytes(schema.dtype)}"
-            )
-        np_arr = parse_bytes(value_bytes, schema.dtype)
-        array = np.ascontiguousarray(np_arr.astype(numpy_dtype(schema.dtype)))
-        array = array.reshape(tuple(schema.shape))
-        return array
-
-    def deserialize_state_tensor(
-        self: "NumpyStateTensorSerializer",
-        schema: P.StateTensorSchema,
-        value: np.ndarray,
-    ) -> P.StateTensorValue:
-        contiguous_value = np.ascontiguousarray(value)
-        return P.StateTensorValue(data=contiguous_value.flatten().tobytes())
-
-
-class NumpySerializer(
-    NumpyJointPositionsSerializer,
-    NumpyJointVelocitiesSerializer,
-    NumpyJointTorquesSerializer,
-    NumpyJointCommandsSerializer,
-    NumpyCameraFrameSerializer,
-    NumpyAudioFrameSerializer,
-    NumpyImuSerializer,
-    NumpyTimestampSerializer,
-    NumpyVectorCommandSerializer,
-    NumpyStateTensorSerializer,
-    Serializer[np.ndarray],
-):
-    def __init__(
-        self: "NumpySerializer",
-        schema: P.ValueSchema,
-        *,
-        dtype: np.dtype | None = None,
-    ) -> None:
-        NumpyBaseSerializer.__init__(self, dtype=dtype)
-        Serializer.__init__(self, schema=schema)
-
-
-class NumpyMultiSerializer(MultiSerializer[np.ndarray]):
-    def __init__(self: "NumpyMultiSerializer", schema: P.IOSchema) -> None:
-        super().__init__([NumpySerializer(schema=s) for s in schema.values])