85 lines
3.1 KiB
Python
85 lines
3.1 KiB
Python
import csv
|
||
import time
|
||
import os
|
||
import numpy as np
|
||
|
||
|
||
class DataLogger:
|
||
def __init__(self, filename_prefix="dataset"):
|
||
"""
|
||
初始化日志记录器
|
||
:param filename_prefix: 文件名前缀，例如 "train_batch"
|
||
"""
|
||
# 1. 自动创建 data 文件夹，保持项目目录整洁
|
||
if not os.path.exists('data'):
|
||
os.makedirs('data')
|
||
print("[Logger] Created 'data' directory.")
|
||
|
||
# 2. 生成带时间戳的文件名，防止覆盖旧数据
|
||
# 格式: data/dataset_20240105_143022.csv
|
||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||
self.filename = f"data/{filename_prefix}_{timestamp}.csv"
|
||
|
||
# 3. 定义标准表头 (Header) - 对应你的12个力 + 3个构型 + 3个标签
|
||
# 传感器命名规则: F(手指ID)_U(单元ID)_(Fn/Ft)
|
||
sensor_headers = []
|
||
for f in range(3): # 3个手指
|
||
for u in range(1, 3): # 2个单元
|
||
sensor_headers.append(f"F{f}_U{u}_Fn")
|
||
sensor_headers.append(f"F{f}_U{u}_Ft")
|
||
|
||
# 完整的列名列表
|
||
self.header = ['Timestamp', 'Config_ID'] + sensor_headers + ['Label_X', 'Label_Y', 'Label_Theta']
|
||
|
||
# 4. 创建文件并写入第一行（表头）
|
||
try:
|
||
with open(self.filename, mode='w', newline='') as f:
|
||
writer = csv.writer(f)
|
||
writer.writerow(self.header)
|
||
print(f"[Logger] Log file initialized: {self.filename}")
|
||
except Exception as e:
|
||
print(f"[Logger] Error creating file: {e}")
|
||
|
||
def log(self, sensor_data, config_id, label_vector):
|
||
"""
|
||
写入一行数据
|
||
:param sensor_data: 12维 numpy 数组或列表 (触觉数据)
|
||
:param config_id: 整数 (0, 1, 2) (当前手指构型)
|
||
:param label_vector: [x, y, theta] (机械臂的真实偏差值)
|
||
"""
|
||
with open(self.filename, mode='a', newline='') as f:
|
||
writer = csv.writer(f)
|
||
|
||
# 拼接数据: 时间 + ID + 传感器数值 + 标签
|
||
# time.time() 获取当前精确时间戳
|
||
row = [f"{time.time():.3f}", int(config_id)] + list(sensor_data) + list(label_vector)
|
||
|
||
# 格式化: 将浮点数保留4位小数，节省空间且美观
|
||
formatted_row = []
|
||
for item in row:
|
||
if isinstance(item, float):
|
||
formatted_row.append(f"{item:.4f}")
|
||
else:
|
||
formatted_row.append(item)
|
||
|
||
writer.writerow(formatted_row)
|
||
|
||
|
||
# === 单元测试 (Unit Test) ===
|
||
# 直接运行此文件，测试是否能生成 CSV
|
||
if __name__ == "__main__":
|
||
print("Testing DataLogger...")
|
||
logger = DataLogger(filename_prefix="test_data")
|
||
|
||
# 模拟写入 5 条假数据
|
||
for i in range(5):
|
||
# 模拟 12 个传感器数据 (0-10之间)
|
||
fake_sensor = np.random.rand(12) * 10
|
||
# 模拟机械臂偏差 (Label)
|
||
fake_label = [0.0, 5.5, -2.0]
|
||
|
||
logger.log(sensor_data=fake_sensor, config_id=0, label_vector=fake_label)
|
||
print(f"Logged row {i + 1}")
|
||
time.sleep(0.1)
|
||
|
||
print("Test done. Please check the 'data' folder.") |