bexus-26-imu-fusion-flight-software/groundstation/vispyHelper.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-

from PyQt5.QtWidgets import *
from vispy import app, visuals, scene, io
import numpy as np
from colorsys import hsv_to_rgb

PlotAxis = scene.visuals.create_visual_node(visuals.AxisVisual)
Plot3DLine = scene.visuals.create_visual_node(visuals.LinePlotVisual)
Plot3DSurface = scene.visuals.create_visual_node(visuals.SurfacePlotVisual)
PlotCube = scene.visuals.create_visual_node(visuals.CubeVisual)
PlotMesh = scene.visuals.create_visual_node(visuals.MeshVisual)


vertices, faces, normals, texcoords = io.read_mesh("gondel2.obj")
arrowVert = np.array([[0,0,0],[10,0,0],[0,10,0],[0,0,10]])
arrowConn = np.array([[0,1],[0,2],[0,3]])
arrows = np.array([[0,0,0,10,0,0],[0,0,0,0,10,0],[0,0,0,0,0,10]])

#vertex_colors = np.random.random(12)
vertex_colors = np.linspace(0, 1, len(vertices))
vertex_colors = np.array([hsv_to_rgb(c, 1, 1) for c in vertex_colors])

class viewQuad():
	def __init__(self, object, follow = True):
		self.follow = follow
		self.canvas = scene.SceneCanvas(keys='interactive', show=False)
		self.grid = self.canvas.central_widget.add_grid(margin=10)
		self.views = []
		
		self.views.append(viewOrtho(self.grid.add_grid(row=0, col=0), 3, self.follow))
		self.views.append(viewOrtho(self.grid.add_grid(row=1, col=0), 2, self.follow))
		self.views.append(viewOrtho(self.grid.add_grid(row=1, col=1), 1, self.follow))
		self.views.append(viewRot(self.grid.add_grid(row=0,col=1), self.follow))
		
		for view in self.views:
			view.grid.border_color = (0.5, 0.5, 0.5, 1)
		
		object.setLayout(QVBoxLayout())
		object.layout().addWidget(self.canvas.native)
	
	def update(self, data3d, rot):
		for view in self.views:
			view.update(data3d, rot)
		
class view3D():

	def __init__(self, object, follow = False):
		self.canvas = scene.SceneCanvas(keys='interactive', show=False)
		self.view = self.canvas.central_widget.add_view()
		self.view.camera = scene.cameras.FlyCamera(parent=self.view.scene, fov=60, name='Fly')
		
		size = 100000
		s = np.linspace(-size, size, 50)
		YY, XX = np.meshgrid(s, s)
		z = -0.1 * np.sqrt(YY*YY + XX*XX)
		c=np.zeros(shape=(50,50,4))
		
		for x in range(50):
			for y in range(50):
				if x==25 and y>=25:
					c[x][y] = np.array([0,1,0,1])
				elif x>=25 and y==25:
					c[x][y] = np.array([1,0,0,1])
				else:
					c[x][y] = np.array([0,0.3,1,0.5])
		
		#print(c)
		# colors=c,
		self.groundPlane = Plot3DSurface(x=s, y=s, z=z, parent=self.view.scene)
		
		self.dataPlot = Plot3DLine(np.array([[0,0,0]], np.float32), parent=self.view.scene)
		self.dataPlot2 = Plot3DLine(np.array([[0,0,0]], np.float32), parent=self.view.scene)

		self.carPlot = PlotMesh(vertices, faces, shading='smooth',
		vertex_colors=vertex_colors, parent=self.view.scene)
		#self.carPlot = PlotCube((1,1.5,0.5), edge_color=(1,1,1,1), face_colors=vertex_colors, parent=self.view.scene)
		self.carPlot.transform = visuals.transforms.MatrixTransform()
	
		
		object.setLayout(QVBoxLayout())
		object.layout().addWidget(self.canvas.native)
		self.follow = follow
		
	def update(self, data3d, rot):
		rotX, rotY, rotZ = rot
		self.dataPlot.set_data(data3d[0], width=2.0, color='red',	marker_size=0,face_color=(0.2, 0.2, 1, 0.8))
		self.dataPlot2.set_data(data3d[1], width=2.0, color='green',	marker_size=0,face_color=(0.2, 0.2, 1, 0.8))
		self.carPlot.transform.reset()
		self.carPlot.transform.rotate(rotX, (1,0,0))
		self.carPlot.transform.rotate(rotY, (0,1,0))
		self.carPlot.transform.rotate(rotZ, (0,0,1))
		
		self.carPlot.transform.translate(data3d[0][-1])
		if self.follow:
			self.view.camera.center = tuple(np.add(data3d[0][-1], [5,5,5]))
		

