Pixel weight calculator

pixel_weights.py

import numpy as np

# Author: Ethan Trewhitt <ethan@trewhitt.org>
#
# Given a space of dimension 'dims', determine the length of a line segment between p1 and p2 in terms of its length within
# every pixel in the space. Yields a NumPy array of shape 'dims' where each value is the length of the line segment through
# the corresponding pixel. Maximum output value is sqrt(2), which means the segment went from one corner of the pixel to the
# other. Minimum output value is 0, which means the line segment did not go through that pixel.
#
# dims = NumPy compatible dimension specification (duple)
# p1 = duple containing (x, y) of one end of line segment (can be floats)
# p2 = duple containing (x, y) of other end
#
# Notes:
#    p1 and p2 must lay within the bounds of dims.
#    This code may divide by zero, which is okay. Ignore with np.seterr(divide='ignore')
#
# Sample input values:
#    dims = (100, 50)
#    p1 = (4, 2.6)
#    p2 = (83.4, 23)
#
def pixel_weights(dims, p1, p2):
    TOLERANCE = 1e-12       # used to compensate for rounding
    
    p1_x = float(p1[0])
    p1_y = float(p1[1])
    p2_x = float(p2[0])
    p2_y = float(p2[1])
    
    m = np.divide((p2_y - p1_y), (p2_x - p1_x))

    # Find all X intersections and the corresponding Y values
    new_x = np.arange( (math.floor(min(p1_x, p2_x) + TOLERANCE) + 1), (math.ceil(max(p1_x, p2_x) - TOLERANCE)) )
    new_xy = np.vstack(( new_x, (m * (new_x - p1_x) + p1_y) ))

    # Find all Y intersections and the corresponding X values
    new_y = np.arange( (math.floor(min(p1_y, p2_y) + TOLERANCE) + 1), (math.ceil(max(p1_y, p2_y) - TOLERANCE)) )
    new_yx = np.vstack(( ( np.divide((new_y - p1_y), m) + p1_x), new_y ))

    # Combine the two lists along with the start and end points
    points = np.hstack(( np.vstack((p1_x, p1_y)), np.vstack((p2_x, p2_y)), new_xy, new_yx ))

    # Sort by Y and X (in case it's a horizontal line)
    points = points[:,points[1,:].argsort(kind='mergesort')[::-1]]
    points = points[:,points[0,:].argsort(kind='mergesort')[::-1]]
    
    # Round so that we can combine any really close points
    points = np.round(points / TOLERANCE) * TOLERANCE

    # Eliminate duplicate points (this approach works because the points are sorted)
    points = points[:, ( np.hstack(( points[0,0:-1], np.NAN )) != np.hstack(( points[0,1:], np.NAN )) ) | ( np.hstack(( points[1,0:-1], np.NAN )) != np.hstack(( points[1,1:], np.NAN )) ) ]

    # Get the segments between each pair of points
    starts = points[:, 0:-1]        # all points but the last
    ends = points[:, 1:]            # all points but the first

    # Attribute each segment to the pixel occupied by its midpoint
    idx = np.floor((starts + ends) / 2).astype(int)

    weights = np.zeros(dims)
    weights[idx[0,:], idx[1,:]] = np.sqrt( np.power(starts[0,:] - ends[0,:], 2) + np.power(starts[1,:] - ends[1,:], 2) )        # Pythagorean

    return weights