Prompt a user to interactively pick N points using OpenCV and Python

cv_pick_points.py

"""
Prompts the user to pick num_points from the figure displayed in window_name

Promts the user to pick points from a displayed image. Users can cancel the
point selection by pressing Escape.

Args:
    window_name:    The name of the window the user should pick from
    num_points:     Number of points the user should pick (defaults to 4)
    mouse_event:    The event to use for picking points (defaults to left
                    mouse button single click)

Returns:
    None if the picking is cancelled with Escape, otherwise a float32 Numpy array
    of shape (num_points, 2) containing the picked points in pixel space
"""
def pick_points(window_name, num_points = 4, mouse_event=cv2.EVENT_LBUTTONDOWN):

    KEY_ESCAPE = 27
    KEY_WAIT_MS = 20

    points = np.zeros((0, 2), dtype = "float32")
    cancelled = False

    # Mouse callback function
    def on_click(event, x, y, flags, param):
        nonlocal points

        if event == mouse_event:
            points = np.append(points, [[x, y]], axis=0)

    cv2.setMouseCallback(window_name, on_click)

    # Wait for 4 clicks, or escape pressed
    while(1):
        if len(points) >= num_points:
            break
        
        k = cv2.waitKey(KEY_WAIT_MS)
        if k == KEY_ESCAPE:
            # Escape pressed - return nothing
            cancelled = True
            break
    
    # Return the points clicked
    # TODO ajs 27/Dec/17 Remove mouse callback somehow?
    # https://stackoverflow.com/questions/47988877/python-how-to-remove-mousecallback-in-opencv
    #cv2.setMouseCallback(window_name, None)
    if cancelled:
        return None
    return points