OpenCV Contour Tree Iteration Class

CvContourNode.py

# import for instances returned from class methods
from __future__ import annotations

import sys
from typing import Optional, List, Any

import cv2

from cars.colab.model.Global import CvRect
from cars.colab.util import Units
from cars.colab.util.cv import Cv, Colors


class CvContourNode:
    _hierarchy = None
    _contours = None

    @classmethod
    def parse(cls, hierarchy, contours) -> CvContourNode:

        sys.setrecursionlimit(10000)

        # get the actual inner list of hierarchy descriptions
        hierarchy = hierarchy[0]

        cls._hierarchy = hierarchy
        cls._contours = contours

        root = CvContourNode()
        cls._parse_hierarchy_item(0, root)

        return root

    @classmethod
    def _parse_hierarchy_item(cls, index, parent) \
            -> CvContourNode:
        h_item = cls._hierarchy[index]
        h_next = h_item[0]
        h_first_child = h_item[2]

        node = CvContourNode(parent)
        parent.children.append(node)
        node.contour = cls._contours[index]
        node.level = parent.level + 1

        # process siblings
        has_sibling = h_next >= 0
        if has_sibling:
            cls._parse_hierarchy_item(h_next, node.parent)

        # process children
        has_child = h_first_child >= 0
        if has_child:
            cls._parse_hierarchy_item(h_first_child, node)

        return node

    def __init__(self, parent: Optional[CvContourNode] = None) -> None:
        self.parent: Optional[CvContourNode] = parent
        self.children: List[CvContourNode] = []
        self.level: int = 0
        self.contour: Any = None

    def flatten(self, level=-1, level_modulo=0, result=[]):
        """

        :param level:
        :param result:
        :param level_modulo: ALL = 0, ODD = 1, EVEN = 2
        :return: flat array of contours
        """
        for node in self.children:
            result.append(node.contour)
            node.flatten(level, level_modulo, result)

        return result

    def draw(self, image, level_offset=0, level_modulo=1, level_color=None,
             thickness: int = 1):

        if level_color is None:
            level_color = [Colors.MAGENTA, Colors.RED, Colors.BLUE,
                           Colors.GREEN]
            for i in range(4, 100):
                level_color.append(Cv.random_color())

        # draw the current contour
        if self.contour is not None:
            is_valid = (self.level + level_offset) % level_modulo == 0
            if is_valid:
                rect: CvRect = cv2.boundingRect(self.contour)
                p1, p2, = Units.rect_to_points(rect)
                cv2.rectangle(image, p1, p2, level_color[self.level], thickness)

        # iterate through children
        for node in self.children:
            node.draw(image, level_offset, level_modulo, level_color, thickness)