import cv2
import numpy as np
cv_vers = int(cv2.__version__[0])
def _channel_check(img, type_):
"""
Convenience function for raising an exception if the input image is not
of the correct format.
"""
exc_info = {2: "grayscale", 3: "BGR"}
channels = len(np.asarray(img).shape)
if channels != type_:
raise ValueError(
f"Input image needs to be {exc_info[type_]} or have {type_} channels. Instead an image with {channels} channels was provided."
)
[docs]def find_contours(img, *args, **kwargs):
"""
Convenience function for contour detection.
This function accounts for the OpenCV version supplied and executes
accordingly.
Parameters
----------
img : ndarray
Grayscale image to perform contour detection on.
*args : tuple
Additional arguments `cv2.findContours` will require.
*kwargs : tuple
Additional keyword arguments `cv2.findContours` will require.
Returns
-------
contours : ndarray
An array of contours detected.
hierarchy : ndarray
Corresponding hierarchy information to the contours detected.
Raises
------
ValueError
If the supplied image is not grayscale or has greater than 2 channels.
"""
_channel_check(img, 2)
if cv_vers == 4:
contours, hierarchy = cv2.findContours(img, *args, **kwargs)
else:
_, contours, hierarchy = cv2.findContours(img, *args, **kwargs)
return contours, hierarchy
[docs]def fit_circles(contours):
"""
Fit minimum enclosing circles to contour(s).
Parameters
----------
contours : ndarray or list
Contour(s) to fit circles to.
Returns
-------
circles : ndarray or list
An array or list corresponding to dimensions to circles fitted.
"""
if isinstance(contours, list) or isinstance(contours, tuple):
ret = [cv2.minEnclosingCircle(c) for c in contours]
else:
ret = cv2.minEnclosingCircle(contours)
return ret
[docs]def fit_rectangles(contours, rotate=False):
"""
Fit bounding boxes to contour(s).
Parameters
----------
contours : ndarray or list
Contour(s) to fit circles to.
rotate : bool
Whether to fit rotated bounding boxes, default is False.
Returns
-------
rectangles : ndarray or list
An array or list corresponding to dimensions to bounding boxes fitted.
"""
if isinstance(contours, list) or isinstance(contours, tuple):
if rotate:
ret = [cv2.minAreaRect(c) for c in contours]
else:
ret = [cv2.boundingRect(c) for c in contours]
else:
if rotate:
ret = cv2.minAreaRect(contours)
else:
ret = cv2.boundingRect(contours)
return ret
[docs]def fit_ellipses(contours):
"""
Fit ellipses to contour(s).
Parameters
----------
contours : ndarray or list
Contour(s) to fit ellipses to.
Returns
-------
ellipses : ndarray or list
An array or list corresponding to dimensions to ellipses fitted.
"""
if isinstance(contours, list) or isinstance(contours, tuple):
ret = [cv2.fitEllipse(c) for c in contours]
else:
ret = cv2.fitEllipse(contours)
return ret
[docs]def draw_contours(img, contours, *args, **kwargs):
"""
Convenience function for drawing contour(s) on an image.
Parameters
----------
img : ndarray
Image to draw contours on.
contours : ndarray or list
Contour(s) to draw on the supplied image.
*args : tuple
Additional arguments `cv2.drawContours` will require.
*kwargs : tuple
Additional keyword arguments `cv2.drawContours` will require.
See Also
--------
draw_rectangles
draw_circles
draw_ellipses
"""
if isinstance(contours, list) or isinstance(contours, tuple):
for c in contours:
cv2.drawContours(img, [c], *args, **kwargs)
else:
cv2.drawContours(img, [contours], *args, **kwargs)
[docs]def draw_rectangles(img, dims, *args, **kwargs):
"""
Convenience function for drawing rectangle(s) on an image.
Parameters
----------
img : ndarray
Image to draw contours on.
dims : tuple or list
Rectangle(s) to draw on the supplied image. Note these should be
supplied in (x,y,w,h) format.
*args : tuple
Additional arguments `cv2.rectangle` will require.
*kwargs : tuple
Additional keyword arguments `cv2.rectangle` will require.
See Also
--------
draw_contours
draw_circles
draw_ellipses
"""
def _draw(rect):
x,y,w,h = rect
cv2.rectangle(img, (x, y), (x + w, y + h), *args, **kwargs)
if isinstance(dims, list):
for r in dims: _draw(r)
else:
_draw(dims)
[docs]def draw_circles(img, dims, *args, **kwargs):
"""
Convenience function for drawing circle(s) on an image.
Parameters
----------
img : ndarray
Image to draw contours on.
dims : tuple or list
Circle(s) to draw on the supplied image. Note these should be
supplied in ((x,y),r) format.
*args : tuple
Additional arguments `cv2.circle` will require.
*kwargs : tuple
Additional keyword arguments `cv2.circle` will require.
See Also
--------
draw_contours
draw_rectangles
draw_ellipses
"""
def _draw(circ):
(x,y),r = circ
x,y,r = map(int, (x,y,r))
cv2.circle(img, (x,y), r, *args, **kwargs)
if isinstance(dims, list):
for c in dims: _draw(c)
else:
_draw(c)
[docs]def draw_ellipses(img, dims, *args, **kwargs):
"""
Convenience function for drawing ellipse(s) on an image.
Parameters
----------
img : ndarray
Image to draw contours on.
dims : tuple or list
Ellipse(s) to draw on the supplied image. Note these should be
supplied in ((x,y),(w,h),a) format.
*args : tuple
Additional arguments `cv2.ellipse` will require.
*kwargs : tuple
Additional keyword arguments `cv2.ellipse` will require.
See Also
--------
draw_contours
draw_rectangles
draw_circles
"""
if isinstance(dims, list):
for c in dims:
cv2.ellipse(img, c, *args, **kwargs)
else:
cv2.ellipse(img, dims, *args, **kwargs)
[docs]def gray(frame):
"""
Convert a frame to grayscale.
Parameters
----------
frame : ndarray
Frame to grayscale.
Returns
-------
frame : ndarray
Grayscale frame.
Raises
------
ValueError
If the supplied image is not of BGR format.
See Also
--------
bgr
hsv
"""
_channel_check(frame, 3)
return cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
[docs]def bgr(frame):
"""
Convert a frame to BGR format.
Parameters
----------
frame : ndarray
Frame to convert.
Returns
-------
frame : ndarray
Frame of BGR format.
Raises
------
ValueError
If the supplied image is not grayscale.
See Also
--------
gray
hsv
"""
_channel_check(frame, 2)
return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
[docs]def hsv(frame):
"""
Convert a frame to HSV format.
Parameters
----------
frame : ndarray
Frame to convert.
Returns
-------
frame : ndarray
Frame of HSV format.
Raises
------
ValueError
If the supplied image is not of BGR format.
See Also
--------
gray
bgr
"""
_channel_check(frame, 3)
return cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
[docs]class Mask:
"""
Convenience class for performing segmentation to a mask.
``Mask`` enables the creation of a callable that will perform
``bitwise_and`` on a given frame with the mask supplied at initiation.
This can be useful for performing the same operation on a series of frames.
Parameters
----------
mask : ndarray
Mask to segment frame(s) to.
Returns
-------
mask : Mask
Class object with ``__call__`` method that performs ``bitwise_and``
with the mask provided at initiation.
Raises
------
ValueError
If the supplied image is not grayscale.
See Also
--------
shrink
rect_mask
circle_mask
contour_mask
"""
[docs] def __init__(self, mask):
_channel_check(mask, 2)
self.mask = mask
def __call__(self, frame) -> "Mask":
"""
Mask a frame using the pre-defined mask.
Parameters
----------
frame : ndarray.
Frame to mask.
Returns
-------
masked_frame : ndarray
Masked frame.
Raises
------
ValueError
If the supplied image is not grayscale.
"""
_channel_check(frame, 2)
return cv2.bitwise_and(frame, frame, mask=self.mask)
[docs]def cast_contours(contours, x, y):
"""
Convenience function to cast contour(s) to a given x,y.
Parameters
----------
contours : list or ndarray
Contour(s) to cast.
x : int
x position to cast to.
y : int
y position to cast to.
Returns
-------
contours : list of ndarray
Contours casted to new coordinates.
See Also
--------
find_contours
draw_contours
"""
if isinstance(contours, list) or isinstance(contours, tuple):
for i, c in enumerate(contours):
c[:, 0, 0] += x
c[:, 0, 1] += y
contours[i] = c
else:
contours[:, 0, 0] += x
contours[:, 0, 1] += y
return contours
[docs]def shrink(img, by=50):
"""
Mask an image to a new roi.
Parameters
----------
img : ndarray
Grayscale image to reduce the roi for.
by : int
Number of pixels to reduce roi by, default is 50.
Returns
-------
img : ndarray
Image with reduced roi.
Raises
------
ValueError
If the supplied image is not grayscale.
Notes
-----
This function will shrink the roi by a uniform amount on all four
sides of an image, and not to a specific region within the image.
See Also
--------
Mask
rect_mask
circle_mask
contouer_mask
"""
_channel_check(img, 2)
mask = np.zeros_like(img)
h, w = img.shape
mask[by:h-by, by:w-by] = 1
img = cv2.bitwise_and(img, img, mask=mask)
return img
[docs]def rect_mask(img, dims):
"""
Create a rectangular mask.
Parameters
----------
img : ndarray
Grayscale image to produce the rectangular mask for.
dims : tuple or list
Rectangle(s) to base mask on. Note these should be
supplied in (x,y,w,h) format.
Returns
-------
mask : ndarray
Rectangular mask.
Raises
------
ValueError
If the supplied image is not grayscale.
See Also
--------
Mask
shrink
circle_mask
contour_mask
"""
_channel_check(img, 2)
rect_mask_ = np.zeros_like(img)
draw_rectangles(rect_mask_, dims, 255, -1)
return rect_mask_
[docs]def circle_mask(img, dims):
"""
Create a circular mask.
Parameters
----------
img : ndarray
Grayscale image to produce circular mask for.
dims : tuple or list
Circle(s) to base mask on. Note these should be
supplied in ((x,y),r) format.
Returns
-------
mask : ndarray
Circular mask.
Raises
------
ValueError
If the supplied image is not grayscale.
See Also
--------
Mask
shrink
rect_mask
contour_mask
"""
_channel_check(img, 2)
circle_mask_ = np.zeros_like(img)
draw_circles(circle_mask_, dims, 255, -1)
return circle_mask_
[docs]def ellipse_mask(img, dims):
"""
Create an ellipse mask.
Parameters
----------
img : ndarray
Grayscale image to produce elliptical mask for.
dims : tuple or list
Ellipse(s) to base mask on. Note these should be
supplied in ((x,y),(w,h),a) format.
Returns
-------
mask : ndarray
Elliptical mask.
Raises
------
ValueError
If the supplied image is not grayscale.
See Also
--------
Mask
shrink
rect_mask
circle_mask
contour_mask
"""
_channel_check(img, 2)
ell_mask_ = np.zeros_like(img)
draw_ellipses(ell_mask_, dims, 255, -1)
return ell_mask_
[docs]def contour_mask(img, contours):
"""
Create a contour mask at the dimensions of contour(s).
Parameters
----------
img : ndarray
Grayscale image to produce circular mask for.
contours : list or ndarray
Contour(s) to create mask with.
Returns
-------
mask : ndarray
Contour mask at the dimensions of the contour(s) supplied.
Raises
------
ValueError
If the supplied image is not grayscale.
See Also
--------
Mask
shrink
rect_mask
circle_mask
"""
_channel_check(img, 2)
cnt_mask = np.zeros_like(img)
draw_contours(cnt_mask, contours, -1, 255, -1)
return cnt_mask
# Contour filters ----------------------------------------------------------------------------
def _contours_area(contours):
"""Convenience function for computing the area of contour(s). """
if isinstance(contours, list) or isinstance(contours, tuple):
areas = [cv2.contourArea(c) for c in contours]
else:
areas = cv2.contourArea(contours)
return areas
[docs]def smallest(contours):
"""
Contour filter that returns the smallest contour by area.
Parameters
----------
contours : list or ndarray
Contour(s) to filter.
Returns
-------
contour : ndarray
Smallest contour by area.
See Also
--------
largest
parents
Area
Eccentricity
Solidity
"""
areas = _contours_area(contours)
return contours[np.argmin(areas)]
[docs]def largest(contours):
"""
Contour filter that returns the largest contour by area.
Parameters
----------
contours : list or ndarray
Contour(s) to filter.
Returns
-------
contour : ndarray
Largest contour by area.
See Also
--------
smallest
parents
Area
Eccentricity
Solidity
"""
areas = _contours_area(contours)
return contours[np.argmax(areas)]
[docs]def parents(contours, hierarchy):
"""
Contour filter that returns only parent contours.
Parameters
----------
contours : list or ndarray
Contour(s) to filter.
hierarchy : ndarray
Corresponding hierarchy information for contours.
Returns
-------
contour : ndarray
All parent contours.
See Also
--------
smallest
largest
Area
Eccentricity
Solidity
"""
parentLvl = hierarchy[0, :, 3].tolist()
return [contours[i] for i, e in enumerate(parentLvl) if e == -1]
class _ByLimit:
"""
Filter values based on pre-defined limits.
"""
def __init__(self, min=None, max=None):
if min is None and max is None:
raise ValueError("Limit(s) required for filtering contours.")
self.min = min
self.max = max
def __call__(self, x):
"""
Filter x based on a lower or upper limit, or both. If x is between either
of the pre-defined limits then it will be returned.
"""
min, max = (self.min, self.max)
if min and max:
if x >= min and x <= max:
return x
else:
if min and not max:
if x >= min:
return x
elif max and not min:
if x <= max:
return x
[docs]class Area:
"""
Filter contours by area based on a lower or upper limit, or both.
Parameters
----------
min : int or float
Lower limit to filter contour areas.
max: int or float
Upper limit to filter contour areas.
Returns
-------
filter : Area
Class object with ``__call__`` method that filters contours based on
the pre-defined area limits provided at initiation.
See Also
--------
smallest
largest
parents
Eccentricity
Solidity
"""
[docs] def __init__(self, min=None, max=None):
self._filter = _ByLimit(min, max)
def __call__(self, contours) -> "Area":
"""
Filter contours by area based on pre-defined limits.
Parameters
----------
contours : list
List of contours to filter.
Returns
-------
contours : list
Filtered contours.
"""
areas = _contours_area(contours)
_filtered = []
for i, a in enumerate(areas):
if self._filter(a):
_filtered.append(contours[i])
return _filtered
[docs]class Eccentricity:
"""
Filter contours by eccentricity based on a lower or upper limit, or both.
Parameters
----------
min : int or float
Lower eccentricity limit to filter contours.
max: int or float
Upper eccentricity limit to filter contours.
Returns
-------
filter : Eccentricity
Class object with ``__call__`` method that filters contours based on
the pre-defined eccentricity limits provided at initiation.
See Also
--------
smallest
largest
parents
Area
Eccentricity
"""
[docs] def __init__(self, min=None, max=None):
self._filter = _ByLimit(min, max)
def _eccentricity(self, contours):
"""Convenience method for computing the eccentricity of contours. """
for i, c in enumerate(contours):
if c.shape[0] >= 5: # filter to this size for below calculations
center, axes, orientation = cv2.fitEllipse(c)
major, minor = (max(axes), min(axes))
eccentricity = np.sqrt(1 - (minor / major) ** 2)
yield (i, eccentricity)
def __call__(self, contours) -> "Eccentricity":
"""
Filter contours by eccentricity based on pre-defined limits.
Parameters
----------
contours : list
List of contours to filter.
Returns
-------
contours : list
Filtered contours.
"""
_filtered = []
for i, e in self._eccentricity(contours):
if self._filter(e):
_filtered.append(contours[i])
return _filtered
[docs]class Solidity:
"""
Filter contours by solidity based on a lower or upper limit, or both.
Parameters
----------
min : int or float
Lower solidity limit to filter contours.
max: int or float
Upper solidity limit to filter contours.
Returns
-------
filter : Solidity
Class object with ``__call__`` method that filters contours based on
the pre-defined solidity limits provided at initiation.
Notes
-----
Note that here the solidity of a contour is based on its area relative
to its corresponding convex hull.
See Also
--------
smallest
largest
parents
Area
Eccentricity
"""
[docs] def __init__(self, min=None, max=None):
self._filter = _ByLimit(min, max)
def _solidity(self, contours):
"""Convenience method for computing the solidity of contours. """
for i, c in enumerate(contours):
c_area = cv2.contourArea(c)
hull = cv2.convexHull(c)
hullarea = cv2.contourArea(hull)
try:
yield (i, c_area / hullarea)
except ZeroDivisionError:
pass
def __call__(self, contours) -> "Solidity":
"""
Filter contours by solidity based on pre-defined limits.
Parameters
----------
contours : list
List of contours to filter.
Returns
-------
contours : list
Filtered contours.
"""
_filtered = []
for i, s in self._solidity(contours):
if self._filter(s):
_filtered.append(contours[i])
return _filtered