kinfer 0.3.1__cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl

kinfer/rust/src/serializer.rs

@@ -0,0 +1,221 @@
+use crate::kinfer_proto::{
+    AudioFrameSchema, AudioFrameValue, CameraFrameSchema, CameraFrameValue, ImuSchema, ImuValue,
+    JointCommandsSchema, JointCommandsValue, JointPositionUnit, JointPositionsSchema,
+    JointPositionsValue, JointTorqueUnit, JointTorquesSchema, JointTorquesValue,
+    JointVelocitiesSchema, JointVelocitiesValue, JointVelocityUnit, ProtoIO, ProtoIOSchema,
+    ProtoValue, StateTensorSchema, StateTensorValue, TimestampSchema, TimestampValue, ValueSchema,
+    VectorCommandSchema, VectorCommandValue,
+};
+
+use ort::value::Value as OrtValue;
+use std::error::Error;
+
+pub trait JointPositionsSerializer {
+    fn serialize_joint_positions(
+        &self,
+        schema: &JointPositionsSchema,
+        value: JointPositionsValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_joint_positions(
+        &self,
+        schema: &JointPositionsSchema,
+        value: OrtValue,
+    ) -> Result<JointPositionsValue, Box<dyn Error>>;
+}
+
+pub trait JointVelocitiesSerializer {
+    fn serialize_joint_velocities(
+        &self,
+        schema: &JointVelocitiesSchema,
+        value: JointVelocitiesValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_joint_velocities(
+        &self,
+        schema: &JointVelocitiesSchema,
+        value: OrtValue,
+    ) -> Result<JointVelocitiesValue, Box<dyn Error>>;
+}
+
+pub trait JointTorquesSerializer {
+    fn serialize_joint_torques(
+        &self,
+        schema: &JointTorquesSchema,
+        value: JointTorquesValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_joint_torques(
+        &self,
+        schema: &JointTorquesSchema,
+        value: OrtValue,
+    ) -> Result<JointTorquesValue, Box<dyn Error>>;
+}
+
+pub trait JointCommandsSerializer {
+    fn serialize_joint_commands(
+        &self,
+        schema: &JointCommandsSchema,
+        value: JointCommandsValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_joint_commands(
+        &self,
+        schema: &JointCommandsSchema,
+        value: OrtValue,
+    ) -> Result<JointCommandsValue, Box<dyn Error>>;
+}
+
+pub trait CameraFrameSerializer {
+    fn serialize_camera_frame(
+        &self,
+        schema: &CameraFrameSchema,
+        value: CameraFrameValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_camera_frame(
+        &self,
+        schema: &CameraFrameSchema,
+        value: OrtValue,
+    ) -> Result<CameraFrameValue, Box<dyn Error>>;
+}
+
+pub trait AudioFrameSerializer {
+    fn serialize_audio_frame(
+        &self,
+        schema: &AudioFrameSchema,
+        value: AudioFrameValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_audio_frame(
+        &self,
+        schema: &AudioFrameSchema,
+        value: OrtValue,
+    ) -> Result<AudioFrameValue, Box<dyn Error>>;
+}
+
+pub trait ImuSerializer {
+    fn serialize_imu(
+        &self,
+        schema: &ImuSchema,
+        value: ImuValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_imu(
+        &self,
+        schema: &ImuSchema,
+        value: OrtValue,
+    ) -> Result<ImuValue, Box<dyn Error>>;
+}
+
+pub trait TimestampSerializer {
+    fn serialize_timestamp(
+        &self,
+        schema: &TimestampSchema,
+        value: TimestampValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_timestamp(
+        &self,
+        schema: &TimestampSchema,
+        value: OrtValue,
+    ) -> Result<TimestampValue, Box<dyn Error>>;
+}
+
+pub trait VectorCommandSerializer {
+    fn serialize_vector_command(
+        &self,
+        schema: &VectorCommandSchema,
+        value: VectorCommandValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_vector_command(
+        &self,
+        schema: &VectorCommandSchema,
+        value: OrtValue,
+    ) -> Result<VectorCommandValue, Box<dyn Error>>;
+}
+
+pub trait StateTensorSerializer {
+    fn serialize_state_tensor(
+        &self,
+        schema: &StateTensorSchema,
+        value: StateTensorValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize_state_tensor(
+        &self,
+        schema: &StateTensorSchema,
+        value: OrtValue,
+    ) -> Result<StateTensorValue, Box<dyn Error>>;
+}
+
+pub trait Serializer:
+    JointPositionsSerializer
+    + JointVelocitiesSerializer
+    + JointTorquesSerializer
+    + JointCommandsSerializer
+    + CameraFrameSerializer
+    + AudioFrameSerializer
+    + ImuSerializer
+    + TimestampSerializer
+    + VectorCommandSerializer
+    + StateTensorSerializer
+{
+    fn serialize(
+        &self,
+        schema: &ValueSchema,
+        value: ProtoValue,
+    ) -> Result<OrtValue, Box<dyn Error>>;
+
+    fn deserialize(
+        &self,
+        schema: &ValueSchema,
+        value: OrtValue,
+    ) -> Result<ProtoValue, Box<dyn Error>>;
+}
+
+pub fn convert_position(
+    value: f32,
+    from_unit: JointPositionUnit,
+    to_unit: JointPositionUnit,
+) -> Result<f32, Box<dyn Error>> {
+    match (from_unit, to_unit) {
+        (JointPositionUnit::Radians, JointPositionUnit::Degrees) => {
+            Ok(value * 180.0 / std::f32::consts::PI)
+        }
+        (JointPositionUnit::Degrees, JointPositionUnit::Radians) => {
+            Ok(value * std::f32::consts::PI / 180.0)
+        }
+        (a, b) if a == b => Ok(value),
+        _ => Err("Unsupported position unit conversion".into()),
+    }
+}
+
+pub fn convert_velocity(
+    value: f32,
+    from_unit: JointVelocityUnit,
+    to_unit: JointVelocityUnit,
+) -> Result<f32, Box<dyn Error>> {
+    match (from_unit, to_unit) {
+        (JointVelocityUnit::RadiansPerSecond, JointVelocityUnit::DegreesPerSecond) => {
+            Ok(value * 180.0 / std::f32::consts::PI)
+        }
+        (JointVelocityUnit::DegreesPerSecond, JointVelocityUnit::RadiansPerSecond) => {
+            Ok(value * std::f32::consts::PI / 180.0)
+        }
+        (a, b) if a == b => Ok(value),
+        _ => Err("Unsupported velocity unit conversion".into()),
+    }
+}
+
+pub fn convert_torque(
+    value: f32,
+    from_unit: JointTorqueUnit,
+    to_unit: JointTorqueUnit,
+) -> Result<f32, Box<dyn Error>> {
+    match (from_unit, to_unit) {
+        (a, b) if a == b => Ok(value),
+        _ => Err("Unsupported torque unit conversion".into()),
+    }
+}

kinfer/rust/src/tests/onnx_serializer_tests.rs

@@ -0,0 +1,212 @@
+use crate::{
+    kinfer_proto::{
+        self as P, AudioFrameSchema, AudioFrameValue, CameraFrameSchema, CameraFrameValue, DType,
+        ImuAccelerometerValue, ImuGyroscopeValue, ImuMagnetometerValue, ImuSchema, ImuValue,
+        JointCommandValue, JointCommandsSchema, JointCommandsValue, JointPositionUnit,
+        JointPositionValue, JointPositionsSchema, JointPositionsValue, JointTorqueUnit,
+        JointTorqueValue, JointTorquesSchema, JointTorquesValue, JointVelocitiesSchema,
+        JointVelocitiesValue, JointVelocityUnit, JointVelocityValue, ProtoValue, StateTensorSchema,
+        StateTensorValue, TimestampSchema, TimestampValue, ValueSchema, VectorCommandSchema,
+        VectorCommandValue,
+    },
+    onnx_serializer::OnnxSerializer,
+    serializer::{
+        AudioFrameSerializer, CameraFrameSerializer, ImuSerializer, JointCommandsSerializer,
+        JointPositionsSerializer, JointTorquesSerializer, JointVelocitiesSerializer,
+        StateTensorSerializer, TimestampSerializer, VectorCommandSerializer,
+    },
+};
+
+use ndarray::Array;
+use ort::value::Value as OrtValue;
+use std::f32::consts::PI;
+
+#[test]
+fn test_serialize_joint_positions() {
+    let joint_names = vec![
+        "joint_1".to_string(),
+        "joint_2".to_string(),
+        "joint_3".to_string(),
+    ];
+    let schema = ValueSchema {
+        value_name: "test".to_string(),
+        value_type: Some(P::proto::value_schema::ValueType::JointPositions(
+            JointPositionsSchema {
+                unit: JointPositionUnit::Degrees as i32,
+                joint_names: joint_names.clone(),
+            },
+        )),
+    };
+
+    let serializer = OnnxSerializer::new(schema.clone());
+
+    // Test with matching units
+    let value = JointPositionsValue {
+        values: vec![
+            JointPositionValue {
+                joint_name: "joint_1".to_string(),
+                value: 60.0,
+                unit: JointPositionUnit::Degrees as i32,
+            },
+            JointPositionValue {
+                joint_name: "joint_2".to_string(),
+                value: 30.0,
+                unit: JointPositionUnit::Degrees as i32,
+            },
+            JointPositionValue {
+                joint_name: "joint_3".to_string(),
+                value: 90.0,
+                unit: JointPositionUnit::Degrees as i32,
+            },
+        ],
+    };
+
+    let result = match schema.value_type.as_ref().unwrap() {
+        P::proto::value_schema::ValueType::JointPositions(schema) => {
+            serializer.serialize_joint_positions(schema, value.clone())
+        }
+        _ => panic!("Wrong schema type"),
+    }
+    .unwrap();
+
+    // Verify tensor shape and values
+    let tensor = result.try_extract_tensor::<f32>().unwrap();
+    let array = tensor.view();
+    assert_eq!(array.shape(), &[3]);
+    assert_eq!(array[[0]], 60.0); // joint_1
+    assert_eq!(array[[1]], 30.0); // joint_2
+    assert_eq!(array[[2]], 90.0); // joint_3
+
+    let deserialized = match schema.value_type.as_ref().unwrap() {
+        P::proto::value_schema::ValueType::JointPositions(schema) => {
+            serializer.deserialize_joint_positions(schema, result)
+        }
+        _ => panic!("Wrong schema type"),
+    }
+    .unwrap();
+
+    // Verify full deserialization
+    assert_eq!(deserialized.values.len(), value.values.len());
+    for (expected, actual) in value.values.iter().zip(deserialized.values.iter()) {
+        assert_eq!(expected.joint_name, actual.joint_name);
+        assert_eq!(expected.value, actual.value);
+        assert_eq!(expected.unit, actual.unit);
+    }
+
+    // Test unit conversion
+    let value_radians = JointPositionsValue {
+        values: vec![JointPositionValue {
+            joint_name: "joint_1".to_string(),
+            value: PI / 6.0,
+            unit: JointPositionUnit::Radians as i32,
+        }],
+    };
+
+    let schema_radians = ValueSchema {
+        value_name: "test".to_string(),
+        value_type: Some(P::proto::value_schema::ValueType::JointPositions(
+            JointPositionsSchema {
+                unit: JointPositionUnit::Radians as i32,
+                joint_names: vec!["joint_1".to_string()],
+            },
+        )),
+    };
+
+    let serializer = OnnxSerializer::new(schema_radians.clone());
+    let result = match schema_radians.value_type.as_ref().unwrap() {
+        P::proto::value_schema::ValueType::JointPositions(schema) => {
+            serializer.serialize_joint_positions(schema, value_radians.clone())
+        }
+        _ => panic!("Wrong schema type"),
+    }
+    .unwrap();
+
+    let tensor = result.try_extract_tensor::<f32>().unwrap();
+    let array = tensor.view();
+    assert!((array[[0]] - PI / 6.0).abs() < 1e-6);
+}
+
+#[test]
+fn test_serialize_joint_positions_errors() {
+    let schema = ValueSchema {
+        value_name: "test".to_string(),
+        value_type: Some(P::proto::value_schema::ValueType::JointPositions(
+            JointPositionsSchema {
+                unit: JointPositionUnit::Degrees as i32,
+                joint_names: vec!["joint_1".to_string(), "joint_2".to_string()],
+            },
+        )),
+    };
+
+    let serializer = OnnxSerializer::new(schema.clone());
+
+    // Test cases that should fail:
+
+    // Case 1: Wrong number of joints
+    let value_wrong_count = JointPositionsValue {
+        values: vec![JointPositionValue {
+            joint_name: "joint_1".to_string(),
+            value: 60.0,
+            unit: JointPositionUnit::Degrees as i32,
+        }],
+    };
+
+    let result = match schema.value_type.as_ref().unwrap() {
+        P::proto::value_schema::ValueType::JointPositions(schema) => {
+            serializer.serialize_joint_positions(schema, value_wrong_count)
+        }
+        _ => panic!("Wrong schema type"),
+    };
+    assert!(
+        result.is_err(),
+        "Should fail when joint count doesn't match"
+    );
+
+    // Case 2: Wrong joint names
+    let value_wrong_names = JointPositionsValue {
+        values: vec![
+            JointPositionValue {
+                joint_name: "wrong_joint_1".to_string(),
+                value: 60.0,
+                unit: JointPositionUnit::Degrees as i32,
+            },
+            JointPositionValue {
+                joint_name: "wrong_joint_2".to_string(),
+                value: 30.0,
+                unit: JointPositionUnit::Degrees as i32,
+            },
+        ],
+    };
+
+    let result = match schema.value_type.as_ref().unwrap() {
+        P::proto::value_schema::ValueType::JointPositions(schema) => {
+            serializer.serialize_joint_positions(schema, value_wrong_names)
+        }
+        _ => panic!("Wrong schema type"),
+    };
+    assert!(result.is_err(), "Should fail when joint names don't match");
+
+    // Case 3: Wrong unit type
+    let value_wrong_unit = JointPositionsValue {
+        values: vec![
+            JointPositionValue {
+                joint_name: "joint_1".to_string(),
+                value: 60.0,
+                unit: JointPositionUnit::Radians as i32,
+            },
+            JointPositionValue {
+                joint_name: "joint_2".to_string(),
+                value: 30.0,
+                unit: JointPositionUnit::Radians as i32,
+            },
+        ],
+    };
+
+    let result = match schema.value_type.as_ref().unwrap() {
+        P::proto::value_schema::ValueType::JointPositions(schema) => {
+            serializer.serialize_joint_positions(schema, value_wrong_unit)
+        }
+        _ => panic!("Wrong schema type"),
+    };
+    assert!(result.is_err(), "Should fail when units don't match");
+}

kinfer/rust_bindings/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+
+name = "rust_bindings"
+version.workspace = true
+edition.workspace = true
+description.workspace = true
+authors.workspace = true
+repository.workspace = true
+license.workspace = true
+readme.workspace = true
+
+[lib]
+
+name = "rust_bindings"
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+pyo3 = { version = ">= 0.21.0", features = ["extension-module"] }
+pyo3-stub-gen = ">= 0.6.0"

kinfer/rust_bindings/pyproject.toml

@@ -0,0 +1,7 @@
+[build-system]
+requires = ["maturin>=1.1,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "rust_bindings"
+requires-python = ">=3.11"

kinfer/rust_bindings/src/bin/stub_gen.rs

@@ -0,0 +1,7 @@
+use pyo3_stub_gen::Result;
+
+fn main() -> Result<()> {
+    let stub = rust_bindings::stub_info()?;
+    stub.generate()?;
+    Ok(())
+}

kinfer/rust_bindings/src/lib.rs

@@ -0,0 +1,17 @@
+use pyo3::prelude::*;
+use pyo3_stub_gen::define_stub_info_gatherer;
+use pyo3_stub_gen::derive::gen_stub_pyfunction;
+
+#[pyfunction]
+#[gen_stub_pyfunction]
+fn get_version() -> String {
+    env!("CARGO_PKG_VERSION").to_string()
+}
+
+#[pymodule]
+fn rust_bindings(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(get_version, m)?)?;
+    Ok(())
+}
+
+define_stub_info_gatherer!(stub_info);

kinfer/rust_bindings.pyi

@@ -0,0 +1,7 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401
+
+
+def get_version() -> str:
+    ...
+

kinfer/serialize/__init__.py

@@ -0,0 +1,36 @@
+"""Defines an interface for instantiating serializers."""
+
+from typing import Literal
+
+from kinfer import proto as P
+
+from .base import MultiSerializer, Serializer
+from .json import JsonMultiSerializer, JsonSerializer
+from .numpy import NumpyMultiSerializer, NumpySerializer
+from .pytorch import PyTorchMultiSerializer, PyTorchSerializer
+
+SerializerType = Literal["json", "numpy", "pytorch"]
+
+
+def get_serializer(schema: P.ValueSchema, serializer_type: SerializerType) -> Serializer:
+    match serializer_type:
+        case "json":
+            return JsonSerializer(schema=schema)
+        case "numpy":
+            return NumpySerializer(schema=schema)
+        case "pytorch":
+            return PyTorchSerializer(schema=schema)
+        case _:
+            raise ValueError(f"Unsupported serializer type: {serializer_type}")
+
+
+def get_multi_serializer(schema: P.IOSchema, serializer_type: SerializerType) -> MultiSerializer:
+    match serializer_type:
+        case "json":
+            return JsonMultiSerializer(schema=schema)
+        case "numpy":
+            return NumpyMultiSerializer(schema=schema)
+        case "pytorch":
+            return PyTorchMultiSerializer(schema=schema)
+        case _:
+            raise ValueError(f"Unsupported serializer type: {serializer_type}")