class viewRot():

	def __init__(self, grid, follow = True):
		self.grid = grid
		self.view = self.grid.add_view(row=0,col=0)
		self.view.camera = scene.cameras.TurntableCamera(parent=self.view.scene)
		
		self.dataPlot = Plot3DLine(np.array([[0,0,0]], np.float32), parent=self.view.scene)
		self.dataPlot2 = Plot3DLine(np.array([[0,0,0]], np.float32), parent=self.view.scene)
		
		#self.gridlines = scene.visuals.GridLines(scale=(10,10), color=(0,1,0), parent=self.view.scene)

		self.carPlot = PlotMesh(vertices, faces,
		vertex_colors=vertex_colors, shading='smooth', parent=self.view.scene)
		self.carPlot.transform = visuals.transforms.MatrixTransform()
		
		self.follow = follow
		
	def update(self, data3d, rot):
		rotX, rotY, rotZ = rot
		self.dataPlot.set_data(data3d[0], width=2.0, color='red', marker_size=0)
		self.dataPlot2.set_data(data3d[1], width=2.0, color='green', marker_size=0)
		print(data3d[1])
		self.carPlot.transform.reset()
		self.carPlot.transform.rotate(rotX, (1,0,0))
		self.carPlot.transform.rotate(rotY, (0,1,0))
		self.carPlot.transform.rotate(rotZ, (0,0,1))
		self.carPlot.transform.translate(data3d[0][-1])

		if self.follow:
			self.view.camera.center = tuple(data3d[0][-1])
			
class viewOrtho():
	def __init__(self, grid, up=3, follow = True):
		self.follow = follow
		self.up = up
		self.grid = grid
		
		if self.up == 3: #top
			la='Y'
			lb='X'
			ti="Top"
		elif self.up == 2:
			la='Y'
			lb='Z'
			ti="Front"
			
		else:
			la='X'
			lb='Z'
			ti="Side"
		
		title = scene.Label(ti, color='white')
		title.height_max = 40
		self.grid.add_widget(title, row=0, col=0, col_span=2)
		
		yaxis = scene.AxisWidget(orientation='left',
												 axis_label=lb+' Axis',
												 axis_font_size=12,
												 axis_label_margin=50,
												 tick_label_margin=5)
		yaxis.width_max = 80
		self.grid.add_widget(yaxis, row=1,col=0)

		xaxis = scene.AxisWidget(orientation='bottom',
												 axis_label=la+' Axis',
												 axis_font_size=12,
												 axis_label_margin=50,
												 tick_label_margin=5)

		xaxis.height_max = 80
		self.grid.add_widget(xaxis, row=2, col=1)
		
		right_padding = self.grid.add_widget(row=1, col=2, row_span=1)
		right_padding.width_max = 50
		
		self.view = self.grid.add_view(row=1,col=1)
		
		self.view.camera = scene.cameras.PanZoomCamera(rect=(-5,-5,5,5), aspect=1, parent=self.view.scene)#flip=(0,1,0),

		
		yaxis.link_view(self.view)
		xaxis.link_view(self.view)
		
		self.gridlines = scene.visuals.GridLines(scale=(1,1), color=(0,1,0), parent=self.view.scene)
		
		self.dataPlot = Plot3DLine(np.array([[0,0]], np.float32), parent=self.view.scene)
		self.dataPlot2 = Plot3DLine(np.array([[0,0]], np.float32), parent=self.view.scene)

		self.carPlot = PlotMesh(vertices, faces,
		vertex_colors=vertex_colors, shading='smooth', parent=self.view.scene)
		self.carPlot.transform = visuals.transforms.MatrixTransform()
		
	def update(self, data3d, rot):
		rotX, rotY, rotZ = rot
		if self.up == 3: #top
			data2d = data3d[0][:,:2]    # X Y 
			data2d2 = data3d[1][:,:2]    # X Y 
		elif self.up == 2: #front
			data2d = data3d[0][:,0:3:2] # X Z
			data2d2 = data3d[1][:,0:3:2] # X Z
		else: #side
			data2d = data3d[0][:,1:]    # Y Z
			data2d2 = data3d[1][:,1:]    # Y Z
		
		self.dataPlot.set_data(data2d, width=2.0, color='red', marker_size=0)
		self.dataPlot2.set_data(data2d2, width=2.0, color='green', marker_size=0)
		self.carPlot.transform.reset()
		self.carPlot.transform.set_ortho(-1, 1, -1, 1, -1, 1)
		
		rotX, rotY, rotZ = (-rotX), (-rotY), (rotZ)
		
		self.carPlot.transform.rotate(rotX, (1,0,0))
		self.carPlot.transform.rotate(rotY, (0,1,0))
		self.carPlot.transform.rotate(rotZ, (0,0,1))
		
		if self.up <= 2:#front
			self.carPlot.transform.rotate(90, (1,0,0))
		
		if self.up == 1:#side
			self.carPlot.transform.rotate(90, (0,1,0))

		self.carPlot.transform.translate(data2d[-1])
		if self.follow:
			self.view.camera.center = tuple(data2d[-1])