Added color support to shapes and methods

jhavl · jhavl · commit 4e3201b033af · 2021-02-15T16:19:47.000+10:00
diff --git a/examples/gimbal.py b/examples/gimbal.py
@@ -0,0 +1,80 @@
+# #!/usr/bin/env python
+# """
+# @author Jesse Haviland
+# """
+
+import roboticstoolbox as rtb
+import spatialmath as sm
+import numpy as np
+import pathlib
+import os
+
+path = os.path.realpath('.')
+
+# Launch the simulator Swift
+env = rtb.backends.Swift()
+env.launch()
+
+path = pathlib.Path(path) / 'roboticstoolbox' / 'data'
+
+g1 = rtb.Mesh(
+    filename=str(path / 'gimbal-ring1.stl'),
+    base=sm.SE3(0, 0, 1.3),
+    color=[34, 143, 201]
+)
+# g1.v = [0, 0, 0, 0.4, 0, 0]
+
+g2 = rtb.Mesh(
+    filename=str(path / 'gimbal-ring2.stl'),
+    base=sm.SE3(0, 0, 1.3),
+    color=[31, 184, 72]
+)
+# g2.v = [0, 0, 0, 0.4, 0.0, 0]
+
+g3 = rtb.Mesh(
+    filename=str(path / 'gimbal-ring3.stl'),
+    base=sm.SE3(0, 0, 1.3),
+    color=[240, 103, 103]
+)
+# g3.v = [0, 0, 0, 0.4, 0, 0]
+
+env.add(g1)
+env.add(g2)
+env.add(g3)
+
+
+def set_one(x):
+    g1.base = sm.SE3(0, 0, 1.3) * sm.SE3.Rz(np.deg2rad(float(x)))
+
+def set_two(x):
+    g2.base = sm.SE3(0, 0, 1.3) * sm.SE3.Ry(np.deg2rad(float(x)))
+
+def set_three(x):
+    g3.base = sm.SE3(0, 0, 1.3) * sm.SE3.Rx(np.deg2rad(float(x)))
+    # g2.base = sm.SE3(0, 0, 1.3) * sm.SE3.Rz(np.pi/2) * sm.SE3.Ry(-np.deg2rad(float(x)))
+    # g2.base = g3.base 
+
+
+env.add_slider(
+    set_one,
+    min=-180, max=180,
+    step=1, value=0,
+    desc='Ring One')
+
+env.add_slider(
+    set_two,
+    min=-180, max=180,
+    step=1, value=0,
+    desc='Ring Two')
+
+env.add_slider(
+    set_three,
+    min=-180, max=180,
+    step=1, value=0,
+    desc='Ring Three')
+
+
+
+while(True):
+    env.process_events()
+    env.step(0)
diff --git a/roboticstoolbox/robot/Shape.py b/roboticstoolbox/robot/Shape.py
@@ -162,14 +162,7 @@ def color(self, value):
 
         default_color = (0.95, 0.5, 0.25, 1.0)
 
-        if isinstance(value, tuple):
-            value = list(value)
-
-        if isinstance(value, list):
-            if len(value) == 3:
-                value.append(1.0)
-                value = tuple(value)
-        elif isinstance(value, str):
+        if isinstance(value, str):
             if _mpl:
                 try:
                     value = mpc.to_rgba(value)
@@ -183,6 +176,17 @@ def color(self, value):
                     'Install using: pip install matplotlib')
         elif value is None:
             value = default_color
+        else:
+
+            value = np.array(value)
+
+            if np.any(value > 1.0):
+                value = value / 255.0
+
+            if value.shape[0] == 3:
+                value = np.r_[value, 1.0]
+
+            value = tuple(value)
 
         self._color = value
 
@@ -324,10 +328,10 @@ class Mesh(Shape):
 
     """
 
-    def __init__(self, filename=None, scale=[1, 1, 1], base=None):
+    def __init__(self, filename=None, scale=[1, 1, 1], base=None, color=None):
         super(Mesh, self).__init__(
             filename=filename, base=base,
-            scale=scale, stype='mesh')
+            scale=scale, stype='mesh', color=color)
 
     def _init_pob(self):
         name, file_extension = os.path.splitext(self.filename)
@@ -359,10 +363,10 @@ class Cylinder(Shape):
 
     """
 
-    def __init__(self, radius, length, base=None):
+    def __init__(self, radius, length, base=None, color=None):
         super(Cylinder, self).__init__(
             base=base, radius=radius, length=length,
-            stype='cylinder')
+            stype='cylinder', color=color)
 
     def _init_pob(self):
         col = p.createCollisionShape(
@@ -388,9 +392,9 @@ class Sphere(Shape):
 
     """
 
-    def __init__(self, radius, base=None):
+    def __init__(self, radius, base=None, color=None):
         super(Sphere, self).__init__(
-            base=base, radius=radius, stype='sphere')
+            base=base, radius=radius, stype='sphere', color=color)
 
     def _init_pob(self):
         col = p.createCollisionShape(
@@ -415,8 +419,9 @@ class Box(Shape):
 
     """
 
-    def __init__(self, scale, base=None):
-        super(Box, self).__init__(base=base, scale=scale, stype='box')
+    def __init__(self, scale, base=None, color=None):
+        super(Box, self).__init__(
+            base=base, scale=scale, stype='box', color=color)
 
     def _init_pob(self):
 
