opencv_contours_to_polygons.py

"""OpenCV contours to Shapely polygons converter

This script allows the user to convert the contours hirarchy generated by cv2.find_contours()
into a list of Shapely polygons.

This script requires that `cv2` `numpy` and `shapely` be installed within the Python
environment you are running this script in.

This file can also be imported as a module and contains the following
functions:

    * generate_polygons - generate the list of Shapely polygons from cv2 contours hierarchy
    * main - the main function of the script
"""

import cv2
import numpy
from shapely.geometry import Polygon

def _DFS(polygons, contours, hierarchy, sibling_id, is_outer, siblings):
  while sibling_id != -1:
    contour = contours[sibling_id].squeeze(axis=1)
    if len(contour) >= 3:
      first_child_id = hierarchy[sibling_id][2]
      children = [] if is_outer else None
      _DFS(polygons, contours, hierarchy, first_child_id, not is_outer, children)

      if is_outer:
        polygon = Polygon(contour, holes=children)
        polygons.append(polygon)
      else:
        siblings.append(contour)

    sibling_id = hierarchy[sibling_id][0]

def generate_polygons(contours, hierarchy):
  """Generates a list of Shapely polygons from the contours hirarchy returned by cv2.find_contours().
     The list of polygons is generated by performing a depth-first search on the contours hierarchy tree.

  Parameters
  ----------
  contours : list
    The contours returned by cv2.find_contours()
  hierarchy : list
    The hierarchy returned by cv2.find_contours()

  Returns
  -------
  list
    The list of generated Shapely polygons
  """
  
  hierarchy = hierarchy[0]
  polygons = []
  _DFS(polygons, contours, hierarchy, 0, True, [])
  return polygons

def main(args):
  img = cv2.imread(str(args.input_img), cv2.IMREAD_GRAYSCALE)
  contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_TC89_L1)
  polygons = generate_polygons(contours, hierarchy)
  for p in polygons: 
    print(p)
  

if __name__ == "__main__":

  from argparse import ArgumentParser
  from pathlib import Path
  
  parser = ArgumentParser(description=__doc__)
  parser.add_argument('input_img', type=Path, help="Input image")
  args = parser.parse_args()
  main(args)