Triangle rasterization using a scanline algorithm for numpy array indexing

Question

I need to access those elements of a large numpy array that lie in a small triangular area. The brute force solutions of using matplotlib.path.contains_points or ImageDraw.Draw.Polygon are too slow for my liking because they return a boolean array instead of a list of indices for accessing the relevant elements.

So instead I've decided to implement a scanline algorithm that rasterizes the triangle and returns an array of the indices of the elements inside the triangle. I used the excellent description of the "Standard Algorithm" on the Sunshine "Triangle Rasterization" page, with three changes:

• The vertices of the triangle are not part of the returned array (because I don't want to use those vertices in the subsequent processing steps)
• the points on all edges of the triangle are part of the returned array. This means that if I process two neighbouring triangles, the edge points would be part of both triangles. I do this because in my case it's important to not miss any points, so I'll rather process a couple of points twice.
• Instead of using two separate functions for flat bottom and flat top triangles, I used a single function that either moves up or down.

The code works, but I'm not sure that this is a particularly elegant implementation. I have two main questions:

• Should I be using numpy arrays and associated functions like argsort for small structures like the triangle vertices, or would it be more efficient (and better practice) to use tuples and lists for this?
• I used nested for loops for the "heart" of the algorithm. I know that numpy only shines when loops are vectorized, but in this case I wouldn't know how to go about that. Is it advisable to vectorize this loop, and if yes, how do I do it?
• Is there a more elegant way to order the two horizontal vertices in rasterize_horizontal_triangle than the tri[1, 0], tri[2, 0] = tri[2, 0], tri[1, 0] expression I'm currently using?

I'm not a very experienced programmer, let alone a professional one, so I'd welcome any and all feedback on the code, be it about the style, the design or the efficiency.

triangle_rasterization.py:

# Use real division everywhere
from __future__ import division

import numpy as np


def rasterize_triangle(tri):
    """
    Given a 3x2 numpy array TRI describing the integer vertices of a general
    triangle, return an array containing all the points that lie within this
    triangle or on the triangle's edge, but not the triangle vertices
    themselves.
    This code is based on the description given in
    http://www.sunshine2k.de/coding/java/TriangleRasterization/TriangleRasterization.html
    """
    # Sort by increasing y coordinate
    tri = tri[tri[:, 1].argsort()]

    # Check for triangles with horizontal edge
    if tri[1, 1] == tri[2, 1]:
        # Bottom is horizontal
        points = rasterize_flat_triangle(tri)
    elif tri[0, 1] == tri[1, 1]:
        # Top is horizontal
        points = rasterize_flat_triangle(tri[(2, 0, 1), :])
    else:
        # General triangle.
        # We'll split this into two triangles with horizontal edges and process
        # them separately.
        # Find the additional vertex that splits the triangle.
        helper_point = np.array([tri[0, 0] + (tri[1, 1] - tri[0, 1]) /
                                         (tri[2, 1] - tri[0, 1]) *
                                         (tri[2, 0] - tri[0, 0]),
                                         tri[1, 1]]).round()
        # Top triangle
        points = rasterize_flat_triangle(tri[(0, 1), :],
            helper_point=helper_point)
        # Bottom triangle
        points = np.vstack([points, rasterize_flat_triangle(tri[(2, 1), :],
            helper_point=helper_point)])
    return points


def rasterize_flat_triangle(tri, helper_point=None):
    '''
    Given a 3x2 numpy array TRI describing the vertices of a triangle where the
    second and third vertex have the same y coordinate, return an array
    containing all the points that lie within this triangle or
    on the triangle's edge, but not the triangle vertices themselves.
    Or, given a 2x2 numpy array TRI containing two vertices and HELPER_POINT
    containing the third vertex, again return the same points as before, but
    additionally return the helper_point as well (used when treating a
    general triangle that's split into two triangles with horizontal edges)
    '''
    # Is the triangle we're treating part of a split triangle?
    if helper_point is not None:
        tri = np.vstack([tri, helper_point])

    # Is the bottom or the top edge horizontal?
    ydir = np.sign(tri[1, 1] - tri[0, 1])

    # Make sure that the horizontal edge is left-right oriented
    if tri[1, 0] > tri[2, 0]:
        tri[1, 0], tri[2, 0] = tri[2, 0], tri[1, 0]

    # Find the inverse slope (dx/dy) for the two non-horizontal edges
    invslope1 = ydir * (tri[1, 0] - tri[0, 0]) / (tri[1, 1] - tri[0, 1])
    invslope2 = ydir * (tri[2, 0] - tri[0, 0]) / (tri[2, 1] - tri[0, 1])

    # Initialize the first scan line, which is one y-step below or above the
    # first vertex
    curx1 = tri[0, 0] + invslope1
    curx2 = tri[0, 0] + invslope2
    points = []

    # Step vertically. Don't include the first row, because that row only
    # contains the first vertex and we don't want to return the vertices
    for y in np.arange(tri[0, 1] + ydir, tri[1, 1], ydir):
        for x in np.arange(curx1.round(), curx2.round() + 1):
            points.extend([(x, y)])
        curx1 += invslope1
        curx2 += invslope2

    # If we're dealing with the first half of a split triangle, add the
    # helper point (because that's not a "real" vertex of the triangle)
    if helper_point is not None and ydir == 1:
        points.extend([tuple(helper_point)])

    # If we're not dealing with a split triangle, or if we're dealing with the
    # first half of a split triangle, add the last line (but without the end
    # points, because they're the vertices of the triangle
    if helper_point is None or ydir == 1:
        for x in np.arange(tri[1, 0] + 1, tri[2, 0]):
            points.extend([(x, tri[1, 1])])

    return np.array(points, dtype='int')

And here's a small test script:

import triangle_rasterization as tr
import matplotlib.pyplot as plt

triangleA = np.array([[5,15],[5,1],[14,8]])
pointsA = tr.rasterize_triangle(triangleA)

triangleB = np.array([[14,8],[5,1],[18, 1]])
pointsB = tr.rasterize_triangle(triangleB)

triangleC = np.array([[5,15],[14,15],[14,8]])
pointsC = tr.rasterize_triangle(triangleC)

array = np.zeros([20,20])
array[pointsA[:,1], pointsA[:,0]] = 1
array[pointsB[:,1], pointsB[:,0]] = 2
array[pointsC[:,1], pointsC[:,0]] = 3

plt.imshow(array, interpolation='none')
plt.scatter(*triangleA.T, c='white')
plt.scatter(*triangleB.T, c='white')
plt.scatter(*triangleC.T, c='white')

Answer 1

1. In rasterize_triangle, the docstring says:

   return an array containing all the points […]

   but I think it would be clearer to write:

   return an array containing the coordinates of the points […]

2. In Numpy, it's usually more convenient for functions that return coordinates to return a tuple of arrays, rather than a multi-dimensional array. See for example numpy.diag_indices or numpy.triu_indices. That's because you can use a tuple of coordinate arrays to index an array:

   >>> a = np.arange(16).reshape(4,4)
   >>> np.diag_indices(4)
   (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
   >>> a[_]
   array([ 0,  5, 10, 15])
   

   Returning an n×2 array is less convenient: you have have to call numpy.unravel_index on the result before you can use it to index an array.

3. In rasterize_triangle the docstring says:

   This code is based on the description given in […]

   Docstrings are best written from the user's point of view (how do I use this function? what arguments do I pass? what does it return?), but this sentence is from the implementer's point of view and would be better made into a comment.

4. The function rasterize_flat_triangle has a long and complicated docstring, which suggests that its interface is not quite right. There are two ways of calling this function: pass a 2×2 array and a helper point (in which case the array and the helper point are stacked to form a triangle, and the scanline including the helper point is included in the result if the bottom is horizontal), or pass a 3×2 array and no helper point.

   It would simplify the interface if you let the caller worry about the stacking of the input and the inclusion or exclusion of the scanline containing the helper point. This only happens in one place (in rasterize_triangle where the triangle is split), so it would be an overall simplification.

5. There are no examples or test cases. These are the kind of self-contained function where doctests would work well.

6. The x-coordinate of the helper point is computed using the expression:

   tri[0, 0] + (tri[1, 1] - tri[0, 1]) / (tri[2, 1] - tri[0, 1]) * (tri[2, 0] - tri[0, 0])
   

   but I think numpy.interp would make the intention clearer (that this point is interpolated along the line between points 0 and 2):

   np.interp(tri[1, 1], tri[(0, 2), 1], tri[(0, 2), 0])
   

7. The base of the triangle is made horizontal (if necessary) like this:

   tri[(2, 0, 1), :]
   

   but I think numpy.roll would make this operation clearer:

   np.roll(tri, 1, 0)
   

8. For rasterize_flat_triangle to work, it must be the case that the second and third vertices have the same y coordinate, so I would assert this:

   assert tri[1, 1] == tri[2, 1]
   

9. In rasterize_flat_triangle the result is returned as np.array(points, dtype='int'). This makes it awkward to use this function with other data types such as 'int64'. Consider taking the datatype of the result as a keyword argument (defaulting to 'int').