Pixel-wise breadth-first search that not only traverses pixels but also returns the corresponding tree. Includes convenience functions.

bfs.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-


"""
Most existing breadth-first search implementations provide a solution to
*traverse* a tree, but they don't *return* that specific tree. This module
provides a tree datastructure and BFS implementation to obtain the BFS tree.
"""


import numpy as np


# ##############################################################################
# # HELPER DATASTRUCTURE
# ##############################################################################
class PixelTree:
    """
    Helper datastructure for pixel breadth-first search on images.
    Acts as a node that can have pixel position, depth, one parent and multiple
    children.
    Provides convenience methods to modify and traverse.
    """
    def __init__(self, yx_pos, depth=1, parent=None):
        """
        :param yx_pos: Pixel position in format ``(y, x)``.
        :param int depth: Integer >= 1.
        :param parent: If this node is not root. Otherwise, it will set as its
          own parent.
        """
        assert depth >= 1, "Depth has to be >= 1!"
        self.y, self.x = yx_pos
        self.depth = depth
        self.parent = self if parent is None else parent
        self.children = []

    def add_child(self, yx_pos):
        """
        Adds a new PixelTree instance with the given y,x values to this
        instance's children and increments its depth by one.
        :returns: The newly created child.
        """
        child = PixelTree(yx_pos, self.depth+1, self)
        self.children.append(child)
        return child

    def flatten(self, as_yxd_list=False):
        """
        Iterative depth-first traverser that returns a list with this node
        and its (recursive) children, in depth-first order.
        :param bool as_yxd_list: If true, the returned elements will be
          triples of ``(y_pos, x_pos, depth)`` instead of class instances.
        """
        result = []
        marked = set()
        queue = [self]+[child for child in self.children]
        while queue:
            current = queue.pop(0)
            if current not in marked:
                marked.add(current)
                queue += current.children
                if as_yxd_list:
                    current = (current.y, current.x, current.depth)
                result.append(current)
        return result

    def as_array(self, depth_offset=0, dtype=np.float32):
        """
        Converts this node and all its recursive children into a numpy array
        with minimal shape to contain all nodes.
        Each node is then assigned to a corresponding ``(y, x)`` array index,
        and the node depth is written at that index.

        :param depth_offset: Optionally adds this extra depth to the in case
          more depth contrast against the background (zero depth) is needed.
        :param dtype: Numeric type of the returned array.
        :returns: A tuple ``(arr, min_y, min_x, max_y, max_x)``. Since the
          array is minimal-shape (i.e. it doesn't contain empty margins),
          the min and max ranges provide context for the array location.
        :raises: ``IndexError`` if any of the ``(y, x)`` positions is not a
          non-negative integer that is a valid array index.
        """
        yxd_list = self.flatten(True)
        y_list, x_list, _ = zip(*yxd_list)
        min_y, max_y = min(y_list), max(y_list)
        min_x, max_x = min(x_list), max(x_list)
        result = np.zeros((1+max_y-min_y, 1+max_x-min_x), dtype=dtype)
        for y, x, d in yxd_list:
            result[y-min_y, x-min_x] = d + depth_offset
        return result, min_y, min_x, max_y, max_x

    def __len__(self):
        """
        """
        return len(self.flatten())


