Add CSV plotting and episode recording functionality
- Implemented `plot_data_csv.py` to read CSV files and generate plots for velocity and linear acceleration signals. - Created `record1.py` for recording ESP32 IMU data, DIN/DOUT states, and RealSense camera images into episodes. - Enhanced `ProviderWorldIMUVelocityEstimator` to include stationary detection logic, resetting velocity when stationary. - Updated `EpisodeWriter` to save episode data with timestamped filenames for better organization.
This commit is contained in:
Brunsmeier
200
plot_data_csv.py
Normal file
200
plot_data_csv.py
Normal file
@ -0,0 +1,200 @@
|
||||
import argparse
|
||||
import csv
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
TARGET_GROUPS = {
|
||||
"velocity": {
|
||||
"vel_x_ms": ["vel_x_ms"],
|
||||
"vel_y_ms": ["vel_y_ms"],
|
||||
"vel_z_ms": ["vel_z_ms"],
|
||||
},
|
||||
"lin_acc": {
|
||||
"lin_acc_x_ms2": ["lin_acc_x_ms2", "lin_acc_xms2"],
|
||||
"lin_acc_y_ms2": ["lin_acc_y_ms2", "lin_acc_yms2"],
|
||||
"lin_acc_z_ms2": ["lin_acc_z_ms2", "lin_acc_zms2"],
|
||||
},
|
||||
}
|
||||
|
||||
TIME_CANDIDATES = ["t_monotonic", "t_wall", "t_esp_ms", "frame_id"]
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Read CSV files from a folder and plot velocity/linear acceleration signals."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-dir",
|
||||
default="data_csv",
|
||||
help="Directory containing CSV files (default: data_csv)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default="png",
|
||||
help="Directory to save output images (default: png)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dpi",
|
||||
type=int,
|
||||
default=140,
|
||||
help="Output image DPI (default: 140)",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def pick_first_existing(header, candidates):
|
||||
for name in candidates:
|
||||
if name in header:
|
||||
return name
|
||||
return None
|
||||
|
||||
|
||||
def to_float(value):
|
||||
if value is None:
|
||||
return float("nan")
|
||||
text = str(value).strip()
|
||||
if text == "":
|
||||
return float("nan")
|
||||
try:
|
||||
return float(text)
|
||||
except ValueError:
|
||||
return float("nan")
|
||||
|
||||
|
||||
def read_csv_data(csv_path):
|
||||
with csv_path.open("r", encoding="utf-8", newline="") as f:
|
||||
reader = csv.DictReader(f)
|
||||
header = reader.fieldnames or []
|
||||
|
||||
time_col = pick_first_existing(header, TIME_CANDIDATES)
|
||||
resolved = {}
|
||||
for group in TARGET_GROUPS.values():
|
||||
for canonical_name, aliases in group.items():
|
||||
resolved[canonical_name] = pick_first_existing(header, aliases)
|
||||
|
||||
rows = list(reader)
|
||||
|
||||
if not rows:
|
||||
return None
|
||||
|
||||
x = []
|
||||
if time_col is None:
|
||||
x = list(range(len(rows)))
|
||||
x_label = "sample_index"
|
||||
else:
|
||||
raw_t = [to_float(r.get(time_col)) for r in rows]
|
||||
t0 = raw_t[0] if raw_t else 0.0
|
||||
if time_col.endswith("_ms"):
|
||||
x = [((v - t0) / 1000.0) for v in raw_t]
|
||||
x_label = f"{time_col} (s, relative)"
|
||||
else:
|
||||
x = [(v - t0) for v in raw_t]
|
||||
x_label = f"{time_col} (relative)"
|
||||
|
||||
data = {}
|
||||
for canonical_name, actual_name in resolved.items():
|
||||
if actual_name is None:
|
||||
data[canonical_name] = None
|
||||
else:
|
||||
data[canonical_name] = [to_float(r.get(actual_name)) for r in rows]
|
||||
|
||||
return {
|
||||
"x": x,
|
||||
"x_label": x_label,
|
||||
"data": data,
|
||||
"resolved": resolved,
|
||||
}
|
||||
|
||||
|
||||
def plot_one_csv(csv_path, output_dir, dpi):
|
||||
parsed = read_csv_data(csv_path)
|
||||
if parsed is None:
|
||||
print(f"[SKIP] {csv_path.name}: empty file")
|
||||
return False
|
||||
|
||||
x = parsed["x"]
|
||||
x_label = parsed["x_label"]
|
||||
data = parsed["data"]
|
||||
resolved = parsed["resolved"]
|
||||
|
||||
has_velocity = any(data[name] is not None for name in TARGET_GROUPS["velocity"].keys())
|
||||
has_lin_acc = any(data[name] is not None for name in TARGET_GROUPS["lin_acc"].keys())
|
||||
|
||||
if not has_velocity and not has_lin_acc:
|
||||
print(f"[SKIP] {csv_path.name}: no target columns found")
|
||||
return False
|
||||
|
||||
fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)
|
||||
fig.suptitle(f"Motion Signals - {csv_path.name}")
|
||||
|
||||
ax_v, ax_a = axes
|
||||
|
||||
for name in TARGET_GROUPS["velocity"].keys():
|
||||
series = data[name]
|
||||
if series is not None:
|
||||
ax_v.plot(x, series, label=name, linewidth=1.1)
|
||||
|
||||
for name in TARGET_GROUPS["lin_acc"].keys():
|
||||
series = data[name]
|
||||
if series is not None:
|
||||
ax_a.plot(x, series, label=name, linewidth=1.1)
|
||||
|
||||
ax_v.set_ylabel("velocity (m/s)")
|
||||
ax_a.set_ylabel("linear acc (m/s^2)")
|
||||
ax_a.set_xlabel(x_label)
|
||||
|
||||
ax_v.grid(True, alpha=0.25)
|
||||
ax_a.grid(True, alpha=0.25)
|
||||
|
||||
if has_velocity:
|
||||
ax_v.legend(loc="upper right")
|
||||
else:
|
||||
ax_v.text(0.5, 0.5, "No velocity columns", transform=ax_v.transAxes, ha="center", va="center")
|
||||
|
||||
if has_lin_acc:
|
||||
ax_a.legend(loc="upper right")
|
||||
else:
|
||||
ax_a.text(0.5, 0.5, "No linear-acc columns", transform=ax_a.transAxes, ha="center", va="center")
|
||||
|
||||
missing = [k for k, v in resolved.items() if v is None]
|
||||
if missing:
|
||||
fig.text(0.01, 0.01, f"Missing columns: {', '.join(missing)}", fontsize=9)
|
||||
|
||||
fig.tight_layout(rect=[0, 0.03, 1, 0.97])
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_path = output_dir / f"{csv_path.stem}_motion.png"
|
||||
fig.savefig(out_path, dpi=dpi)
|
||||
plt.close(fig)
|
||||
|
||||
print(f"[OK] {csv_path.name} -> {out_path}")
|
||||
return True
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
input_dir = Path(args.input_dir)
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
if not input_dir.exists() or not input_dir.is_dir():
|
||||
raise SystemExit(f"Input directory does not exist: {input_dir}")
|
||||
|
||||
csv_files = sorted(input_dir.glob("*.csv"))
|
||||
if not csv_files:
|
||||
raise SystemExit(f"No CSV files found in: {input_dir}")
|
||||
|
||||
ok_count = 0
|
||||
for csv_path in csv_files:
|
||||
if plot_one_csv(csv_path, output_dir, args.dpi):
|
||||
ok_count += 1
|
||||
|
||||
print(f"Finished. Generated charts for {ok_count}/{len(csv_files)} files.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user