kinfer 0.3.3__cp312-cp312-macosx_11_0_arm64.whl → 0.4.1__cp312-cp312-macosx_11_0_arm64.whl

kinfer/serialize/pytorch.py

@@ -1,402 +0,0 @@
-"""Defines a serializer for PyTorch tensors."""
-
-from typing import cast
-
-import numpy as np
-import torch
-from torch import Tensor
-
-from kinfer import proto as K
-from kinfer.serialize.base import (
-    AudioFrameSerializer,
-    CameraFrameSerializer,
-    ImuSerializer,
-    JointCommandsSerializer,
-    JointPositionsSerializer,
-    JointTorquesSerializer,
-    JointVelocitiesSerializer,
-    MultiSerializer,
-    Serializer,
-    StateTensorSerializer,
-    TimestampSerializer,
-    VectorCommandSerializer,
-)
-from kinfer.serialize.utils import (
-    check_names_match,
-    convert_angular_position,
-    convert_angular_velocity,
-    convert_torque,
-    dtype_num_bytes,
-    dtype_range,
-    numpy_dtype,
-    parse_bytes,
-    pytorch_dtype,
-)
-
-
-class PyTorchBaseSerializer:
-    def __init__(
-        self: "PyTorchBaseSerializer",
-        device: str | torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> None:
-        self.device = device
-        self.dtype = dtype
-
-
-class PyTorchJointPositionsSerializer(PyTorchBaseSerializer, JointPositionsSerializer[Tensor]):
-    def serialize_joint_positions(
-        self: "PyTorchJointPositionsSerializer",
-        schema: K.JointPositionsSchema,
-        value: K.JointPositionsValue,
-    ) -> Tensor:
-        value_map = {v.joint_name: v for v in value.values}
-        check_names_match("schema", schema.joint_names, "value", list(value_map.keys()))
-        tensor = torch.tensor(
-            [
-                convert_angular_position(value_map[name].value, value_map[name].unit, schema.unit)
-                for name in schema.joint_names
-            ],
-            dtype=self.dtype,
-            device=self.device,
-        )
-        return tensor
-
-    def deserialize_joint_positions(
-        self: "PyTorchJointPositionsSerializer",
-        schema: K.JointPositionsSchema,
-        value: Tensor,
-    ) -> K.JointPositionsValue:
-        if value.shape != (len(schema.joint_names),):
-            raise ValueError(
-                f"Shape of tensor must match number of joint names: {value.shape} != {len(schema.joint_names)}"
-            )
-        value_list = cast(list[float], value.detach().cpu().numpy().astype(float).tolist())
-        return K.JointPositionsValue(
-            values=[
-                K.JointPositionValue(joint_name=name, value=value_list[i], unit=schema.unit)
-                for i, name in enumerate(schema.joint_names)
-            ]
-        )
-
-
-class PyTorchJointVelocitiesSerializer(PyTorchBaseSerializer, JointVelocitiesSerializer[Tensor]):
-    def serialize_joint_velocities(
-        self: "PyTorchJointVelocitiesSerializer",
-        schema: K.JointVelocitiesSchema,
-        value: K.JointVelocitiesValue,
-    ) -> Tensor:
-        value_map = {v.joint_name: v for v in value.values}
-        check_names_match("schema", schema.joint_names, "value", list(value_map.keys()))
-        tensor = torch.tensor(
-            [
-                convert_angular_velocity(value_map[name].value, value_map[name].unit, schema.unit)
-                for name in schema.joint_names
-            ],
-            dtype=self.dtype,
-            device=self.device,
-        )
-        return tensor
-
-    def deserialize_joint_velocities(
-        self: "PyTorchJointVelocitiesSerializer",
-        schema: K.JointVelocitiesSchema,
-        value: Tensor,
-    ) -> K.JointVelocitiesValue:
-        if value.shape != (len(schema.joint_names),):
-            raise ValueError(
-                f"Shape of tensor must match number of joint names: {value.shape} != {len(schema.joint_names)}"
-            )
-        value_list = cast(list[float], value.detach().cpu().numpy().astype(float).tolist())
-        return K.JointVelocitiesValue(
-            values=[
-                K.JointVelocityValue(joint_name=name, value=value_list[i], unit=schema.unit)
-                for i, name in enumerate(schema.joint_names)
-            ]
-        )
-
-
-class PyTorchJointTorquesSerializer(PyTorchBaseSerializer, JointTorquesSerializer[Tensor]):
-    def serialize_joint_torques(
-        self: "PyTorchJointTorquesSerializer",
-        schema: K.JointTorquesSchema,
-        value: K.JointTorquesValue,
-    ) -> Tensor:
-        value_map = {v.joint_name: v for v in value.values}
-        check_names_match("schema", schema.joint_names, "value", list(value_map.keys()))
-        tensor = torch.tensor(
-            [convert_torque(value_map[name].value, value_map[name].unit, schema.unit) for name in schema.joint_names],
-            dtype=self.dtype,
-            device=self.device,
-        )
-        return tensor
-
-    def deserialize_joint_torques(
-        self: "PyTorchJointTorquesSerializer",
-        schema: K.JointTorquesSchema,
-        value: Tensor,
-    ) -> K.JointTorquesValue:
-        if value.shape != (len(schema.joint_names),):
-            raise ValueError(
-                f"Shape of tensor must match number of joint names: {value.shape} != {len(schema.joint_names)}"
-            )
-        value_list = cast(list[float], value.detach().cpu().numpy().astype(float).tolist())
-        return K.JointTorquesValue(
-            values=[
-                K.JointTorqueValue(joint_name=name, value=value_list[i], unit=schema.unit)
-                for i, name in enumerate(schema.joint_names)
-            ]
-        )
-
-
-class PyTorchJointCommandsSerializer(PyTorchBaseSerializer, JointCommandsSerializer[Tensor]):
-    def _convert_value_to_tensor(
-        self: "PyTorchJointCommandsSerializer",
-        value: K.JointCommandValue,
-        schema: K.JointCommandsSchema,
-    ) -> Tensor:
-        return torch.tensor(
-            [
-                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,
-            device=self.device,
-        )
-
-    def _convert_tensor_to_value(
-        self: "PyTorchJointCommandsSerializer",
-        values: list[float],
-        schema: K.JointCommandsSchema,
-        name: str,
-    ) -> K.JointCommandValue:
-        if len(values) != 5:
-            raise ValueError(f"Shape of tensor must match number of joint commands: {len(values)} != 5")
-        return K.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: "PyTorchJointCommandsSerializer",
-        schema: K.JointCommandsSchema,
-        value: K.JointCommandsValue,
-    ) -> Tensor:
-        value_map = {v.joint_name: v for v in value.values}
-        check_names_match("schema", schema.joint_names, "value", list(value_map.keys()))
-        tensor = torch.stack(
-            [self._convert_value_to_tensor(value_map[name], schema) for name in schema.joint_names],
-            dim=0,
-        )
-        return tensor
-
-    def deserialize_joint_commands(
-        self: "PyTorchJointCommandsSerializer",
-        schema: K.JointCommandsSchema,
-        value: Tensor,
-    ) -> K.JointCommandsValue:
-        if value.shape != (len(schema.joint_names), 5):
-            raise ValueError(
-                "Shape of tensor must match number of joint names and commands: "
-                f"{value.shape} != ({len(schema.joint_names)}, 5)"
-            )
-        value_list = cast(list[list[float]], value.detach().cpu().numpy().astype(float).tolist())
-        return K.JointCommandsValue(
-            values=[
-                self._convert_tensor_to_value(value_list[i], schema, name) for i, name in enumerate(schema.joint_names)
-            ]
-        )
-
-
-class PyTorchCameraFrameSerializer(PyTorchBaseSerializer, CameraFrameSerializer[Tensor]):
-    def serialize_camera_frame(
-        self: "PyTorchCameraFrameSerializer", schema: K.CameraFrameSchema, value: K.CameraFrameValue
-    ) -> Tensor:
-        np_arr = parse_bytes(value.data, K.DType.UINT8)
-        tensor = torch.from_numpy(np_arr).to(self.device, self.dtype) / 255.0
-        if tensor.numel() != schema.channels * schema.height * schema.width:
-            raise ValueError(
-                "Length of data must match number of channels, height, and width: "
-                f"{tensor.numel()} != {schema.channels} * {schema.height} * {schema.width}"
-            )
-        tensor = tensor.view(schema.channels, schema.height, schema.width)
-        return tensor
-
-    def deserialize_camera_frame(
-        self: "PyTorchCameraFrameSerializer", schema: K.CameraFrameSchema, value: Tensor
-    ) -> K.CameraFrameValue:
-        np_arr = (value * 255.0).detach().cpu().flatten().numpy().astype(np.uint8)
-        return K.CameraFrameValue(data=np_arr.tobytes())
-
-
-class PyTorchAudioFrameSerializer(PyTorchBaseSerializer, AudioFrameSerializer[Tensor]):
-    def serialize_audio_frame(
-        self: "PyTorchAudioFrameSerializer", schema: K.AudioFrameSchema, value: K.AudioFrameValue
-    ) -> Tensor:
-        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)
-        tensor = torch.from_numpy(np_arr).to(self.device, self.dtype)
-        tensor = tensor.view(schema.channels, -1)
-        tensor = tensor / max_value
-        return tensor
-
-    def deserialize_audio_frame(
-        self: "PyTorchAudioFrameSerializer", schema: K.AudioFrameSchema, value: Tensor
-    ) -> K.AudioFrameValue:
-        _, max_value = dtype_range(schema.dtype)
-        np_arr = (value * max_value).detach().cpu().flatten().numpy().astype(numpy_dtype(schema.dtype))
-        return K.AudioFrameValue(data=np_arr.tobytes())
-
-
-class PyTorchImuSerializer(PyTorchBaseSerializer, ImuSerializer[Tensor]):
-    def serialize_imu(self: "PyTorchImuSerializer", schema: K.ImuSchema, value: K.ImuValue) -> Tensor:
-        vectors: list[Tensor] = []
-        if schema.use_accelerometer:
-            vectors.append(
-                torch.tensor(
-                    [value.linear_acceleration.x, value.linear_acceleration.y, value.linear_acceleration.z],
-                    dtype=self.dtype,
-                    device=self.device,
-                )
-            )
-        if schema.use_gyroscope:
-            vectors.append(
-                torch.tensor(
-                    [value.angular_velocity.x, value.angular_velocity.y, value.angular_velocity.z],
-                    dtype=self.dtype,
-                    device=self.device,
-                )
-            )
-        if schema.use_magnetometer:
-            vectors.append(
-                torch.tensor(
-                    [value.magnetic_field.x, value.magnetic_field.y, value.magnetic_field.z],
-                    dtype=self.dtype,
-                    device=self.device,
-                )
-            )
-        if not vectors:
-            raise ValueError("IMU has nothing to serialize")
-        return torch.stack(vectors, dim=0)
-
-    def deserialize_imu(self: "PyTorchImuSerializer", schema: K.ImuSchema, value: Tensor) -> K.ImuValue:
-        vectors = value.tolist()
-        imu_value = K.ImuValue()
-        if schema.use_accelerometer:
-            (x, y, z), vectors = vectors[0], vectors[1:]
-            imu_value.linear_acceleration.x = x
-            imu_value.linear_acceleration.y = y
-            imu_value.linear_acceleration.z = z
-        if schema.use_gyroscope:
-            (x, y, z), vectors = vectors[0], vectors[1:]
-            imu_value.angular_velocity.x = x
-            imu_value.angular_velocity.y = y
-            imu_value.angular_velocity.z = z
-        if schema.use_magnetometer:
-            (x, y, z), vectors = vectors[0], vectors[1:]
-            imu_value.magnetic_field.x = x
-            imu_value.magnetic_field.y = y
-            imu_value.magnetic_field.z = z
-        return imu_value
-
-
-class PyTorchTimestampSerializer(PyTorchBaseSerializer, TimestampSerializer[Tensor]):
-    def serialize_timestamp(
-        self: "PyTorchTimestampSerializer", schema: K.TimestampSchema, value: K.TimestampValue
-    ) -> Tensor:
-        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 torch.tensor([total_elapsed_seconds], dtype=self.dtype, device=self.device, requires_grad=False)
-
-    def deserialize_timestamp(
-        self: "PyTorchTimestampSerializer", schema: K.TimestampSchema, value: Tensor
-    ) -> K.TimestampValue:
-        total_elapsed_seconds = value.item()
-        elapsed_seconds = int(total_elapsed_seconds)
-        elapsed_nanos = int((total_elapsed_seconds - elapsed_seconds) * 1_000_000_000)
-        return K.TimestampValue(seconds=elapsed_seconds, nanos=elapsed_nanos)
-
-
-class PyTorchVectorCommandSerializer(PyTorchBaseSerializer, VectorCommandSerializer[Tensor]):
-    def serialize_vector_command(
-        self: "PyTorchVectorCommandSerializer", schema: K.VectorCommandSchema, value: K.VectorCommandValue
-    ) -> Tensor:
-        return torch.tensor(value.values, dtype=self.dtype, device=self.device)
-
-    def deserialize_vector_command(
-        self: "PyTorchVectorCommandSerializer", schema: K.VectorCommandSchema, value: Tensor
-    ) -> K.VectorCommandValue:
-        if value.shape != (schema.dimensions,):
-            raise ValueError(f"Shape of tensor must match number of dimensions: {value.shape} != {schema.dimensions}")
-        values = cast(list[float], value.tolist())
-        return K.VectorCommandValue(values=values)
-
-
-class PyTorchStateTensorSerializer(PyTorchBaseSerializer, StateTensorSerializer[Tensor]):
-    def serialize_state_tensor(
-        self: "PyTorchStateTensorSerializer", schema: K.StateTensorSchema, value: K.StateTensorValue
-    ) -> Tensor:
-        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)
-        tensor = torch.from_numpy(np_arr).to(self.device, pytorch_dtype(schema.dtype))
-        tensor = tensor.view(tuple(schema.shape))
-        return tensor
-
-    def deserialize_state_tensor(
-        self: "PyTorchStateTensorSerializer", schema: K.StateTensorSchema, value: Tensor
-    ) -> K.StateTensorValue:
-        return K.StateTensorValue(data=value.detach().cpu().flatten().numpy().tobytes())
-
-
-class PyTorchSerializer(
-    PyTorchJointPositionsSerializer,
-    PyTorchJointVelocitiesSerializer,
-    PyTorchJointTorquesSerializer,
-    PyTorchJointCommandsSerializer,
-    PyTorchCameraFrameSerializer,
-    PyTorchAudioFrameSerializer,
-    PyTorchImuSerializer,
-    PyTorchTimestampSerializer,
-    PyTorchVectorCommandSerializer,
-    PyTorchStateTensorSerializer,
-    Serializer[Tensor],
-):
-    def __init__(
-        self: "PyTorchSerializer",
-        schema: K.ValueSchema,
-        *,
-        device: str | torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> None:
-        PyTorchBaseSerializer.__init__(self, device=device, dtype=dtype)
-        Serializer.__init__(self, schema=schema)
-
-
-class PyTorchMultiSerializer(MultiSerializer[Tensor]):
-    def __init__(self: "PyTorchMultiSerializer", schema: K.IOSchema) -> None:
-        super().__init__([PyTorchSerializer(schema=s) for s in schema.values])

kinfer/serialize/schema.py

@@ -1,125 +0,0 @@
-"""Defines utility functions for the schema."""
-
-import numpy as np
-
-from kinfer import proto as K
-from kinfer.serialize.utils import dtype_num_bytes
-
-
-def get_dummy_value(value_schema: K.ValueSchema) -> K.Value:
-    value_type = value_schema.WhichOneof("value_type")
-
-    match value_type:
-        case "joint_positions":
-            return K.Value(
-                joint_positions=K.JointPositionsValue(
-                    values=[
-                        K.JointPositionValue(
-                            joint_name=joint_name,
-                            value=0.0,
-                            unit=value_schema.joint_positions.unit,
-                        )
-                        for joint_name in value_schema.joint_positions.joint_names
-                    ]
-                ),
-            )
-        case "joint_velocities":
-            return K.Value(
-                joint_velocities=K.JointVelocitiesValue(
-                    values=[
-                        K.JointVelocityValue(
-                            joint_name=joint_name,
-                            value=0.0,
-                            unit=value_schema.joint_velocities.unit,
-                        )
-                        for joint_name in value_schema.joint_velocities.joint_names
-                    ]
-                ),
-            )
-        case "joint_torques":
-            return K.Value(
-                joint_torques=K.JointTorquesValue(
-                    values=[
-                        K.JointTorqueValue(
-                            joint_name=joint_name,
-                            value=0.0,
-                            unit=value_schema.joint_torques.unit,
-                        )
-                        for joint_name in value_schema.joint_torques.joint_names
-                    ]
-                ),
-            )
-        case "joint_commands":
-            return K.Value(
-                joint_commands=K.JointCommandsValue(
-                    values=[
-                        K.JointCommandValue(
-                            joint_name=joint_name,
-                            torque=0.0,
-                            velocity=0.0,
-                            position=0.0,
-                            kp=0.0,
-                            kd=0.0,
-                            torque_unit=value_schema.joint_commands.torque_unit,
-                            velocity_unit=value_schema.joint_commands.velocity_unit,
-                            position_unit=value_schema.joint_commands.position_unit,
-                        )
-                        for joint_name in value_schema.joint_commands.joint_names
-                    ]
-                ),
-            )
-        case "camera_frame":
-            return K.Value(
-                camera_frame=K.CameraFrameValue(
-                    data=b"\x00"
-                    * (
-                        value_schema.camera_frame.width
-                        * value_schema.camera_frame.height
-                        * value_schema.camera_frame.channels
-                    )
-                ),
-            )
-        case "audio_frame":
-            return K.Value(
-                audio_frame=K.AudioFrameValue(
-                    data=b"\x00"
-                    * (
-                        value_schema.audio_frame.channels
-                        * value_schema.audio_frame.sample_rate
-                        * dtype_num_bytes(value_schema.audio_frame.dtype)
-                    )
-                ),
-            )
-        case "imu":
-            return K.Value(
-                imu=K.ImuValue(
-                    linear_acceleration=K.ImuAccelerometerValue(x=0.0, y=0.0, z=0.0),
-                    angular_velocity=K.ImuGyroscopeValue(x=0.0, y=0.0, z=0.0),
-                    magnetic_field=K.ImuMagnetometerValue(x=0.0, y=0.0, z=0.0),
-                ),
-            )
-        case "timestamp":
-            return K.Value(
-                timestamp=K.TimestampValue(seconds=1728000000, nanos=0),
-            )
-        case "vector_command":
-            return K.Value(
-                vector_command=K.VectorCommandValue(values=[0.0] * value_schema.vector_command.dimensions),
-            )
-        case "state_tensor":
-            return K.Value(
-                state_tensor=K.StateTensorValue(
-                    data=b"\x00"
-                    * np.prod(value_schema.state_tensor.shape)
-                    * dtype_num_bytes(value_schema.state_tensor.dtype)
-                ),
-            )
-        case _:
-            raise ValueError(f"Invalid value type: {value_type}")
-
-
-def get_dummy_io(schema: K.IOSchema) -> K.IO:
-    io_value = K.IO()
-    for value_schema in schema.values:
-        io_value.values.append(get_dummy_value(value_schema))
-    return io_value

kinfer/serialize/types.py

@@ -1,17 +0,0 @@
-"""Type conversion utilities for serializers."""
-
-from typing import Type, TypeVar, cast
-
-from kinfer import proto as K
-
-T = TypeVar("T")
-
-
-def to_value_type(enum_value: T) -> K.Value:
-    """Convert an enum value to ValueType."""
-    return cast(K.Value, enum_value)
-
-
-def from_value_type(value_type: K.Value, enum_class: Type[T]) -> T:
-    """Convert a ValueType to the specified enum type."""
-    return cast(T, value_type)