60 lines
1.6 KiB
Python
60 lines
1.6 KiB
Python
#!/usr/bin/env python3
|
|
import yaml
|
|
import numpy as np
|
|
import sys
|
|
|
|
|
|
def features(x, y):
|
|
return np.array([1.0, x, y, x * x, x * y, y * y], dtype=np.float64)
|
|
|
|
|
|
def main():
|
|
sample_path = sys.argv[1] if len(sys.argv) > 1 else 'samples.yaml'
|
|
model_path = sys.argv[2] if len(sys.argv) > 2 else 'model.yaml'
|
|
|
|
with open(sample_path, 'r', encoding='utf-8') as f:
|
|
data = yaml.safe_load(f)
|
|
|
|
samples = data['samples']
|
|
if len(samples) < 6:
|
|
print('样本太少，至少需要 6 个，建议 9~16 个')
|
|
return
|
|
|
|
X = []
|
|
Y = []
|
|
|
|
for s in samples:
|
|
X.append(features(s['x_cam'], s['y_cam']))
|
|
Y.append(s['arm_cmd'])
|
|
|
|
X = np.array(X, dtype=np.float64) # [N, 6]
|
|
Y = np.array(Y, dtype=np.float64) # [N, cmd_dim]
|
|
|
|
# 最小二乘
|
|
W, _, _, _ = np.linalg.lstsq(X, Y, rcond=None) # [6, cmd_dim]
|
|
Y_pred = X @ W
|
|
err = Y_pred - Y
|
|
mae = np.mean(np.abs(err), axis=0)
|
|
rmse = np.sqrt(np.mean(err ** 2, axis=0))
|
|
|
|
model = {
|
|
'feature_order': ['1', 'x', 'y', 'x2', 'xy', 'y2'],
|
|
'weights': W.tolist(),
|
|
'cmd_dim': int(Y.shape[1]),
|
|
'sample_count': int(len(samples)),
|
|
'mae_per_dim': mae.tolist(),
|
|
'rmse_per_dim': rmse.tolist(),
|
|
'z_ref_median': float(np.median([s['z_cam'] for s in samples]))
|
|
}
|
|
|
|
with open(model_path, 'w', encoding='utf-8') as f:
|
|
yaml.safe_dump(model, f, allow_unicode=True, sort_keys=False)
|
|
|
|
print(f'模型已保存到: {model_path}')
|
|
print('每维 MAE:', mae)
|
|
print('每维 RMSE:', rmse)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|