AHRS visualize

  1'''
  2AHRS - AHRS_visualize.py with asynchronous mode.
  3
  4This example visualize AHRS by using 3D model.
  5
  6-------------------------------------------------------------------------------------
  7Please change correct serial number or IP and port number BEFORE you run example code.
  8
  9For other examples please check:
 10    https://github.com/WPC-Systems-Ltd/WPC_Python_driver_release/tree/main/examples
 11See README.md file to get detailed usage of this example.
 12
 13Copyright (c) 2022-2024 WPC Systems Ltd. All rights reserved.
 14'''
 15
 16## Python
 17import asyncio
 18import numpy as np
 19import stl.mesh as mesh
 20import mpl_toolkits.mplot3d as mplot3d
 21from PIL import Image
 22import matplotlib.pyplot as plt
 23import matplotlib.gridspec as gridspec
 24import matplotlib.image as mpimg
 25import matplotlib.font_manager as font_manager
 26from matplotlib import rcParams
 27from matplotlib.transforms import Affine2D
 28
 29## WPC
 30
 31from wpcsys import pywpc
 32
 33
 34################################################################################
 35## Configuration
 36
 37DATA_PATH = 'Material/viz_data/'
 38IMG_PATH = 'Material/viz_data/avion_'
 39
 40for font in font_manager.findSystemFonts(DATA_PATH):
 41    font_manager.fontManager.addfont(font)
 42
 43## Set font family globally
 44rcParams['font.family'] = 'Digital-7 Mono'
 45
 46## Style
 47plt.style.use("bmh")
 48BMH_BURGUNDY = "#a70f34"
 49BMH_BLUE = "#3d90be"
 50BMH_PURPLE = "#7c6ca4"
 51
 52for font in font_manager.findSystemFonts(DATA_PATH):
 53    font_manager.fontManager.addfont(font)
 54
 55## Set font family globally
 56rcParams['font.family'] = 'Digital-7 Mono'
 57
 58IMG_WPC = plt.imread('Material/trademark.jpg')
 59IMG_WPC_PIL = Image.open('Material/trademark.jpg')
 60
 61################################################################################
 62## Plane Images and Constants
 63
 64IMG_YAW = mpimg.imread(f'{IMG_PATH}yaw.png')
 65IMG_PITCH = mpimg.imread(f'{IMG_PATH}pitch.png')
 66IMG_ROLL = mpimg.imread(f'{IMG_PATH}roll.png')
 67
 68SIZE_YAW = np.array([IMG_YAW.shape[0], IMG_YAW.shape[1]])
 69CENTER_YAW = SIZE_YAW / 2
 70DIAG_YAW = np.sqrt(np.sum(SIZE_YAW ** 2))
 71ALPHA_YAW = (DIAG_YAW - SIZE_YAW) / 2
 72
 73
 74SIZE_PITCH = np.array([IMG_PITCH.shape[0], IMG_PITCH.shape[1]])
 75CENTER_PITCH = SIZE_PITCH / 2
 76DIAG_PITCH = np.sqrt(np.sum(SIZE_PITCH ** 2))
 77ALPHA_PITCH = (DIAG_PITCH - SIZE_PITCH) / 2
 78
 79SIZE_ROLL = np.array([IMG_ROLL.shape[0], IMG_ROLL.shape[1]])
 80CENTER_ROLL = SIZE_ROLL / 2
 81DIAG_ROLL = np.sqrt(np.sum(SIZE_ROLL ** 2))
 82ALPHA_ROLL = (DIAG_ROLL - SIZE_ROLL) / 2
 83
 84## Dictionaries
 85PLANE_DICT = {
 86  'yaw': dict(img=IMG_YAW, size=SIZE_YAW, center=CENTER_YAW, diag=DIAG_YAW, alpha=ALPHA_YAW),
 87  'pitch': dict(img=IMG_PITCH, size=SIZE_PITCH, center=CENTER_PITCH, diag=DIAG_PITCH, alpha=ALPHA_PITCH),
 88  'roll': dict(img=IMG_ROLL, size=SIZE_ROLL, center=CENTER_ROLL, diag=DIAG_ROLL, alpha = ALPHA_ROLL),
 89}
 90
 91MODEL_DICT = {
 92  'cat': dict(label='Cat', view=400),
 93  'rat': dict(label='Rat', view=50)
 94}
 95DEFAULT_MODEL = 'rat'
 96
 97################################################################################
 98## Constants
 99
