libs/module: Add attitude_estimator.py.

Signed-off-by: hlym123 <[email protected]>

Author: hlym123

m5stack/libs/attitude_estimator.py

@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: 2025 M5Stack Technology CO LTD
+#
+# SPDX-License-Identifier: MIT
+import math
+import time
+
+
+class AttitudeEstimator:
+    def __init__(self):
+        self.gyro_x_offset = 0.0
+        self.gyro_y_offset = 0.0
+        self.gyro_z_offset = 0.0
+        self.angle_gyro_x = 0.0
+        self.angle_gyro_y = 0.0
+        self.angle_gyro_z = 0.0
+        self.angle_x = 0.0
+        self.angle_y = 0.0
+        self.angle_z = 0.0
+        self.pre_interval = 0.0
+        self.gyro_coef = 0.98
+        self.acc_coef = 0.02
+
+    def calibrate_gyro(self, read_gyro_func, sample_count=100, delay=0.01):
+        sum_gx = 0.0
+        sum_gy = 0.0
+        sum_gz = 0.0
+
+        print("Starting gyroscope calibration, please keep the sensor stationary...")
+        for i in range(sample_count):
+            gx, gy, gz = read_gyro_func()
+            sum_gx += gx
+            sum_gy += gy
+            sum_gz += gz
+            time.sleep(delay)
+
+        self.gyro_x_offset = sum_gx / sample_count
+        self.gyro_y_offset = sum_gy / sample_count
+        self.gyro_z_offset = sum_gz / sample_count
+
+        print("Gyroscope calibration completed:")
+        print("gyro_x_offset = {:.2f} °/s".format(self.gyro_x_offset))
+        print("gyro_y_offset = {:.2f} °/s".format(self.gyro_y_offset))
+        print("gyro_z_offset = {:.2f} °/s".format(self.gyro_z_offset))
+
+    def update_attitude(self, gx, gy, gz, ax, ay, az, dt):
+        """
+        Update attitude angles (fusing accelerometer and gyroscope data)
+
+        :param gx: Gyroscope X-axis angular velocity (°/s)
+        :param gy: Gyroscope Y-axis angular velocity (°/s)
+        :param gz: Gyroscope Z-axis angular velocity (°/s)
+        :param ax: Accelerometer X-axis acceleration (m/s²)
+        :param ay: Accelerometer Y-axis acceleration (m/s²)
+        :param az: Accelerometer Z-axis acceleration (m/s²)
+        :param dt: Sampling time interval (seconds)
+        """
+        angle_acc_x = math.atan2(ay, az + abs(ax)) * 360 / (2 * math.pi)
+        angle_acc_y = math.atan2(ax, az + abs(ay)) * 360 / (-2 * math.pi)
+
+        gyro_x = gx - self.gyro_x_offset
+        gyro_y = gy - self.gyro_y_offset
+        gyro_z = gz - self.gyro_z_offset
+
+        self.angle_gyro_x += gyro_x * dt
+        self.angle_gyro_y += gyro_y * dt
+        self.angle_gyro_z += gyro_z * dt
+
+        self.angle_x = (self.gyro_coef * (self.angle_x + gyro_x * dt)) + (
+            self.acc_coef * angle_acc_x
+        )
+        self.angle_y = (self.gyro_coef * (self.angle_y + gyro_y * dt)) + (
+            self.acc_coef * angle_acc_y
+        )
+        self.angle_z = self.angle_gyro_z
+        self.angle_z = (self.angle_z + 360) % 360
+
+        self.pre_interval += dt
+
+    def get_angles(self):
+        return self.angle_x, self.angle_y, self.angle_z

m5stack/libs/module/bala2.py

@@ -5,11 +5,10 @@
 import struct
 from micropython import const
 from module.mbus import i2c1
-from pid import PIDController
-from imu import AttitudeEstimator
-from driver.mpu6886 import MPU6886
+import pid
+import attitude_estimator
 from driver import mpu6886
-from machine import Timer
+import machine
 
 
 REG_MTR1_CTRL = const(0x00)  # motor 1 control register
@@ -44,13 +43,13 @@ def __init__(self, timer_id: int = 0):
         # imu
         if 0x68 not in self.i2c.scan():
             raise RuntimeError("MPU6886 not found")
-        self.sensor = MPU6886(
+        self.sensor = mpu6886.MPU6886(
             self.i2c,
             accel_fs=mpu6886.ACCEL_FS_SEL_4G,
             gyro_fs=mpu6886.GYRO_FS_SEL_1000DPS,
             gyro_sf=mpu6886.SF_DEG_S,
         )
-        self.imu = AttitudeEstimator()
+        self.imu = attitude_estimator.AttitudeEstimator()
         # PID controller
         self.pid_angle_kp = 30
         self.pid_angle_ki = 0
@@ -60,15 +59,15 @@ def __init__(self, timer_id: int = 0):
         self.pid_speed_ki = 0.18
         self.pid_speed_kd = 0
         self.pid_speed_target = 0
-        self.pid_angle = PIDController(
+        self.pid_angle = pid.PIDController(
             kp=self.pid_angle_kp,
             ki=self.pid_angle_ki,
             kd=self.pid_angle_kd,
             setpoint=self.pid_angle_target,
             direction=1,
         )
         self.pid_angle.set_output_limits(-1023, 1023)
-        self.pid_speed = PIDController(
+        self.pid_speed = pid.PIDController(
             kp=self.pid_speed_kp,
             ki=self.pid_speed_ki,
             kd=self.pid_speed_kd,
@@ -78,7 +77,7 @@ def __init__(self, timer_id: int = 0):
         self.pid_speed.set_integral_limits(-80, 80)
         self.pid_speed.set_output_limits(-1023, 1023)
         # others
-        self.timer = Timer(timer_id)
+        self.timer = machine.Timer(timer_id)
         self.dt = 0.01
         self.motor_speed = 0
         self.prev_enc = 0
@@ -273,7 +272,7 @@ def start(self) -> None:
                 module_bala2_0.start()
         """
         self.timer.init(
-            period=int(self.dt * 1000), mode=Timer.PERIODIC, callback=self._balance_task
+            period=int(self.dt * 1000), mode=machine.Timer.PERIODIC, callback=self._balance_task
         )
 
     def stop(self) -> None: