diff --git a/skimage/draw/_draw.pyx b/skimage/draw/_draw.pyx index c25183a51af..0d384a38ca1 100644 --- a/skimage/draw/_draw.pyx +++ b/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): diff --git a/skimage/draw/draw.py b/skimage/draw/draw.py index 498183ae833..ce4e68fcb64 100644 --- a/skimage/draw/draw.py +++ b/skimage/draw/draw.py @@ -125,7 +125,7 @@ def circle(r, c, radius, shape=None): return ellipse(r, c, radius, radius, shape) -def set_color(img, coords, color, alpha=None): +def set_color(img, coords, color, alpha=1): """Set pixel color in the image at the given coordinates. Coordinates that exceed the shape of the image will be ignored. @@ -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 diff --git a/skimage/draw/tests/test_draw.py b/skimage/draw/tests/test_draw.py index bd295f17b00..214f09b0102 100644 --- a/skimage/draw/tests/test_draw.py +++ b/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)