# ##############################################################################
# # BFS
# ##############################################################################
class PixelBFS:
    """
    Parses a binary mask image via breadth-first search for neighbouring pixels
    and returns the parsed forest.

    :cvar NEIGH_DELTAS: Dictionary mapping neighbouring relationships to the
      corresponding ``(y, x)`` offsets.
    :cvar NEIGH_DELTAS: Dictionary mapping neighbouring relationships to the
      corresponding binary codes.
    :cvar NEIGH_DTYPE: Data type capable of containing the sum of all codes.
    """

    NEIGH_DELTAS = {"b": (1, 0), "r": (0, 1), "l": (0, -1), "t": (-1, 0),
                    "br": (1, 1), "bl": (1, -1), "tr": (-1, 1), "tl": (-1, -1)}
    NEIGH_CODES = {"b": 1, "r": 2, "l": 4, "t": 8,
                   "br": 16, "bl": 32, "tr": 64, "tl": 128}
    NEIGH_4 = ("b", "r", "l", "t")
    NEIGH_8 = ("b", "r", "l", "t", "br", "bl", "tr", "tl")
    NEIGH_DTYPE = np.uint8

    @classmethod
    def get_neighbour_map(cls, y_indexes, x_indexes, with_8con=False):
        """
        :param y_indexes: Integer array of shape ``(N,)`` containing the
          y-indexes of the pixels.
        :param x_indexes: Integer array of shape ``(N,)`` containing the
          x-indexes of the pixels.
        :param bool with_8con: If true, diagonal neighbours will also be
          regarded.
        :returns: Array of shape ``(1+ymax-ymin, 1+xmax-xmin)``, where
          ``result[0, 0]`` encodes the neighbouring relationships for the
          pixel at position ``(ymin, xmin)`` in the form of an encoding.
          This encoding is the sum of the active ``NEIGH_CODES``.
        """
        y_min, y_max = min(y_indexes), max(y_indexes)
        x_min, x_max = min(x_indexes), max(x_indexes)
        yyy1 = 1 + (y_indexes - y_min)  # starts on 1
        xxx1 = 1 + (x_indexes - x_min)  # starts on 1
        # vectorized hist to store neigh rels (+2 margins, +1 zero-idx)
        result = np.zeros((3 + y_max - y_min, 3 + x_max - x_min),
                          dtype=cls.NEIGH_DTYPE)
        # populate histogram following binary code from NEIGH_MASK
        # Check with e.g. check with x&0b100!=0
        neighs = cls.NEIGH_8 if with_8con else cls.NEIGH_4
        for n in neighs:
            delta_y, delta_x = cls.NEIGH_DELTAS[n]
            code = cls.NEIGH_CODES[n]
            result[(yyy1 - delta_y, xxx1 - delta_x)] += code
        # remove margins and return
        result = result[1:-1, 1:-1]
        return result

    @classmethod
    def __call__(cls, yyy_xxx, with_8con=False):
        """
        :param yyy_xxx: Pair of integer arrays of same shape ``(yyy, xxx)``,
          containing the pixel locations to be traversed.
        :param bool with_8con: If true, diagonal neighbours will also be
          regarded.
        :returns: A list of disjoint ``PixelTree`` objects, each one composed
          of neighbouring pixels arranged in breadth-first order.
        """
        # convenience aliases
        neighs = cls.NEIGH_8 if with_8con else cls.NEIGH_4
        yyy, xxx = yyy_xxx
        y_min, x_min = min(yyy), min(xxx)
        # Prepare globals for the BFS:
        bfs_forest = []
        marked = set()
        # compute pixel neighbour map and start BFS
        neigh_map = cls.get_neighbour_map(yyy, xxx)

        for i, yx in enumerate(zip(yyy, xxx)):
            if yx not in marked:
                tree = PixelTree(yx)  # start a new tree
                # queue won't be empty until all adj. pixels have been explored
                tree_queue = [tree]
                # hashable copy of tree queue needed for faster lookup
                queue_hash = {yx}
                #
                while tree_queue:
                    node = tree_queue.pop(0)
                    current_yx = (node.y, node.x)
                    if current_yx not in marked:
                        # if current is newly visited, mark and fetch neighs
                        marked.add(current_yx)
                        queue_hash.remove(current_yx)  # will never visit again
                        yy, xx = current_yx
                        neigh_code = neigh_map[yy - y_min, xx - x_min]

                        for n in neighs:
                            delta_y, delta_x = cls.NEIGH_DELTAS[n]
                            neigh_idx = (yy + delta_y, xx + delta_x)
                            # check if neigh_code contains this neighbour
                            is_neigh = neigh_code & cls.NEIGH_CODES[n]
                            if (is_neigh and (neigh_idx not in tree_queue) and
                               (neigh_idx not in marked)):
                                # if new+valid, create and queue new node
                                child = node.add_child(neigh_idx)
                                tree_queue.append(child)
                                queue_hash.add(neigh_idx)

                # add tree to forest
                bfs_forest.append(tree)
        #
        return bfs_forest


# ##############################################################################
# # CONVENIENCE FUNCTIONS
# ##############################################################################
def make_cyclic_contour(pixel_tree):
    """
    Assuming we extracted a BFS tree from a ring-shaped 1D object, the tree
    will have 2 main deep branch. This function gathers the 2 deepest paths,
    flips one of them and concatenates them, resulting in a cyclic list of
    locations.
    """
    # find out 2 deepest leaves
    leaves_by_depth = sorted([node for node in pixel_tree.flatten()
                              if not node.children], key=lambda x: x.depth)
    assert len(leaves_by_depth) >= 2, \
        "We assume contours have 2 deep leaves! This shouldn't happen"
    leaf_a, leaf_b = leaves_by_depth[-2:]
    # traverse up leaf until reaching root and collect path
    path_a, path_b = [(leaf_a.y, leaf_a.x)], [(leaf_b.y, leaf_b.x)]
    while leaf_a.parent != leaf_a:
        leaf_a = leaf_a.parent
        path_a.append((leaf_a.y, leaf_a.x))
    while leaf_b.parent != leaf_b:
        leaf_b = leaf_b.parent
        path_b.append((leaf_b.y, leaf_b.x))
    # finally append both paths in circular form
    result = path_a + list(reversed(path_b))
    return result