Initial commit

arm_plane_calib/fit_plane_calib.py

@@ -0,0 +1,60 @@
+#!/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()
+