Source code for draugr.opencv_utilities.drawing.opencv_draw

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

__author__ = "Christian Heider Nielsen"
__doc__ = r"""

           Created on 19/03/2020
           """

from typing import Sequence, Tuple, Union

import cv2
import numpy
from PIL import Image

from draugr.opencv_utilities.namespaces.enums import (
    LineTypeEnum,
    ContourRetrievalModeEnum,
)
from warg import RGB, compute_color_for_labels

__all__ = ["find_contours", "draw_masks"]


[docs]def find_contours(*args, **kwargs) -> Tuple: """ Wraps cv2.findContours to maintain compatibility between versions 3 and 4 Returns: contours, hierarchy""" if cv2.__version__.startswith("4"): contours, hierarchy = cv2.findContours(*args, **kwargs) elif cv2.__version__.startswith("3"): _, contours, hierarchy = cv2.findContours(*args, **kwargs) else: raise AssertionError("cv2 must be either version 3 or 4 to call this method") return contours, hierarchy
[docs]def draw_masks( image: Union[Image.Image, numpy.ndarray], masks: Union[Image.Image, numpy.ndarray], *, labels: Sequence = None, border: bool = True, border_width: int = 1, # If it is negative, the contour interiors are drawn. border_color: Tuple = RGB(255, 255, 255), alpha: float = 0.5, color: Tuple = None, line_type: LineTypeEnum = LineTypeEnum.anti_aliased ) -> numpy.ndarray: """ Args: image: numpy array image, shape should be (height, width, channel) masks: (N, 1, Height, Width) labels: mask label border: draw border on mask border_width: border width border_color: border color alpha: mask alpha color: mask color Returns: numpy.ndarray""" line_type = LineTypeEnum(line_type) if isinstance(image, Image.Image): image = numpy.array(image) assert isinstance(image, numpy.ndarray) if isinstance(masks, Image.Image): masks = numpy.array(masks) assert isinstance(masks, numpy.ndarray) # TODO: ASSERT 3/4 CHANNELS! if labels is None: labels = list(range(masks.shape[0])) for i, mask in enumerate(masks): mask = mask.squeeze()[..., None].astype(numpy.bool) label = labels[i] mask_color = compute_color_for_labels(label) if color is None else tuple(color) image = numpy.where( mask, mask * numpy.array(mask_color) * alpha + image * (1 - alpha), image ) if border: contours, hierarchy = find_contours( mask.astype(numpy.uint8), ContourRetrievalModeEnum.tree.value, cv2.CHAIN_APPROX_SIMPLE, ) image = cv2.drawContours( image, contours, -1, border_color, thickness=border_width, lineType=line_type.value, ) return image.astype(numpy.uint8)
if __name__ == "__main__": pass