100DEGREE_TO_RADIAN = np.pi / 180.0
101RADIAN_TO_DEGREE = 180.0 / np.pi
102
103################################################################################
104## Functions - utilities
105
106def WPC_initializeFigure(nb_rows=1, nb_col=1, share_x='all', share_y='all', option='none'):
107  fig = plt.figure(figsize=(10, 6))
108  fig.clf()
109  gs = gridspec.GridSpec(3, 3, width_ratios=[3, 1, 1], height_ratios=[1, 1, 1])
110
111  ## If `option` is `3D`, return a 3-D axis.
112  if option == '3D':
113    ax_main = fig.add_subplot(gs[:, 0], projection='3d')
114    ax_main.set_title('Main 3D Plot')
115    ax_plane_yaw = fig.add_subplot(gs[0, 1])
116    ax_plane_pitch = fig.add_subplot(gs[1, 1])
117    ax_plane_roll = fig.add_subplot(gs[2, 1])
118    ax_value_yaw = fig.add_subplot(gs[0, 2])
119    ax_value_pitch = fig.add_subplot(gs[1, 2])
120    ax_value_roll = fig.add_subplot(gs[2, 2])
121
122    Pil_width = IMG_WPC_PIL.width
123    Pil_height = IMG_WPC_PIL.height
124    new_pil_image = IMG_WPC_PIL.resize((int(Pil_width//3), int(Pil_height//3)))
125    fig.figimage(new_pil_image, xo=fig.bbox.xmin, yo=fig.bbox.ymin, alpha=0.8, zorder = 0)
126
127    ## fig.figimage(IMG_WPC, xo=fig.bbox.xmin, yo=fig.bbox.ymin, alpha=1)
128    ax_dict = {
129      'main': ax_main,
130      'plane_yaw': ax_plane_yaw,
131      'plane_pitch': ax_plane_pitch,
132      'plane_roll': ax_plane_roll,
133      'value_yaw': ax_value_yaw,
134      'value_pitch': ax_value_pitch,
135      'value_roll': ax_value_roll,
136    }
137
138    return fig, ax_dict
139
140  ## If `option` is `grid`, return gridspec for further usage.
141  if option == 'grid':
142    ax_grid = gridspec.GridSpec(nb_rows, nb_col, figure=fig)
143    return fig, None, None, ax_grid
144
145  ## If empty rows or columns, return `None`.
146  if nb_rows == 0 or nb_col == 0:
147    return fig, None, None, None
148
149  ## Return subplots
150  ax_mat = fig.subplots(nb_rows, nb_col, sharex=share_x, sharey=share_y, squeeze=False)
151
152  ## sharex/sharey if the subplot would share same axis, squeeze=False to ensure ax_mat is 2D array
153  ax_arr = ax_mat.flatten()
154  ax = ax_arr[0]
155
156  ## Add grid to the figure
157  fig.grid(True)
158  return fig, ax, ax_arr, ax_mat
159
160def WPC_initialize_plane(angle_type, fig, ax):
161  ## Load image
162  ax.set_axis_off()
163  d_plane = PLANE_DICT[angle_type]
164  img_plane = d_plane['img']
165  im = ax.imshow(img_plane)
166
167  size_plane = d_plane['size']
168  center_plane = d_plane['center']
169  diag_plane = d_plane['diag']
170  alpha_plane = d_plane['alpha']
171  tt = 1.15 ## Tolerance threshold of 20 percent to not have cut circle
172
173  ax.set_xlim(-alpha_plane[1]*tt, (size_plane[1] + alpha_plane[1])*tt)
174  ax.set_ylim(-alpha_plane[0]*tt, (size_plane[0] + alpha_plane[0])*tt)
175
176  ## display circle
177  cc = plt.Circle((center_plane[1], center_plane[0]), diag_plane/2, fill=False, color='black', linewidth=2, linestyle='dashdot')
178  ax.add_artist(cc)
179  return  im
180
181def WPC_showFigure(fig):
182  w, h = fig.get_size_inches()
183  fig.set_size_inches(w, h, forward=True)
184  plt.ion() ## Turn on interactive mode
185  plt.show()
186  return
187
188def WPC_getRotMat(roll, pitch, yaw, use_deg=False):
189  if use_deg:
190    roll *= DEGREE_TO_RADIAN
191    pitch *= DEGREE_TO_RADIAN
192    yaw *= DEGREE_TO_RADIAN
193  rot_mat_x = np.array([[1, 0, 0], [0, np.cos(roll), np.sin(roll)], [0, -np.sin(roll), np.cos(roll)]])
194  rot_mat_y = np.array([[np.cos(pitch), 0, np.sin(pitch)], [0, 1, 0], [-np.sin(pitch), 0, np.cos(pitch)]])
195  rot_mat_z = np.array([[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]])
196  rot_mat = rot_mat_x.dot(rot_mat_y).dot(rot_mat_z) ## Dot product from the back RxRyRz
197  return rot_mat
198
199def WPC_showEmpty(tag=DEFAULT_MODEL, save=0):
200  fig, ax_dict = WPC_initializeFigure(option='3D')
201  im_yaw = WPC_initialize_plane('yaw', fig, ax_dict.get('plane_yaw'))
202  im_pitch = WPC_initialize_plane('pitch', fig, ax_dict.get('plane_pitch'))
203  im_roll = WPC_initialize_plane('roll', fig, ax_dict.get('plane_roll'))
204
205  im_dict = {
206    'yaw': im_yaw,
207    'pitch': im_pitch,
208    'roll': im_roll,
209  }
210
211  d = MODEL_DICT[tag] #model dict dictionary
212
213  ## Load
214  model = mesh.Mesh.from_file(f'{DATA_PATH}{tag}.stl')
215  data_orig = model.vectors ## 3D array
216
217  ## Plot
218  poly = mplot3d.art3d.Poly3DCollection([], color=BMH_BLUE, edgecolor='k', lw=0.2) #or lightsteelblue
219  ax_dict.get('main').add_collection3d(poly)
220
221  ## Settings
222  scale = [-0.6*d['view'], 0.6*d['view']]
223  ax_dict.get('main').view_init(elev=0, azim=0, roll=0)
224  ax_dict.get('main').auto_scale_xyz(scale, scale, scale)
225  ax_dict.get('main').set_axis_off()
226  ax_dict.get('main').set_title(f'Sensor fusion {tag}', weight='bold')
227
228  ## Save
229  fig.set_size_inches(6, 6)
230  fig.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95)
231  fig.canvas.manager.set_window_title('WPC AHRS visualization')
232
233  WPC_showFigure(fig)
234  return fig, ax_dict, im_dict, data_orig, poly
235
236def WPC_draw3DModel(fig, ax, data_orig, poly, roll, pitch, yaw):
237  ## Rotate
238  rot_mat = WPC_getRotMat(roll, pitch, yaw, use_deg=True)
239  data = data_orig.dot(rot_mat)
240
241  ## Plot
242  poly.set_verts(data)
243  return
244
245def WPC_plot_plane(fig, ax, angle_type, im, center):
246  ## Display image
247  im.remove()
248
249  ## Calculate rotation angle in radians
250  angle = np.deg2rad(angle_type)  #Convert in  en radians
251
252  ## Apply rotation to the image
253  rotation_matrix = Affine2D().rotate_deg_around(center[1], center[0], np.rad2deg(angle)+180)
254  im.set_transform(rotation_matrix + ax.transData)
255
256  ## Show the image again
257  ax.add_artist(im)
258  return
259
260def WPC_text_button(fig, Roll, Pitch, Yaw, ax_value_roll, ax_value_pitch, ax_value_yaw):
261  for ax in [ax_value_roll, ax_value_pitch, ax_value_yaw]:
262    ax.clear()
263    ax.axis('off')
264    ax.add_patch(plt.Rectangle((0, 0), 1, 1, facecolor='#ebebeb', transform=ax.transAxes, zorder=-1))
265    ax.grid(False)
266    roll = "{:7.2f}".format(Roll)
267    pitch = "{:7.2f}".format(Pitch)
268    yaw = "{:7.2f}".format(Yaw)
269    ax_value_roll.text(0.5, 0.5, f' Roll:\n{roll} deg', fontsize=32, weight='bold', color=BMH_BLUE,ha='center', va='center')
270    ax_value_pitch.text(0.5, 0.5, f' Pitch:\n{pitch} deg', fontsize=32, weight='bold', color=BMH_BURGUNDY, ha='center', va='center')
271    ax_value_yaw.text(0.5, 0.5, f' Yaw:\n{yaw} deg', fontsize=32, weight='bold', color=BMH_PURPLE, ha='center', va='center')
272  return
273
274async def main():
275    ## Get Python driver version
276    print(f'{pywpc.PKG_FULL_NAME} - Version {pywpc.__version__}')
277
278    ## Create device handle
279    dev = pywpc.WifiDAQE3AH()
280
281    ## Show empty
282    fig, ax_dict, im_dict, data_orig, poly = WPC_showEmpty()
283
284    ## Connect to device
285    try:
286      dev.connect("192.168.5.38") ## Depend on your device
287    except Exception as err:
288      pywpc.printGenericError(err)
289      ## Release device handle
290      dev.close()
291      return
292
293    try:
294      ## Parameters setting
295      port = 0 ## Depend on your device
296      mode = 0 ## 0: Orientation, 1: Acceleration, 2: Orientation + Acceleration
297
298      ## Get firmware model & version
299      driver_info = await dev.Sys_getDriverInfo_async()
300      print("Model name: " + driver_info[0])
301      print("Firmware version: " + driver_info[-1])
302
303      ## Open AHRS and update rate is 333 HZ
304      err = await dev.AHRS_open_async(port)
305      print(f"AHRS_open_async in port {port}, status: {err}")
306
307      ## Start AHRS
308      err = await dev.AHRS_start_async(port)
309      print(f"AHRS_start_async in port {port}, status: {err}")
310
311      while plt.fignum_exists(fig.number):
312          ahrs_list = await dev.AHRS_getEstimate_async(port, mode)
313          if len(ahrs_list) > 0:
314              WPC_draw3DModel(fig, ax_dict.get('main'), data_orig, poly, ahrs_list[0], ahrs_list[1], ahrs_list[2])
315              WPC_plot_plane(fig, ax_dict.get('plane_roll'), ahrs_list[0], im_dict.get('roll'), CENTER_ROLL)
316              WPC_plot_plane(fig, ax_dict.get('plane_pitch'), ahrs_list[1], im_dict.get('pitch'), CENTER_PITCH)
317              WPC_plot_plane(fig, ax_dict.get('plane_yaw'), ahrs_list[2], im_dict.get('yaw'), CENTER_YAW)
318              WPC_text_button(fig, ahrs_list[0], ahrs_list[1], ahrs_list[2], ax_dict.get('value_roll'), ax_dict.get('value_pitch'),  ax_dict.get('value_yaw'))
319              plt.tight_layout()
320              plt.pause(2**-5)
321
322    except KeyboardInterrupt:
323      print("Press keyboard")
324
325    except Exception as err:
326      pywpc.printGenericError(err)
327
328    finally:
329      ## Stop AHRS
330      err = await dev.AHRS_stop_async(port)
331      print(f"AHRS_stop_async in port {port}, status: {err}")
332
333      ## Close AHRS
334      err = await dev.AHRS_close_async(port)
335      print(f"AHRS_close_async in port {port}, status: {err}")
336
337      ## Disconnect device
338      dev.disconnect()
339
340      ## Release device handle
341      dev.close()
342    return
343def main_for_spyder(*args):
344    if asyncio.get_event_loop().is_running():
345        return asyncio.create_task(main(*args)).result()
346    else:
347        return asyncio.run(main(*args))
348if __name__ == '__main__':
349    asyncio.run(main()) ## Use terminal
350    # await main() ## Use Jupyter or IPython(>=7.0)
351    # main_for_spyder() ## Use Spyder