Clean up implementation

skimage/draw/_draw.pyx

@@ -35,10 +35,7 @@ def _coords_inside_image(rr, cc, shape, val=None):
     if val is None:
         return rr[mask], cc[mask]
     else:
-        if np.isscalar(val):
-           return rr[mask], cc[mask], val
-        else:
-           return rr[mask], cc[mask], val[mask]
+        return rr[mask], cc[mask], val[mask]
 
 
 def line(Py_ssize_t y0, Py_ssize_t x0, Py_ssize_t y1, Py_ssize_t x1):

skimage/draw/draw.py

@@ -140,7 +140,7 @@ def set_color(img, coords, color, alpha=None):
         Color to be assigned to coordinates in the image.
     alpha : scalar or (N,) ndarray
         Alpha values used to blend color with image.  0 is transparent,
-        1 is opaque.
+        1 is opaque (default).
 
     Returns
     -------
@@ -168,36 +168,28 @@ def set_color(img, coords, color, alpha=None):
     """
     rr, cc = coords
 
+    if img.ndim == 2:
+        img = img[..., np.newaxis]
+
+    color = np.array(color, ndmin=1, copy=False)
+
+    if img.shape[-1] != color.shape[-1]:
+        raise ValueError('Color shape ({}) must match last '
+                         'image dimension ({}).'.format(color.shape[0],
+                                                        img.shape[-1]))
+
     if alpha is None:
+        alpha = 1
+
+    if np.isscalar(alpha):
         rr, cc = _coords_inside_image(rr, cc, img.shape)
+        alpha = np.ones_like(rr) * alpha
     else:
         rr, cc, alpha = _coords_inside_image(rr, cc, img.shape, val=alpha)
 
-    if not np.isscalar(color):
-        color = np.array(color)
-        if color.shape[0] != img.shape[-1]:
-            raise ValueError('Color shape ({}) must match last '
-                             'image dimension ({}).'.format(color.shape[0],
-                                                            img.shape[-1]))
+    alpha = alpha[..., np.newaxis]
 
-    if alpha is None:
-        img[rr, cc] = color
-    else:
-        if np.isscalar(alpha) or np.isscalar(color):
-            color = color * alpha
-        else:
-            color = color * alpha[:, None]
-
-
-        # Strategy: try operation directly.  If scalar / correctly shaped
-        # vector, this will work.  If not (e.g. vector alpha but color image),
-        # try broadcasting.
-        try:
-            vals = img[rr, cc] * (1 - alpha)
-        except ValueError:
-            vals = img[rr, cc] * (1 - alpha[:, None])
-
-        try:
-            img[rr, cc] = vals + color
-        except ValueError:
-            img[rr, cc] = vals + color[:, None]
+    color = color * alpha
+    vals = img[rr, cc] * (1 - alpha)
+
+    img[rr, cc] = vals + color

skimage/draw/tests/test_draw.py

@@ -31,9 +31,6 @@ def test_set_color_with_alpha():
 
     img = np.zeros((10, 10, 3))
 
-    rr, cc, alpha = line_aa(0, 0, 0, 30)
-    set_color(img, (rr, cc), 1, alpha=alpha)
-
     rr, cc, alpha = line_aa(0, 0, 0, 30)
     set_color(img, (rr, cc), (1, 0, 0), alpha=alpha)