Merge pull request #8 from jorenham/development

Version 0.2

Author: jorenham

CHANGELOG.md

@@ -3,6 +3,11 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [0.2] - 2019-09-29
+### Added
+- The `Turtle` class now has color support.
+- The examples work with RGBA images now.
+
 ## [0.1] - 2018-06-26
 ### Added
 - The `Turtle` class as wrapper for a two-dimensional (grayscale) NumPy array. 

numpy_turtle/examples/fractal_plant.py

@@ -1,5 +1,6 @@
 import numpy as np
-from scipy.misc import toimage
+from matplotlib import pyplot as plt
+from skimage.io import imsave
 
 from numpy_turtle import Turtle, l_system
 
@@ -17,12 +18,13 @@ def main():
     padding = 32
     n = 6
 
-    a = np.zeros((rows, cols))
+    a = np.zeros((rows, cols, 4))
     s = l_system.grow(axiom, rules, n)
 
-    t = Turtle(a)
+    t = Turtle(a, aa=True)
     t.position = rows - padding, padding
     t.rotate(np.pi - angle)
+    t.color = (0, 1, 0, 1)
 
     for s_n in s:
         if s_n == 'F':
@@ -36,7 +38,7 @@ def main():
         elif s_n == ']':
             t.pop()
 
-    toimage(a).show()
+    imsave('images/fractal_plant.png', a)
 
 
 if __name__ == '__main__':

numpy_turtle/examples/images/fractal_plant.png

@@ -1,5 +1,5 @@
 import numpy as np
-from scipy.misc import toimage
+from skimage.io import imsave
 
 from numpy_turtle import Turtle, l_system
 
@@ -20,12 +20,13 @@ def main():
     rows = int(np.ceil(cols * np.sin(angle / 2)))
     n = 8
 
-    a = np.zeros((rows, cols))
+    a = np.zeros((rows, cols, 4))
     s = l_system.grow(axiom, rules, n)
 
     t = Turtle(a)
     t.position = rows, 0
     t.rotate(np.pi / 2)
+    t.color = (0, 0, 0, 1)
 
     for s_n in s:
         if s_n == 'F' or s_n == 'G':
@@ -35,7 +36,7 @@ def main():
         elif s_n == '+':
             t.rotate(-angle)
 
-    toimage(a).show()
+    imsave('images/sierpinski_triangle.png', a)
 
 
 if __name__ == '__main__':

numpy_turtle/examples/images/sierpinski_triangle.png

@@ -1,13 +1,15 @@
-from typing import Tuple
+from typing import Tuple, Union
 
 import numpy as np
 from skimage.draw import line, line_aa
 
 _TAU = np.pi * 2
 
+Color = Union[int, float, Tuple[int, ...], Tuple[float, ...]]
+
 
 class Turtle:
-    def __init__(self, array, deg: bool=False, aa: bool=False):
+    def __init__(self, array: np.ndarray, deg: bool = False, aa: bool = False):
         """Draw on a NumPy array using turtle graphics.
 
         Starts at (0, 0) (top-left corner) with a direction of 0 (pointing
@@ -16,18 +18,23 @@ def __init__(self, array, deg: bool=False, aa: bool=False):
         Parameters
         ----------
         array: np.ndarray
-            The 2D array to write to
+            The 2D array to write to. Can be either of shape h x w (grayscale),
+            h x w x c (e.g. rgb for c=3 channels).
+            The dtype is used to determine the color depth of each channel:
+
+              * `bool` for 2 colors.
+              * All `np.integer` subtypes for discrete depth, ranging from 0 to
+                its max value (e.g. `np.uint8` for values in 0 - 255).
+              * All `np.floating` subtypes for continuous depth, ranging from 0
+                to 1.
+
         deg : :obj:`bool`, optional
             Use degrees instead of radians.
         aa : :obj:`bool`, optional
             Enable anti-aliasing.
         """
-        assert type(array) is np.ndarray, 'Array should be a NumPy ndarray'
-        assert array.ndim == 2, 'Only 2D arrays are supported'
-        assert (
-                np.issubdtype(array.dtype, np.integer) or
-                np.issubdtype(array.dtype, np.floating)
-        ), '{} is unsupported'.format(array.dtype)
+        if type(array) is not np.ndarray:
+            raise TypeError('Array should be a NumPy ndarray')
 
         self.array = array
         self.aa = aa
