kinfer 0.4.3__cp311-cp311-macosx_11_0_arm64.whl → 0.5.2__cp311-cp311-macosx_11_0_arm64.whl

kinfer/rust_bindings.pyi

@@ -8,14 +8,21 @@ import typing
 
 class ModelProviderABC:
     def __new__(cls) -> ModelProviderABC: ...
-    def get_joint_angles(self, joint_names:typing.Sequence[builtins.str]) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_joint_angular_velocities(self, joint_names:typing.Sequence[builtins.str]) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_projected_gravity(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_accelerometer(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_gyroscope(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_command(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def get_time(self) -> numpy.typing.NDArray[numpy.float32]: ...
-    def take_action(self, joint_names:typing.Sequence[builtins.str], action:numpy.typing.NDArray[numpy.float32]) -> None: ...
+    def get_inputs(self, input_types:typing.Sequence[builtins.str], metadata:PyModelMetadata) -> builtins.dict[builtins.str, numpy.typing.NDArray[numpy.float32]]: ...
+    def take_action(self, action:numpy.typing.NDArray[numpy.float32], metadata:PyModelMetadata) -> None: ...
+
+class PyInputType:
+    def __new__(cls, input_type:builtins.str) -> PyInputType: ...
+    def get_name(self) -> builtins.str: ...
+    def get_shape(self, metadata:PyModelMetadata) -> builtins.list[builtins.int]: ...
+    def __repr__(self) -> builtins.str: ...
+    def __eq__(self, other:typing.Any) -> builtins.bool: ...
+
+class PyModelMetadata:
+    def __new__(self, joint_names:typing.Sequence[builtins.str], num_commands:typing.Optional[builtins.int], carry_size:typing.Sequence[builtins.int]) -> PyModelMetadata: ...
+    def to_json(self) -> builtins.str: ...
+    def __repr__(self) -> builtins.str: ...
+    def __eq__(self, other:typing.Any) -> builtins.bool: ...
 
 class PyModelProvider:
     ...
@@ -35,3 +42,5 @@ class PyModelRuntime:
 
 def get_version() -> builtins.str: ...
 
+def metadata_from_json(json:builtins.str) -> PyModelMetadata: ...
+

kinfer/scripts/plot_ndjson.py

@@ -0,0 +1,177 @@
+"""Plot NDJSON logs saved by kinfer."""
+
+import argparse
+import json
+import logging
+from pathlib import Path
+from typing import Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Set up logger
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+
+
+def read_ndjson(filepath: str) -> list[dict]:
+    """Read NDJSON file and return list of parsed objects."""
+    data = []
+    with open(filepath, "r") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                data.append(json.loads(line))
+    return data
+
+
+def skip_initial_data(data: list[dict], skip_seconds: float) -> list[dict]:
+    """Skip the first n seconds of data based on timestamps."""
+    if not data or skip_seconds <= 0.0:
+        return data
+
+    # Extract timestamps and convert to seconds relative to first timestamp
+    timestamps = [d["t_us"] for d in data]
+    t_start = timestamps[0]
+    times = [(t - t_start) / 1e6 for t in timestamps]  # Convert to seconds
+
+    # Find indices where time >= skip_seconds
+    skip_indices = [i for i, t in enumerate(times) if t >= skip_seconds]
+    if not skip_indices:
+        logger.info("All data points are within the skip period (%.2f seconds). No data to plot.", skip_seconds)
+        return []
+
+    # Filter data
+    start_idx = skip_indices[0]
+    filtered_data = data[start_idx:]
+    logger.info("Skipped first %.2f seconds (%d data points)", skip_seconds, start_idx)
+
+    return filtered_data
+
+
+def plot_data(data: list[dict], save_path: Optional[Union[str, Path]] = None) -> None:
+    """Plot all data fields from the NDJSON."""
+    if not data:
+        logger.info("No data to plot")
+        return
+
+    # Extract timestamps and convert to seconds relative to first timestamp
+    timestamps = [d["t_us"] for d in data]
+    t_start = timestamps[0]
+    times = [(t - t_start) / 1e6 for t in timestamps]  # Convert to seconds
+
+    # Extract data arrays
+    joint_angles = np.array([d["joint_angles"] for d in data if d["joint_angles"] is not None])
+    joint_vels = np.array([d["joint_vels"] for d in data if d["joint_vels"] is not None])
+    projected_g = np.array([d["projected_g"] for d in data if d["projected_g"] is not None])
+    accel = np.array([d["accel"] for d in data if d["accel"] is not None])
+    command = np.array([d["command"] for d in data if d["command"] is not None])
+    output = np.array([d["output"] for d in data if d["output"] is not None])
+
+    # Create subplots
+    fig, axes = plt.subplots(3, 2, figsize=(15, 12))
+    fig.suptitle("Robot Data Over Time", fontsize=16)
+
+    if len(joint_angles) > 0:
+        ax = axes[0, 0]
+        for i in range(joint_angles.shape[1]):
+            ax.plot(times[: len(joint_angles)], joint_angles[:, i], alpha=0.7, linewidth=0.8)
+        ax.set_title("Joint Angles")
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Angle (rad)")
+        ax.grid(True, alpha=0.3)
+
+    if len(joint_vels) > 0:
+        ax = axes[0, 1]
+        for i in range(joint_vels.shape[1]):
+            ax.plot(times[: len(joint_vels)], joint_vels[:, i], alpha=0.7, linewidth=0.8)
+        ax.set_title("Joint Velocities")
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Velocity (rad/s)")
+        ax.grid(True, alpha=0.3)
+
+    if len(projected_g) > 0:
+        ax = axes[1, 0]
+        labels = ["X", "Y", "Z"]
+        for i in range(projected_g.shape[1]):
+            ax.plot(times[: len(projected_g)], projected_g[:, i], label=labels[i], linewidth=1.5)
+        ax.set_title("Projected Gravity")
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Acceleration (m/s²)")
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+
+    if len(accel) > 0:
+        ax = axes[1, 1]
+        labels = ["X", "Y", "Z"]
+        for i in range(accel.shape[1]):
+            ax.plot(times[: len(accel)], accel[:, i], label=labels[i], linewidth=1.5)
+        ax.set_title("Acceleration")
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Acceleration (m/s²)")
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+
+    if len(command) > 0:
+        ax = axes[2, 0]
+        for i in range(command.shape[1]):
+            ax.plot(times[: len(command)], command[:, i], label=f"Cmd {i}", linewidth=1.2)
+        ax.set_title("Command")
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Command Value")
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+
+    if len(output) > 0:
+        ax = axes[2, 1]
+        for i in range(output.shape[1]):
+            ax.plot(times[: len(output)], output[:, i], alpha=0.7, linewidth=0.8)
+        ax.set_title("Output")
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Output Value")
+        ax.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches="tight")
+        logger.info("Plot saved to: %s", save_path)
+        plt.close()
+    else:
+        plt.show()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Plot NDJSON logs saved by kinfer")
+    parser.add_argument("filepath", help="Path to the NDJSON file to plot")
+    parser.add_argument("--skip", type=float, default=0.0, help="Skip the first n seconds of data")
+    parser.add_argument("--save", action="store_true", help="Save the plot to a PNG file in a plots folder")
+    args = parser.parse_args()
+
+    filepath = args.filepath
+    if not Path(filepath).exists():
+        logger.info("File not found: %s", filepath)
+        return
+
+    logger.info("Reading data from %s...", filepath)
+    data = read_ndjson(filepath)
+    logger.info("Loaded %d data points", len(data))
+
+    filtered_data = skip_initial_data(data, args.skip)
+
+    save_path = None
+    if args.save:
+        # Create save path in plots folder with same name but .png extension
+        input_path = Path(filepath)
+        plots_dir = input_path.parent / "plots"
+        plots_dir.mkdir(exist_ok=True)
+
+        # Change extension from .ndjson to .png
+        filename = input_path.stem + ".png"
+        save_path = str(plots_dir / filename)
+
+    plot_data(filtered_data, save_path)
+
+
+if __name__ == "__main__":
+    main()

{kinfer-0.4.3.dist-info → kinfer-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kinfer
-Version: 0.4.3
+Version: 0.5.2
 Summary: Tool to make it easier to run a model on a real robot
 Home-page: https://github.com/kscalelabs/kinfer.git
 Author: K-Scale Labs
@@ -10,6 +10,9 @@ License-File: LICENSE
 Requires-Dist: onnx
 Requires-Dist: onnxruntime==1.20.0
 Requires-Dist: pydantic
+Requires-Dist: matplotlib
+Requires-Dist: numpy
+Requires-Dist: pathlib
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: darglint; extra == "dev"

kinfer-0.5.2.dist-info/RECORD

@@ -0,0 +1,26 @@
+kinfer/rust_bindings.cpython-311-darwin.so,sha256=BuoTMXGqu9a6LNcthv_SxZ1BWZvkU_pIAliqHTSK0XM,2090064
+kinfer/requirements.txt,sha256=e8RzCNVZGJ9Jb874Eom0psUCOqp0wGf8LJZ86ysto0w,127
+kinfer/__init__.py,sha256=i5da6ND827Cgn8PFKzDCmEBk14ptQasLa_9fdof4Y9c,398
+kinfer/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kinfer/rust_bindings.pyi,sha256=lPGHMDsydr4Ue-CNHRjQtJWmq5sE5daVN8PbxJgOZY4,2048
+kinfer/rust/Cargo.toml,sha256=NoWaLluRJHV6mOCeVXk8q8Pb4028Z3ZPqLspBMHAPTQ,703
+kinfer/rust/src/runtime.rs,sha256=7BRZmyMoV6FvYs2_XnLk0DFrxt7qgz35loCUJRSsuxM,3448
+kinfer/rust/src/logger.rs,sha256=iDb8hiRjxnH8l-Qzth5XcLk2ICnnX1JlBxyOzaZ4t64,3919
+kinfer/rust/src/types.rs,sha256=KNIM6sHvX1gnQZc74uRTgmzbgS4FZjxll5jcuNGIDc8,2768
+kinfer/rust/src/lib.rs,sha256=25LuAdONp9c3lI0_g2kAf73yrtzbYc9GtJjm1zl9IwU,120
+kinfer/rust/src/model.rs,sha256=RQR0y7u9lPoSBNPnBWYesfxDkBSBBJpkyiH6qqCeQRI,11504
+kinfer/scripts/plot_ndjson.py,sha256=eg-wwuDLRDcuVdr1Jc5W4yeJG5Bf--UjaK6kk9XoC4w,6179
+kinfer/export/serialize.py,sha256=QjsvTn7cIRroZXXONXQUrEd1tR69oLrszojuKoK-Xuk,3208
+kinfer/export/__init__.py,sha256=6a8cClA2H3AAjAhY6ucmxEnfCRZ30FDn2L7SArCu7nY,56
+kinfer/export/jax.py,sha256=zB3HJ_Ao-UtBKpFddywHUQstBpvc5OG4DOJk79VTczE,1423
+kinfer/export/pytorch.py,sha256=nRo9TKClRhPcLJQDR3oHNCbTYjf_3gLBxGfzZBScPxc,1303
+kinfer/rust_bindings/Cargo.toml,sha256=i1RGB9VNd9Q4FJ6gGwjZJQYo8DBBvpVWf3GJ95EfVgM,637
+kinfer/rust_bindings/pyproject.toml,sha256=jLcJuHCnQRh9HWR_R7a9qLHwj6LMBgnHyeKK_DruO1Y,135
+kinfer/rust_bindings/rust_bindings.pyi,sha256=lPGHMDsydr4Ue-CNHRjQtJWmq5sE5daVN8PbxJgOZY4,2048
+kinfer/rust_bindings/src/lib.rs,sha256=Uqi76Na_fGiecMeW5mSSmtYdz2wU8clLxYawbYomNOc,12662
+kinfer/rust_bindings/src/bin/stub_gen.rs,sha256=hhoVGnaSfazbSfj5a4x6mPicGPOgWQAfsDmiPej0B6Y,133
+kinfer-0.5.2.dist-info/RECORD,,
+kinfer-0.5.2.dist-info/WHEEL,sha256=sunMa2yiYbrNLGeMVDqEA0ayyJbHlex7SCn1TZrEq60,136
+kinfer-0.5.2.dist-info/top_level.txt,sha256=6mY_t3PYr3Dm0dpqMk80uSnArbvGfCFkxOh1QWtgDEo,7
+kinfer-0.5.2.dist-info/METADATA,sha256=VbHb0kFyouhwwF4mQ0ajRJQScLpdGGBYNOn_RiBsMFc,1831
+kinfer-0.5.2.dist-info/licenses/LICENSE,sha256=Qw-Z0XTwS-diSW91e_jLeBPX9zZbAatOJTBLdPHPaC0,1069

{kinfer-0.4.3.dist-info → kinfer-0.5.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.7.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: false
 Tag: cp311-cp311-macosx_11_0_arm64
 Generator: delocate 0.13.0

kinfer/common/types.py

@@ -1,12 +0,0 @@
-"""Defines common types."""
-
-__all__ = [
-    "Metadata",
-]
-
-from pydantic import BaseModel
-
-
-class Metadata(BaseModel):
-    joint_names: list[str]
-    num_commands: int | None

kinfer/export/common.py

@@ -1,44 +0,0 @@
-"""Defines common utilities for exporting models."""
-
-
-def get_shape(
-    name: str,
-    num_joints: int | None = None,
-    num_commands: int | None = None,
-    carry_shape: tuple[int, ...] | None = None,
-) -> tuple[int, ...]:
-    match name:
-        case "joint_angles":
-            if num_joints is None:
-                raise ValueError("`num_joints` must be provided when using `joint_angles`")
-            return (num_joints,)
-
-        case "joint_angular_velocities":
-            if num_joints is None:
-                raise ValueError("`num_joints` must be provided when using `joint_angular_velocities`")
-            return (num_joints,)
-
-        case "projected_gravity":
-            return (3,)
-
-        case "accelerometer":
-            return (3,)
-
-        case "gyroscope":
-            return (3,)
-
-        case "command":
-            if num_commands is None:
-                raise ValueError("`num_commands` must be provided when using `command`")
-            return (num_commands,)
-
-        case "carry":
-            if carry_shape is None:
-                raise ValueError("`carry_shape` must be provided for `carry`")
-            return carry_shape
-
-        case "time":
-            return (1,)
-
-        case _:
-            raise ValueError(f"Unknown tensor name: {name}")

kinfer-0.4.3.dist-info/RECORD

@@ -1,26 +0,0 @@
-kinfer/rust_bindings.cpython-311-darwin.so,sha256=ZUWCReCTDKhzqBteg3k4nJ-15ubYFMCSBYQdADqcvd0,1953904
-kinfer/requirements.txt,sha256=j08HO4ptA5afuj99j8FlAP2qla5Zf4_OiEBtgAmF7Jg,90
-kinfer/__init__.py,sha256=i5da6ND827Cgn8PFKzDCmEBk14ptQasLa_9fdof4Y9c,398
-kinfer/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kinfer/rust_bindings.pyi,sha256=lyBw-t_9ONYwVSaUTeeYDWpBbhKod1wf9Dvc4MgnC1g,1776
-kinfer/rust/Cargo.toml,sha256=uxjvuM1Sm82UgCdHW3R6PbbiJjXvfE4VCLvjGPqp2mY,606
-kinfer/rust/src/runtime.rs,sha256=l9wG0XbZmpHSqho3SfGuW54P7YXs8eXZ365kOj7UqUE,3321
-kinfer/rust/src/lib.rs,sha256=Z3dWdhKhqhWqPJee6vWde4ptqquOYW6W9wB0wyaKCyk,71
-kinfer/rust/src/model.rs,sha256=zY_XmkCZOq8VhVE_LpgTxKDrV-FceY8tnkkqS1tFa8c,11890
-kinfer/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kinfer/common/types.py,sha256=jiuCfoTEumOJm4VZ03VoFVIL5PMJ1tCG7b2Kr66qJQA,176
-kinfer/export/serialize.py,sha256=7rcSrXawVnJinxeJp26wovUwoR04BnJ5CZAV9vwKUUo,3474
-kinfer/export/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kinfer/export/jax.py,sha256=gZWWUJKb3MPEtfNYeKfPnypeK0elJwGuDYfpCW0GqNY,1413
-kinfer/export/pytorch.py,sha256=rwHGOTcd08IeL5-HZ_7JVUMPMHleiHI8-TkShAoRrYM,1612
-kinfer/export/common.py,sha256=AYN_gwmxmSlZ7BuswGUKYcY8uVUCh9UyXnWw-k8RVys,1290
-kinfer/rust_bindings/Cargo.toml,sha256=i1RGB9VNd9Q4FJ6gGwjZJQYo8DBBvpVWf3GJ95EfVgM,637
-kinfer/rust_bindings/pyproject.toml,sha256=jLcJuHCnQRh9HWR_R7a9qLHwj6LMBgnHyeKK_DruO1Y,135
-kinfer/rust_bindings/rust_bindings.pyi,sha256=lyBw-t_9ONYwVSaUTeeYDWpBbhKod1wf9Dvc4MgnC1g,1776
-kinfer/rust_bindings/src/lib.rs,sha256=pjsRej2BxTsecG9Y7GdBpW6GGR-pgbBk3yrKqEEQ3K0,11547
-kinfer/rust_bindings/src/bin/stub_gen.rs,sha256=hhoVGnaSfazbSfj5a4x6mPicGPOgWQAfsDmiPej0B6Y,133
-kinfer-0.4.3.dist-info/RECORD,,
-kinfer-0.4.3.dist-info/WHEEL,sha256=3lrG374qykB8NIZqJ4ApcNfFLoAGB9tcfdD4_UMfO40,136
-kinfer-0.4.3.dist-info/top_level.txt,sha256=6mY_t3PYr3Dm0dpqMk80uSnArbvGfCFkxOh1QWtgDEo,7
-kinfer-0.4.3.dist-info/METADATA,sha256=DZKaMDT-DToGkw_Ge_nSYBoo3ZAzxatlv6sR0x9kv-A,1761
-kinfer-0.4.3.dist-info/licenses/LICENSE,sha256=Qw-Z0XTwS-diSW91e_jLeBPX9zZbAatOJTBLdPHPaC0,1069