drawRobotics/drawRobotics.py at master · PinkWink/drawRobotics · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import matplotlib.pyplot as plt
import numpy as np
from Arrow3D import *

conv2Rad = lambda x : x*np.pi/180

def drawVector(fig, pointA, pointB, **kwargs):
	ms = kwargs.get('mutation_scale', 20)
	ars = kwargs.get('arrowstyle', '-|>')
	lc = kwargs.get('lineColor', 'k')
	pc = kwargs.get('projColor', 'k')
	pointEnable = kwargs.get('pointEnable', True)
	projOn = kwargs.get('proj', True)
	lineStyle = kwargs.get('lineStyle', '-')
	annotationString = kwargs.get('annotationString', '')
	lineWidth = kwargs.get('lineWidth', 1)

	if (3 <= pointA.size <= 4):
		xs = [pointA[0], pointB[0]]
		ys = [pointA[1], pointB[1]]
		zs = [pointA[2], pointB[2]]
	else:
		xs = [pointA[0,3], pointB[0,3]]
		ys = [pointA[1,3], pointB[1,3]]
		zs = [pointA[2,3], pointB[2,3]]

	out = Arrow3D(xs, ys, zs, mutation_scale=ms, arrowstyle=ars, color=lc, linestyle=lineStyle, linewidth=lineWidth)
	fig.add_artist(out)

	if pointEnable: fig.scatter(xs[1], ys[1], zs[1], color='k', s=50)

	if projOn:
		fig.plot(xs, ys, [0, 0], color=pc, linestyle='--')
		fig.plot([xs[0], xs[0]], [ys[0], ys[0]], [0, zs[0]], color=pc, linestyle='--')
		fig.plot([xs[1], xs[1]], [ys[1], ys[1]], [0, zs[1]], color=pc, linestyle='--')

	if annotationString != '':
		fig.text(xs[1], ys[1], zs[1], annotationString, size=15, zorder=1, color='k')

def drawPointWithAxis(fig, *args, **kwargs):
	ms = kwargs.get('mutation_scale', 20)
	ars = kwargs.get('arrowstyle', '->')
	pointEnable = kwargs.get('pointEnable', True)
	axisEnable = kwargs.get('axisEnable', True)
	lineStyle = kwargs.get('lineStyle', '-')
	lineWidth = kwargs.get('lineWidth', 1)
	vectorLength = kwargs.get('vectorLength', 1)

	if len(args) == 4:
		ORG = args[0]
		hat_X = args[1]
		hat_Y = args[2]
		hat_Z = args[3]
		xs_n = [ORG[0], ORG[0] + hat_X[0]*vectorLength]
		ys_n = [ORG[1], ORG[1] + hat_X[1]*vectorLength]
		zs_n = [ORG[2], ORG[2] + hat_X[2]*vectorLength]
		xs_o = [ORG[0], ORG[0] + hat_Y[0]*vectorLength]
		ys_o = [ORG[1], ORG[1] + hat_Y[1]*vectorLength]
		zs_o = [ORG[2], ORG[2] + hat_Y[2]*vectorLength]
		xs_a = [ORG[0], ORG[0] + hat_Z[0]*vectorLength]
		ys_a = [ORG[1], ORG[1] + hat_Z[1]*vectorLength]
		zs_a = [ORG[2], ORG[2] + hat_Z[2]*vectorLength]
	else:
		tmp = args[0]
		ORG = tmp[:3,3:]
		hat_X = tmp[:3,0:1]
		hat_Y = tmp[:3,1:2]
		hat_Z = tmp[:3,2:3]
		xs_n = [ORG[0, 0], ORG[0, 0] + hat_X[0, 0]*vectorLength]
		ys_n = [ORG[1, 0], ORG[1, 0] + hat_X[1, 0]*vectorLength]
		zs_n = [ORG[2, 0], ORG[2, 0] + hat_X[2, 0]*vectorLength]
		xs_o = [ORG[0, 0], ORG[0, 0] + hat_Y[0, 0]*vectorLength]
		ys_o = [ORG[1, 0], ORG[1, 0] + hat_Y[1, 0]*vectorLength]
		zs_o = [ORG[2, 0], ORG[2, 0] + hat_Y[2, 0]*vectorLength]
		xs_a = [ORG[0, 0], ORG[0, 0] + hat_Z[0, 0]*vectorLength]
		ys_a = [ORG[1, 0], ORG[1, 0] + hat_Z[1, 0]*vectorLength]
		zs_a = [ORG[2, 0], ORG[2, 0] + hat_Z[2, 0]*vectorLength]

	if pointEnable: fig.scatter(xs_n[0], ys_n[0], zs_n[0], color='k', s=50)

	if axisEnable:
		n = Arrow3D(xs_n, ys_n, zs_n, mutation_scale=ms, arrowstyle=ars, color='r', linestyle=lineStyle, linewidth=lineWidth)
		o = Arrow3D(xs_o, ys_o, zs_o, mutation_scale=ms, arrowstyle=ars, color='g', linestyle=lineStyle, linewidth=lineWidth)
		a = Arrow3D(xs_a, ys_a, zs_a, mutation_scale=ms, arrowstyle=ars, color='b', linestyle=lineStyle, linewidth=lineWidth)
		fig.add_artist(n)
		fig.add_artist(o)
		fig.add_artist(a)

def RotX(phi):
    return np.array([[1, 0, 			0],
                     [0, np.cos(phi), 	-np.sin(phi)],
                     [0, np.sin(phi), 	np.cos(phi)]])

def RotY(theta):
    return np.array([[np.cos(theta), 	0, np.sin(theta)],
                     [0, 				1, 0],
                     [-np.sin(theta), 	0, np.cos(theta)]])

def RotZ(psi):
    return np.array([[np.cos(psi), 	-np.sin(psi), 	0],
                     [np.sin(psi), 	np.cos(psi), 	0],
                     [0, 			0, 				1]])

def D_q(dx, dy, dz):
	return np.array([[1, 0, 0, dx],
					 [0, 1, 0, dy],
					 [0, 0, 1, dz],
					 [0, 0, 0, 1]])