@@ -37,10 +44,32 @@ def __init__(self, array, deg: bool=False, aa: bool=False):
         self.__r, self.__c = 0, 0
         self.__stack = []
 
-        if np.issubdtype(array.dtype, np.integer):
-            self.__color = np.iinfo(array.dtype).max
+        if array.ndim == 2:
+            self.__channels = 1
+        elif array.ndim == 3:
+            self.__channels = array.shape[2]
+        else:
+            raise TypeError('Array does not have 2 or 3 dimensions')
+
+        if array.dtype == np.dtype(bool):
+            self.__depth = 1
+            self.__dtype = bool
+        elif np.issubdtype(array.dtype, np.integer):
+            self.__depth = np.iinfo(array.dtype).max
+            self.__dtype = int
         elif np.issubdtype(array.dtype, np.floating):
-            self.__color = np.finfo(array.dtype).max
+            self.__depth = 1.0
+            self.__dtype = float
+        else:
+            raise TypeError(
+                'Array should have a bool, int-like, or float-like dtype'
+            )
+
+        # color initially the max depth (white).
+        if self.__channels == 1:
+            self.__color = self.__depth
+        else:
+            self.__color = np.full(self.__channels, self.__depth, self.__dtype)
 
     def __in_array(self, r=None, c=None):
         r = self.__r if r is None else r
@@ -58,10 +87,15 @@ def __draw_line(self, new_c, new_r):
 
         if self.aa:
             rr, cc, val = line_aa(r0, c0, r1, c1)
-            self.array[rr, cc] = (val / 255 * self.__color).astype(self.array.dtype)
         else:
             rr, cc = line(r0, c0, r1, c1)
-            self.array[rr, cc] = self.__color
+            val = 1
+
+        if self.__channels == 1:
+            self.array[rr, cc] = val * self.__color
+        else:
+            for c in range(self.__channels):
+                self.array[rr, cc, c] = val * self.__color[c]
 
     def forward(self, distance: float):
         """Move in the current direction and draw a line with Euclidian
@@ -123,10 +157,19 @@ def position(self, rc: Tuple[int, int]):
         self.__r, self.__c = rc
 
     @property
-    def color(self) -> float:
-        """float: Grayscale color"""
-        return self.__color
+    def color(self) -> Color:
+        """int, float, tuple of int or tuple of float: Grayscale color"""
+        if self.__channels == 1:
+            return self.__color
+        else:
+            return tuple(self.__color)
 
     @color.setter
-    def color(self, c: float):
-        self.__color = c
+    def color(self, c: Color):
+        if not np.isscalar(c) and len(c) != self.__channels:
+            raise TypeError('Invalid amount of color values')
+        for _c in [c] if np.isscalar(c) else c:
+            if _c < 0 or _c > self.__depth:
+                raise ValueError('Color value out of range')
+
+        self.__color = np.array(c, dtype=self.__dtype)

numpy_turtle/examples/sierpinski_triangle.py

@@ -9,7 +9,7 @@
 class TestTurtleUint8(TestCase):
     def setUp(self):
         self.size = 10
-        self.array = np.zeros((self.size, self.size), dtype=np.uint8)
+        self.array = np.zeros((self.size, self.size, 3), dtype=np.uint8)
         self.turtle = Turtle(self.array)
 
     def test_pos_forward(self):
@@ -39,5 +39,9 @@ def test_eye(self):
         self.turtle.rotate(0.25 * np.pi)
         self.turtle.forward(np.sqrt(2 * 10**2))
 
-        eye = np.eye(self.size, dtype=self.array.dtype) * self.turtle.color
-        npt.assert_array_equal(self.array, eye)
+        eye = np.eye(self.size, dtype=self.array.dtype)
+        eye3 = np.zeros_like(self.array)
+        for i, c in enumerate(self.turtle.color):
+            eye3[:, :, i] = eye * c
+
+        npt.assert_array_equal(self.array, eye3)

numpy_turtle/numpy_turtle.py

@@ -2,7 +2,7 @@
 
 setup(
     name='numpy-turtle',
-    version='0.1',
+    version='0.2',
     packages=['numpy_turtle'],
     python_requires='>=3.5',
     install_requires=['numpy>=1.13.1', 'scikit_image>=0.13.1'],