robot_sf_ll7/examples/plotting/data_analysis_example.py at main · ll7/robot_sf_ll7 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
"""Perform a showcase on how to use the data analysis module.

Purpose:
    Demonstrate converting simulation recordings to JSON, plotting analytics,
    and saving figures via the `robot_sf.data_analysis` helpers.

Usage:
    uv run python examples/plotting/data_analysis_example.py

Prerequisites:
    - Recording pickle in `examples/recordings/` (script selects the latest)
    - Optional dataset directory `examples/datasets/` for JSON exports

Expected Output:
    - Plots written under `robot_sf/data_analysis/plots`
    - In smoke mode, JSON exports written under `output/examples/datasets/`
    - Console logs describing generated assets

Limitations:
    - Requires example recordings to be present; otherwise logs an error.
"""

import os
from pathlib import Path

from loguru import logger

from examples.demo_utils import fast_demo_enabled
from robot_sf.common.artifact_paths import get_artifact_root
from robot_sf.data_analysis.extract_json_from_pickle import (
    extract_timestamp,
    plot_all_data_json,
    save_to_json,
)
from robot_sf.data_analysis.extract_obj_from_pickle import (
    ensure_dir_exists,
    plot_all_data_pkl,
)
from robot_sf.render.playback_recording import load_states


def show_from_pkl(filename: str, unique_id: str):
    """
    Extract and plot data from a pickle file.

    Args:
        filename (str): Path to the pickle file
        unique_id (str): Unique identifier for plot filenames
    Returns:
        None
    """
    # Load states and map definition
    states, map_def = load_states(filename)

    # Plot all available data
    plot_all_data_pkl(states, map_def, unique_id)


def show_from_json(filename: str, unique_id: str):
    """
    Convert recording file into json and plot the data.

    Args:
        filename (str): Path to the JSON file
        unique_id (str): Unique identifier for plot filenames
    Returns:
        None
    """
    dataset_dir = Path("examples/datasets")
    if dataset_dir.exists():
        # Convert recording to json
        save_to_json(filename, f"{dataset_dir}/{unique_id}.json")

        latest_file = max(dataset_dir.glob("*.json"), key=os.path.getctime, default=None)

        if latest_file:
            # Plot all available data
            plot_all_data_json(str(latest_file), unique_id)

        else:
            logger.error("No json files found in the 'examples/datasets' directory")
    else:
        logger.error("'examples/datasets' directory not found")


if __name__ == "__main__":
    # Example usage
    # Ensure the plots directory exists
    PLOTS_DIR = "robot_sf/data_analysis/plots"
    ensure_dir_exists(PLOTS_DIR)

    # Find the most recent recording file
    recording_dir = Path("examples/recordings")
    if recording_dir.exists():
        recording_path = max(recording_dir.glob("*.pkl"), key=os.path.getctime, default=None)

        if recording_path:
            recording_id = extract_timestamp(str(recording_path))

            if fast_demo_enabled():
                dataset_dir = get_artifact_root() / "examples" / "datasets"
                dataset_dir.mkdir(parents=True, exist_ok=True)
                json_path = dataset_dir / f"{recording_id}.json"
                save_to_json(str(recording_path), str(json_path))
                logger.info(
                    "Fast demo mode: exported JSON and skipped plotting "
                    f"(recording_id={recording_id}, recording_path={recording_path}, "
                    f"json_path={json_path})"
                )
            else:
                show_from_json(str(recording_path), recording_id)
                show_from_pkl(str(recording_path), recording_id)
                logger.info(f"Successfully extracted and plotted data from {recording_path}")
                logger.info(f"Plots saved in {os.path.abspath(PLOTS_DIR)}")
        else:
            logger.error("No recording files found in the 'examples/recordings' directory")
    else:
        logger.error("'examples/recordings' directory not found")