kinfer 0.5.4__cp313-cp313-manylinux_2_28_x86_64.whl

kinfer/rust_bindings/src/lib.rs

@@ -0,0 +1,486 @@
+use async_trait::async_trait;
+use kinfer::model::{ModelError, ModelProvider, ModelRunner};
+use kinfer::runtime::ModelRuntime;
+use kinfer::types::{InputType, ModelMetadata};
+use ndarray::{Array, Ix1, IxDyn};
+use numpy::{PyArray1, PyArrayDyn, PyArrayMethods};
+use pyo3::exceptions::PyNotImplementedError;
+use pyo3::prelude::*;
+use pyo3::types::{PyAny, PyAnyMethods};
+use pyo3::{pymodule, types::PyModule, Bound, PyResult, Python};
+use pyo3_stub_gen::define_stub_info_gatherer;
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyfunction, gen_stub_pymethods};
+use std::collections::HashMap;
+use std::hash::Hash;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+type StepResult = (Py<PyArrayDyn<f32>>, Py<PyArrayDyn<f32>>);
+
+// Custom error type for Send/Sync compatibility
+#[derive(Debug)]
+struct SendError(String);
+
+unsafe impl Send for SendError {}
+unsafe impl Sync for SendError {}
+
+impl std::fmt::Display for SendError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+#[pyfunction]
+#[gen_stub_pyfunction]
+fn get_version() -> String {
+    env!("CARGO_PKG_VERSION").to_string()
+}
+
+#[pyclass]
+#[gen_stub_pyclass]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct PyInputType {
+    pub input_type: InputType,
+}
+
+impl From<InputType> for PyInputType {
+    fn from(input_type: InputType) -> Self {
+        Self { input_type }
+    }
+}
+
+impl From<PyInputType> for InputType {
+    fn from(input_type: PyInputType) -> Self {
+        input_type.input_type
+    }
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyInputType {
+    #[new]
+    fn __new__(input_type: &str) -> PyResult<Self> {
+        let input_type = InputType::from_name(input_type).map_or_else(
+            |_| {
+                Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+                    "Invalid input type: {} (must be one of {})",
+                    input_type,
+                    InputType::get_names().join(", "),
+                )))
+            },
+            Ok,
+        )?;
+        Ok(Self { input_type })
+    }
+
+    fn get_name(&self) -> String {
+        self.input_type.get_name().to_string()
+    }
+
+    fn get_shape(&self, metadata: PyModelMetadata) -> Vec<usize> {
+        self.input_type.get_shape(&metadata.into())
+    }
+
+    fn __repr__(&self) -> String {
+        format!("InputType({})", self.get_name())
+    }
+
+    fn __eq__(&self, other: Bound<'_, PyAny>) -> PyResult<bool> {
+        if let Ok(other) = other.extract::<PyInputType>() {
+            Ok(self == &other)
+        } else {
+            Ok(false)
+        }
+    }
+}
+
+#[pyclass]
+#[gen_stub_pyclass]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct PyModelMetadata {
+    #[pyo3(get, set)]
+    pub joint_names: Vec<String>,
+    #[pyo3(get, set)]
+    pub num_commands: Option<usize>,
+    #[pyo3(get, set)]
+    pub carry_size: Vec<usize>,
+}
+
+#[pymethods]
+#[gen_stub_pymethods]
+impl PyModelMetadata {
+    #[new]
+    fn __new__(
+        joint_names: Vec<String>,
+        num_commands: Option<usize>,
+        carry_size: Vec<usize>,
+    ) -> Self {
+        Self {
+            joint_names,
+            num_commands,
+            carry_size,
+        }
+    }
+
+    fn to_json(&self) -> PyResult<String> {
+        let metadata = ModelMetadata {
+            joint_names: self.joint_names.clone(),
+            num_commands: self.num_commands,
+            carry_size: self.carry_size.clone(),
+        }
+        .to_json()
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        Ok(metadata)
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        let json = self.to_json()?;
+        Ok(format!("ModelMetadata({:?})", json))
+    }
+
+    fn __eq__(&self, other: Bound<'_, PyAny>) -> PyResult<bool> {
+        if let Ok(other) = other.extract::<PyModelMetadata>() {
+            Ok(self == &other)
+        } else {
+            Ok(false)
+        }
+    }
+}
+
+#[pyfunction]
+#[gen_stub_pyfunction]
+fn metadata_from_json(json: &str) -> PyResult<PyModelMetadata> {
+    let metadata = ModelMetadata::model_validate_json(json.to_string()).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid model metadata: {}", e))
+    })?;
+    Ok(PyModelMetadata::from(metadata))
+}
+
+impl From<ModelMetadata> for PyModelMetadata {
+    fn from(metadata: ModelMetadata) -> Self {
+        Self {
+            joint_names: metadata.joint_names,
+            num_commands: metadata.num_commands,
+            carry_size: metadata.carry_size,
+        }
+    }
+}
+
+impl From<&ModelMetadata> for PyModelMetadata {
+    fn from(metadata: &ModelMetadata) -> Self {
+        Self {
+            joint_names: metadata.joint_names.clone(),
+            num_commands: metadata.num_commands,
+            carry_size: metadata.carry_size.clone(),
+        }
+    }
+}
+
+impl From<PyModelMetadata> for ModelMetadata {
+    fn from(metadata: PyModelMetadata) -> Self {
+        Self {
+            joint_names: metadata.joint_names,
+            num_commands: metadata.num_commands,
+            carry_size: metadata.carry_size,
+        }
+    }
+}
+
+#[pyclass(subclass)]
+#[gen_stub_pyclass]
+struct ModelProviderABC;
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl ModelProviderABC {
+    #[new]
+    fn __new__() -> Self {
+        ModelProviderABC
+    }
+
+    fn get_inputs<'py>(
+        &self,
+        input_types: Vec<String>,
+        metadata: PyModelMetadata,
+    ) -> PyResult<HashMap<String, Bound<'py, PyArrayDyn<f32>>>> {
+        Err(PyNotImplementedError::new_err(format!(
+            "Must override get_inputs with {} input types {:?} and metadata {:?}",
+            input_types.len(),
+            input_types,
+            metadata
+        )))
+    }
+
+    fn take_action(
+        &self,
+        action: Bound<'_, PyArray1<f32>>,
+        metadata: PyModelMetadata,
+    ) -> PyResult<()> {
+        let n = action.len()?;
+        if metadata.joint_names.len() != n {
+            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+                "Expected {} joints, got {} action elements",
+                metadata.joint_names.len(),
+                n
+            )));
+        }
+        Err(PyNotImplementedError::new_err(format!(
+            "Must override take_action with {} action elements",
+            n
+        )))
+    }
+}
+
+#[gen_stub_pyclass]
+#[pyclass]
+#[derive(Clone)]
+struct PyModelProvider {
+    obj: Arc<Py<ModelProviderABC>>,
+}
+
+#[pymethods]
+impl PyModelProvider {
+    #[new]
+    fn __new__(obj: Py<ModelProviderABC>) -> Self {
+        Self { obj: Arc::new(obj) }
+    }
+}
+
+#[async_trait]
+impl ModelProvider for PyModelProvider {
+    async fn get_inputs(
+        &self,
+        input_types: &[InputType],
+        metadata: &ModelMetadata,
+    ) -> Result<HashMap<InputType, Array<f32, IxDyn>>, ModelError> {
+        let input_names: Vec<String> = input_types
+            .iter()
+            .map(|t| t.get_name().to_string())
+            .collect();
+        let result = Python::with_gil(|py| -> PyResult<HashMap<InputType, Array<f32, IxDyn>>> {
+            let obj = self.obj.clone();
+            let args = (input_names.clone(), PyModelMetadata::from(metadata.clone()));
+            let result = obj.call_method(py, "get_inputs", args, None)?;
+            let dict: HashMap<String, Vec<f32>> = result.extract(py)?;
+            let mut arrays = HashMap::new();
+            for (i, name) in input_names.iter().enumerate() {
+                let array = dict.get(name).ok_or_else(|| {
+                    PyErr::new::<pyo3::exceptions::PyKeyError, _>(format!(
+                        "Missing input: {}",
+                        name
+                    ))
+                })?;
+                arrays.insert(input_types[i], Array::from_vec(array.clone()).into_dyn());
+            }
+            Ok(arrays)
+        })
+        .map_err(|e| ModelError::Provider(e.to_string()))?;
+        Ok(result)
+    }
+
+    async fn take_action(
+        &self,
+        action: Array<f32, IxDyn>,
+        metadata: &ModelMetadata,
+    ) -> Result<(), ModelError> {
+        Python::with_gil(|py| -> PyResult<()> {
+            let obj = self.obj.clone();
+            let action_1d = action
+                .into_dimensionality::<Ix1>()
+                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+            let args = (
+                PyArray1::from_array(py, &action_1d),
+                PyModelMetadata::from(metadata.clone()),
+            );
+            obj.call_method(py, "take_action", args, None)?;
+            Ok(())
+        })
+        .map_err(|e| ModelError::Provider(e.to_string()))?;
+        Ok(())
+    }
+}
+
+#[gen_stub_pyclass]
+#[pyclass]
+#[derive(Clone)]
+struct PyModelRunner {
+    runner: Arc<ModelRunner>,
+    runtime: Arc<tokio::runtime::Runtime>,
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyModelRunner {
+    #[new]
+    fn __new__(model_path: String, provider: Py<ModelProviderABC>) -> PyResult<Self> {
+        let input_provider = Arc::new(PyModelProvider::__new__(provider));
+
+        // Create a single runtime to be reused for all operations
+        let runtime = Arc::new(tokio::runtime::Runtime::new()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?);
+
+        let runner = runtime.block_on(async {
+            ModelRunner::new(model_path, input_provider)
+                .await
+                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
+        })?;
+
+        Ok(Self {
+            runner: Arc::new(runner),
+            runtime,
+        })
+    }
+
+    // Reuse runtime and release GIL
+    fn init(&self) -> PyResult<Py<PyArrayDyn<f32>>> {
+        let runner = self.runner.clone();
+        let runtime = self.runtime.clone();
+
+        let result = Python::with_gil(|py| {
+            // Release GIL during async operation
+            py.allow_threads(|| {
+                runtime.block_on(async {
+                    runner
+                        .init()
+                        .await
+                        .map_err(|e| SendError(e.to_string()))
+                })
+            })
+        })
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.0))?;
+
+        Python::with_gil(|py| {
+            let array = numpy::PyArray::from_array(py, &result);
+            Ok(array.into())
+        })
+    }
+
+    // Reuse runtime and release GIL
+    fn step(&self, carry: Py<PyArrayDyn<f32>>) -> PyResult<StepResult> {
+        let runner = self.runner.clone();
+        let runtime = self.runtime.clone();
+        
+        // Extract the carry array from Python with GIL
+        let carry_array = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
+            let carry_array = carry.bind(py);
+            Ok(carry_array.to_owned_array())
+        })?;
+
+        // Release GIL during computation
+        let result = Python::with_gil(|py| {
+            py.allow_threads(|| {
+                runtime.block_on(async {
+                    runner
+                        .step(carry_array)
+                        .await
+                        .map_err(|e| SendError(e.to_string()))
+                })
+            })
+        })
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.0))?;
+
+        // Reacquire the GIL to convert results back to Python objects
+        Python::with_gil(|py| {
+            let (output, carry) = result;
+            let output_array = numpy::PyArray::from_array(py, &output);
+            let carry_array = numpy::PyArray::from_array(py, &carry);
+            Ok((output_array.into(), carry_array.into()))
+        })
+    }
+
+    // Reuse runtime and release GIL
+    fn take_action(&self, action: Py<PyArrayDyn<f32>>) -> PyResult<()> {
+        let runner = self.runner.clone();
+        let runtime = self.runtime.clone();
+
+        // Extract action data with GIL
+        let action_array = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
+            let action_array = action.bind(py);
+            Ok(action_array.to_owned_array())
+        })?;
+        
+        // Release GIL during computation
+        Python::with_gil(|py| {
+            py.allow_threads(|| {
+                runtime.block_on(async {
+                    runner
+                        .take_action(action_array)
+                        .await
+                        .map_err(|e| SendError(e.to_string()))
+                })
+            })
+        })
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.0))?;
+
+        Ok(())
+    }
+}
+
+#[gen_stub_pyclass]
+#[pyclass]
+#[derive(Clone)]
+struct PyModelRuntime {
+    runtime: Arc<Mutex<ModelRuntime>>,
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyModelRuntime {
+    #[new]
+    fn __new__(model_runner: PyModelRunner, dt: u64) -> PyResult<Self> {
+        Ok(Self {
+            runtime: Arc::new(Mutex::new(ModelRuntime::new(model_runner.runner, dt))),
+        })
+    }
+
+    fn set_slowdown_factor(&self, slowdown_factor: i32) -> PyResult<()> {
+        let mut runtime = self
+            .runtime
+            .lock()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        runtime.set_slowdown_factor(slowdown_factor);
+        Ok(())
+    }
+
+    fn set_magnitude_factor(&self, magnitude_factor: f32) -> PyResult<()> {
+        let mut runtime = self
+            .runtime
+            .lock()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        runtime.set_magnitude_factor(magnitude_factor);
+        Ok(())
+    }
+
+    fn start(&self) -> PyResult<()> {
+        let mut runtime = self
+            .runtime
+            .lock()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        runtime
+            .start()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
+    }
+
+    fn stop(&self) -> PyResult<()> {
+        let mut runtime = self
+            .runtime
+            .lock()
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        runtime.stop();
+        Ok(())
+    }
+}
+
+#[pymodule]
+fn rust_bindings(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(get_version, m)?)?;
+    m.add_class::<PyInputType>()?;
+    m.add_class::<PyModelMetadata>()?;
+    m.add_function(wrap_pyfunction!(metadata_from_json, m)?)?;
+    m.add_class::<ModelProviderABC>()?;
+    m.add_class::<PyModelRunner>()?;
+    m.add_class::<PyModelRuntime>()?;
+    Ok(())
+}
+
+define_stub_info_gatherer!(stub_info);

kinfer/rust_bindings.pyi

@@ -0,0 +1,46 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401
+
+import builtins
+import numpy
+import numpy.typing
+import typing
+
+class ModelProviderABC:
+    def __new__(cls) -> ModelProviderABC: ...
+    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:
+    ...
+
+class PyModelRunner:
+    def __new__(cls, model_path:builtins.str, provider:ModelProviderABC) -> PyModelRunner: ...
+    def init(self) -> numpy.typing.NDArray[numpy.float32]: ...
+    def step(self, carry:numpy.typing.NDArray[numpy.float32]) -> tuple[numpy.typing.NDArray[numpy.float32], numpy.typing.NDArray[numpy.float32]]: ...
+    def take_action(self, action:numpy.typing.NDArray[numpy.float32]) -> None: ...
+
+class PyModelRuntime:
+    def __new__(cls, model_runner:PyModelRunner, dt:builtins.int) -> PyModelRuntime: ...
+    def set_slowdown_factor(self, slowdown_factor:builtins.int) -> None: ...
+    def set_magnitude_factor(self, magnitude_factor:builtins.float) -> None: ...
+    def start(self) -> None: ...
+    def stop(self) -> None: ...
+
+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()