#!/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