# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals
import itertools as it
import numpy as np
import cv2
[docs]def draw_text_on_image(img, text, org, return_info=False, **kwargs):
r"""
Draws multiline text on an image using opencv
Args:
img (ndarray | None | dict):
Generally a numpy image to draw on (inplace).
Otherwise a canvas will be constructed such that the text will fit.
The user may specify a dictionary with keys width and height
to have more control over the constructed canvas.
text (str): text to draw
org (Tuple[int, int]):
The x, y location of the text string "anchor" in the image as
specified by halign and valign. For instance, If valign='bottom',
halign='left', this is the bottom left corner.
return_info (bool, default=False):
if True, also returns information about the positions the text
was drawn on.
**kwargs:
color (tuple): default blue
thickness (int): defaults to 2
fontFace (int): defaults to cv2.FONT_HERSHEY_SIMPLEX
fontScale (float): defaults to 1.0
valign (str, default='bottom'):
either top, center, or bottom.
NOTE: this default may change to "top" in the future.
halign (str, default='left'):
either left, center, or right
border (dict | int):
If specified as an integer, draws a black border with that
given thickness. If specified as a dictionary, draws a border
with color specified parameters.
"color": border color, defaults to "black".
"thickness": border thickness, defaults to 1.
Returns:
ndarray: the image that was drawn on
Note:
The image is modified inplace. If the image is non-contiguous then this
returns a UMat instead of a ndarray, so be carefull with that.
References:
https://stackoverflow.com/questions/27647424/
https://stackoverflow.com/questions/51285616/opencvs-gettextsize-and-puttext-return-wrong-size-and-chop-letters-with-low
Example:
>>> import kwimage
>>> img = kwimage.grab_test_image(space='rgb')
>>> img2 = kwimage.draw_text_on_image(img.copy(), 'FOOBAR', org=(0, 0), valign='top')
>>> assert img2.shape == img.shape
>>> assert np.any(img2 != img)
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> kwplot.imshow(img2)
>>> kwplot.show_if_requested()
Example:
>>> import kwimage
>>> # Test valign
>>> img = kwimage.grab_test_image(space='rgb', dsize=(500, 500))
>>> img2 = kwimage.draw_text_on_image(img, 'FOOBAR\nbazbiz\nspam', org=(0, 0), valign='top', border=2)
>>> img2 = kwimage.draw_text_on_image(img, 'FOOBAR\nbazbiz\nspam', org=(150, 0), valign='center', border=2)
>>> img2 = kwimage.draw_text_on_image(img, 'FOOBAR\nbazbiz\nspam', org=(300, 0), valign='bottom', border=2)
>>> # Test halign
>>> img2 = kwimage.draw_text_on_image(img, 'FOOBAR\nbazbiz\nspam', org=(250, 100), halign='right', border=2)
>>> img2 = kwimage.draw_text_on_image(img, 'FOOBAR\nbazbiz\nspam', org=(250, 250), halign='center', border=2)
>>> img2 = kwimage.draw_text_on_image(img, 'FOOBAR\nbazbiz\nspam', org=(250, 400), halign='left', border=2)
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> kwplot.imshow(img2)
>>> kwplot.show_if_requested()
Example:
>>> # Ensure the function works with float01 or uint255 images
>>> import kwimage
>>> img = kwimage.grab_test_image(space='rgb')
>>> img = kwimage.ensure_float01(img)
>>> img2 = kwimage.draw_text_on_image(img, 'FOOBAR\nbazbiz\nspam', org=(0, 0), valign='top', border=2)
Example:
>>> # Test dictionary border
>>> import kwimage
>>> img = kwimage.draw_text_on_image(None, 'hello\neveryone', org=(100, 100), valign='top', halign='center', border={'color': 'green', 'thickness': 9})
>>> #img = kwimage.draw_text_on_image(None, 'hello\neveryone', org=(0, 0), valign='top')
>>> #img = kwimage.draw_text_on_image(None, 'hello', org=(0, 60), valign='top', halign='center', border=0)
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> kwplot.imshow(img)
>>> kwplot.show_if_requested()
Example:
>>> # Test dictionary image
>>> import kwimage
>>> img = kwimage.draw_text_on_image({'width': 300}, 'good\nPropogate', org=(150, 0), valign='top', halign='center', border={'color': 'green', 'thickness': 0})
>>> print('img.shape = {!r}'.format(img.shape))
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> kwplot.imshow(img)
>>> kwplot.show_if_requested()
"""
import kwimage
if 'color' not in kwargs:
kwargs['color'] = 'red'
# Get the color that is compatible with the input image encoding
if img is None or isinstance(img, dict):
kwargs['color'] = kwimage.Color(kwargs['color']).as255()
else:
kwargs['color'] = kwimage.Color(kwargs['color'])._forimage(img)
if 'thickness' not in kwargs:
kwargs['thickness'] = 2
if 'fontFace' not in kwargs:
kwargs['fontFace'] = cv2.FONT_HERSHEY_SIMPLEX
if 'fontScale' not in kwargs:
kwargs['fontScale'] = 1.0
if 'lineType' not in kwargs:
kwargs['lineType'] = cv2.LINE_AA
if 'bottomLeftOrigin' in kwargs:
raise ValueError('Do not use bottomLeftOrigin, use valign instead')
border = kwargs.pop('border', None)
if border is not None:
if isinstance(border, int):
border = {'color': 'black', 'thickness': border}
subkw = kwargs.copy()
subkw['color'] = border.get('color', 'black')
subkw.pop('return_info', None)
border_thickness = border.get('thickness', 1)
else:
border_thickness = 0
valign = kwargs.pop('valign', None)
halign = kwargs.pop('halign', None)
getsize_kw = {
k: kwargs[k]
for k in ['fontFace', 'fontScale', 'thickness']
if k in kwargs
}
x0, y0 = list(map(int, org))
thickness = kwargs.get('thickness', 2)
vertical_spacing = 4 # space between vertical lines
ypad = thickness + vertical_spacing
lines = text.split('\n')
line_sizes = []
final_baseline = 0
for line in lines:
# TODO: better handling of baseline
# https://en.wikipedia.org/wiki/Baseline_(typography)
(line_width, line_height), baseline = cv2.getTextSize(line, **getsize_kw)
line_sizes.append((line_width, line_height))
final_baseline = baseline
line_sizes = np.array(line_sizes)
line_org = []
y = y0
for w, h in line_sizes:
next_y = y + (h + ypad)
line_org.append((x0, y))
y = next_y
line_org = np.array(line_org)
# the absolute top and bottom position of text
abs_top_y = line_org[0, 1]
abs_bot_y = (line_org[-1, 1] + line_sizes[-1, 1]) + thickness
first_h = line_sizes[0, 1]
total_h = (abs_bot_y - abs_top_y)
total_w = line_sizes.T[0].max()
if valign is not None:
if valign == 'bottom':
# This is the default for the one-line case
# in the multiline case we need to subtract the total
# height of all lines but the first to ensure the last
# line is on the bottom.
line_org[:, 1] -= (total_h - first_h)
elif valign == 'center':
# Change from bottom to center
line_org[:, 1] += first_h - total_h // 2
elif valign == 'top':
# Because bottom is the default we just need to add height of the
# first line.
line_org[:, 1] += first_h
else:
raise KeyError(valign)
if halign is not None:
if halign == 'left':
# This is the default case, no modification needed
pass
elif halign == 'center':
# When the x-orgin should be the center, subtract half of
# the line width to get the leftmost point.
line_org[:, 0] = x0 - (line_sizes[:, 0] / 2)
elif halign == 'right':
# The x-orgin should be the rightmost point, subtract
# the width of each line to find the leftmost point.
line_org[:, 0] = x0 - line_sizes[:, 0]
else:
raise KeyError(halign)
if img is None:
img = {'width': None, 'height': None}
if isinstance(img, dict):
# if image is unspecified allocate just enough space for text
# allow users to specify partial parameters
alloc_w = img.get('width', None)
alloc_h = img.get('height', None)
if alloc_w is None:
abs_left_x = line_org[:, 0].min()
alloc_w = total_w + border_thickness + abs_left_x
if alloc_h is None:
alloc_h = total_h + border_thickness + abs_top_y + final_baseline
img = np.zeros((alloc_h, alloc_w, 3), dtype=np.uint8)
if border_thickness > 0:
# recursive call
basis = list(range(-border_thickness, border_thickness + 1))
for i, j in it.product(basis, basis):
if i == 0 and j == 0:
continue
org = np.array(org)
img = draw_text_on_image(img, text, org=org + [i, j], **subkw)
for i, line in enumerate(lines):
xy = tuple(line_org[i])
img = cv2.putText(img, line, xy, **kwargs)
if return_info:
info = {
'line_org': line_org,
'line_sizes': line_sizes,
}
return img, info
else:
return img
[docs]def draw_clf_on_image(im, classes, tcx=None, probs=None, pcx=None, border=1):
"""
Draws classification label on an image.
Works best with image chips sized between 200x200 and 500x500
Args:
im (ndarray): the image
classes (Sequence | CategoryTree): list of class names
tcx (int, default=None): true class index if known
probs (ndarray): predicted class probs for each class
pcx (int, default=None): predicted class index.
(if None but probs is specified uses argmax of probs)
Example:
>>> # xdoctest: +REQUIRES(module:torch)
>>> import torch
>>> import kwarray
>>> import kwimage
>>> rng = kwarray.ensure_rng(0)
>>> im = (rng.rand(300, 300) * 255).astype(np.uint8)
>>> classes = ['cls_a', 'cls_b', 'cls_c']
>>> tcx = 1
>>> probs = rng.rand(len(classes))
>>> probs[tcx] = 0
>>> probs = torch.FloatTensor(probs).softmax(dim=0).numpy()
>>> im1_ = kwimage.draw_clf_on_image(im, classes, tcx, probs)
>>> probs[tcx] = .9
>>> probs = torch.FloatTensor(probs).softmax(dim=0).numpy()
>>> im2_ = kwimage.draw_clf_on_image(im, classes, tcx, probs)
>>> # xdoctest: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> kwplot.imshow(im1_, colorspace='rgb', pnum=(1, 2, 1), fnum=1, doclf=True)
>>> kwplot.imshow(im2_, colorspace='rgb', pnum=(1, 2, 2), fnum=1)
>>> kwplot.show_if_requested()
"""
import kwimage
im_ = kwimage.atleast_3channels(im)
w, h = im.shape[0:2][::-1]
if pcx is None and probs is not None:
import kwarray
probs = kwarray.ArrayAPI.numpy(probs)
pcx = probs.argmax()
if probs is not None:
pred_score = None if pcx is None else probs[pcx]
true_score = None if tcx is None else probs[tcx]
org1 = np.array((2, h - 5))
org2 = np.array((2, 5))
true_label = None
if tcx is not None:
true_name = classes[tcx]
if pcx == tcx:
true_label = 't:{tcx}:{true_name}'.format(**locals())
elif probs is None:
true_label = 't:{tcx}:\n{true_name}'.format(**locals())
else:
true_label = 't:{tcx}@{true_score:.2f}:\n{true_name}'.format(**locals())
pred_label = None
if pcx is not None:
pred_name = classes[pcx]
if probs is None:
pred_label = 'p:{pcx}:\n{pred_name}'.format(**locals())
else:
pred_label = 'p:{pcx}@{pred_score:.2f}:\n{pred_name}'.format(**locals())
fontkw = {
'fontScale': 1.0,
'thickness': 2
}
color = 'dodgerblue' if pcx == tcx else 'orangered'
# im_ = draw_text_on_image(im_, pred_label, org=org1 - 2,
# color='white', valign='bottom', **fontkw)
# im_ = draw_text_on_image(im_, true_label, org=org2 - 2,
# color='white', valign='top', **fontkw)
if pred_label is not None:
im_ = draw_text_on_image(im_, pred_label, org=org1, color=color,
border=border, valign='bottom', **fontkw)
if true_label is not None:
im_ = draw_text_on_image(im_, true_label, org=org2, color='lawngreen',
valign='top', border=border, **fontkw)
return im_
[docs]def draw_boxes_on_image(img, boxes, color='blue', thickness=1,
box_format=None, colorspace='rgb'):
"""
Draws boxes on an image.
Args:
img (ndarray): image to copy and draw on
boxes (nh.util.Boxes): boxes to draw
colorspace (str): string code of the input image colorspace
Example:
>>> import kwimage
>>> import numpy as np
>>> img = np.zeros((10, 10, 3), dtype=np.uint8)
>>> color = 'dodgerblue'
>>> thickness = 1
>>> boxes = kwimage.Boxes([[1, 1, 8, 8]], 'ltrb')
>>> img2 = draw_boxes_on_image(img, boxes, color, thickness)
>>> assert tuple(img2[1, 1]) == (30, 144, 255)
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl() # xdoc: +SKIP
>>> kwplot.figure(doclf=True, fnum=1)
>>> kwplot.imshow(img2)
"""
import kwimage
import cv2
if not isinstance(boxes, kwimage.Boxes):
if box_format is None:
raise ValueError('specify box_format')
boxes = kwimage.Boxes(boxes, box_format)
color = kwimage.Color(color)._forimage(img, colorspace)
ltrb = boxes.to_ltrb().data
img2 = img.copy()
for x1, y1, x2, y2 in ltrb:
# pt1 = (int(round(x1)), int(round(y1)))
# pt2 = (int(round(x2)), int(round(y2)))
pt1 = (int(x1), int(y1))
pt2 = (int(x2), int(y2))
# Note cv2.rectangle does work inplace
img2 = cv2.rectangle(img2, pt1, pt2, color, thickness=thickness)
return img2
[docs]def draw_line_segments_on_image(
img, pts1, pts2, color='blue', colorspace='rgb', thickness=1,
**kwargs):
"""
Draw line segments between pts1 and pts2 on an image.
Args:
pts1 (ndarray): xy coordinates of starting points
pts2 (ndarray): corresponding xy coordinates of ending points
color (str | List):
color code or a list of colors for each line segment
colorspace (str, default='rgb'): colorspace of image
thickness (int, default=1)
lineType (int, default=cv2.LINE_AA) option for cv2.line
Returns:
ndarray: the modified image (inplace if possible)
Example:
>>> from kwimage.im_draw import * # NOQA
>>> pts1 = np.array([[2, 0], [2, 20], [2.5, 30]])
>>> pts2 = np.array([[10, 5], [30, 28], [100, 50]])
>>> img = np.ones((100, 100, 3), dtype=np.uint8) * 255
>>> color = 'blue'
>>> colorspace = 'rgb'
>>> img2 = draw_line_segments_on_image(img, pts1, pts2, thickness=2)
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl() # xdoc: +SKIP
>>> kwplot.figure(doclf=True, fnum=1)
>>> kwplot.imshow(img2)
Example:
>>> import kwimage
>>> pts1 = kwimage.Points.random(10).scale(512).xy
>>> pts2 = kwimage.Points.random(10).scale(512).xy
>>> img = np.ones((512, 512, 3), dtype=np.uint8) * 255
>>> color = kwimage.Color.distinct(10)
>>> img2 = kwimage.draw_line_segments_on_image(img, pts1, pts2, color=color)
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl() # xdoc: +SKIP
>>> kwplot.figure(doclf=True, fnum=1)
>>> kwplot.imshow(img2)
"""
import cv2
# color = kwimage.Color(color)._forimage(img, colorspace)
num = len(pts1)
colors = _broadcast_colors(color, num, img, colorspace)
if 'lineType' not in kwargs:
kwargs['lineType'] = cv2.LINE_AA
pts1_ = pts1.tolist()
pts2_ = pts2.tolist()
for xy1, xy2, col in zip(pts1_, pts2_, colors):
xy1 = tuple(map(int, xy1))
xy2 = tuple(map(int, xy2))
cv2.line(img, xy1, xy2, color=col, thickness=thickness, **kwargs)
return img
[docs]def _broadcast_colors(color, num, img, colorspace):
"""
Determine if color applies a single color to all ``num`` items, or if it is
a list of colors for each item. Return as a list of colors for each item.
TODO:
- [ ] add as classmethod of kwimage.Color
Example:
>>> img = (np.random.rand(512, 512, 3) * 255).astype(np.uint8)
>>> colorspace = 'rgb'
>>> color = color_str_list = ['red', 'green', 'blue']
>>> color_str = 'red'
>>> num = 3
>>> print(_broadcast_colors(color_str_list, num, img, colorspace))
>>> print(_broadcast_colors(color_str, num, img, colorspace))
>>> colors_tuple_list = _broadcast_colors(color_str_list, num, img, colorspace)
>>> print(_broadcast_colors(colors_tuple_list, num, img, colorspace))
>>> #
>>> # FIXME: This case seems broken
>>> colors_ndarray_list = np.array(_broadcast_colors(color_str_list, num, img, colorspace))
>>> print(_broadcast_colors(colors_ndarray_list, num, img, colorspace))
"""
# Note there is an ambiguity when num=3 and color=[int, int, int]
# that must be resolved by checking num channels in the image
import kwimage
import ubelt as ub
import numbers
needs_broadcast = True # assume the list wasnt given by default
if ub.iterable(color):
first = ub.peek(color)
if len(color) == num:
if len(color) <= 4 and isinstance(first, numbers.Number):
# ambiguous case, interpret as a single broadcastable color
needs_broadcast = True
else:
# This is the only case we dont need broadcast
needs_broadcast = False
if needs_broadcast:
color = kwimage.Color(color)._forimage(img, colorspace)
colors = [color] * num
else:
colors = [kwimage.Color(c)._forimage(img, colorspace) for c in color]
return colors
[docs]def make_heatmask(probs, cmap='plasma', with_alpha=1.0, space='rgb',
dsize=None):
"""
Colorizes a single-channel intensity mask (with an alpha channel)
Args:
probs (ndarray): 2D probability map with values between 0 and 1
cmap (str): mpl colormap
with_alpha (float): between 0 and 1, uses probs as the alpha multipled
by this number.
space (str): output colorspace
dsize (tuple): if not None, then output is resized to W,H=dsize
SeeAlso:
kwimage.overlay_alpha_images
Example:
>>> # xdoc: +REQUIRES(module:matplotlib)
>>> probs = np.tile(np.linspace(0, 1, 10), (10, 1))
>>> heatmask = make_heatmask(probs, with_alpha=0.8, dsize=(100, 100))
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.imshow(heatmask, fnum=1, doclf=True, colorspace='rgb')
>>> kwplot.show_if_requested()
"""
import kwimage
import matplotlib as mpl
import matplotlib.cm # NOQA
assert len(probs.shape) == 2
cmap_ = mpl.cm.get_cmap(cmap)
probs = kwimage.ensure_float01(probs)
heatmask = cmap_(probs).astype(np.float32)
heatmask = kwimage.convert_colorspace(heatmask, 'rgba', space, implicit=True)
if with_alpha is not False and with_alpha is not None:
heatmask[:, :, 3] = (probs * with_alpha) # assign probs to alpha channel
if dsize is not None:
import cv2
heatmask = cv2.resize(
heatmask, tuple(dsize),
interpolation=cv2.INTER_NEAREST)
return heatmask
[docs]def make_orimask(radians, mag=None, alpha=1.0):
"""
Makes a colormap in HSV space where the orientation changes color and mag
changes the saturation/value.
Args:
radians (ndarray): orientation in radians
mag (ndarray): magnitude (must be normalized between 0 and 1)
alpha (float | ndarray):
if False or None, then the image is returned without alpha
if a float, then mag is scaled by this and used as the alpha channel
if an ndarray, then this is explicilty set as the alpha channel
Returns:
ndarray[float32]: an rgb / rgba image in 01 space
SeeAlso:
kwimage.overlay_alpha_images
Example:
>>> # xdoc: +REQUIRES(module:matplotlib)
>>> x, y = np.meshgrid(np.arange(64), np.arange(64))
>>> dx, dy = x - 32, y - 32
>>> radians = np.arctan2(dx, dy)
>>> mag = np.sqrt(dx ** 2 + dy ** 2)
>>> orimask = make_orimask(radians, mag)
>>> # xdoc: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.imshow(orimask, fnum=1, doclf=True, colorspace='rgb')
>>> kwplot.show_if_requested()
"""
import matplotlib as mpl
import matplotlib.cm # NOQA
TAU = np.pi * 2
# Map radians to 0 to 1
ori01 = (radians % TAU) / TAU
cmap_ = mpl.cm.get_cmap('hsv')
color_rgb = cmap_(ori01)[..., 0:3].astype(np.float32)
if mag is not None:
import kwimage
if mag.max() > 1:
mag = mag / mag.max()
color_hsv = kwimage.convert_colorspace(color_rgb, 'rgb', 'hsv')
color_hsv[..., 1:3] = mag[..., None]
color_rgb = kwimage.convert_colorspace(color_hsv, 'hsv', 'rgb')
else:
mag = 1
orimask = np.array(color_rgb, dtype=np.float32)
if isinstance(alpha, np.ndarray):
# Alpha specified as explicit numpy array
orimask = kwimage.ensure_alpha_channel(orimask)
orimask[:, :, 3] = alpha
elif alpha is not False and alpha is not None:
orimask = kwimage.ensure_alpha_channel(orimask)
orimask[:, :, 3] = mag * alpha
return orimask
[docs]def make_vector_field(dx, dy, stride=0.02, thresh=0.0, scale=1.0, alpha=1.0,
color='red', thickness=1, tipLength=0.1, line_type='aa'):
"""
Create an image representing a 2D vector field.
Args:
dx (ndarray): grid of vector x components
dy (ndarray): grid of vector y components
stride (int | float): sparsity of vectors, int specifies stride step in
pixels, a float specifies it as a percentage.
thresh (float): only plot vectors with magnitude greater than thres
scale (float): multiply magnitude for easier visualization
alpha (float): alpha value for vectors. Non-vector regions receive 0
alpha (if False, no alpha channel is used)
color (str | tuple | kwimage.Color): RGB color of the vectors
thickness (int, default=1): thickness of arrows
tipLength (float, default=0.1): fraction of line length
line_type (int): either cv2.LINE_4, cv2.LINE_8, or cv2.LINE_AA
Returns:
ndarray[float32]: vec_img: an rgb/rgba image in 0-1 space
SeeAlso:
kwimage.overlay_alpha_images
DEPRECATED USE: draw_vector_field instead
Example:
>>> x, y = np.meshgrid(np.arange(512), np.arange(512))
>>> dx, dy = x - 256.01, y - 256.01
>>> radians = np.arctan2(dx, dy)
>>> mag = np.sqrt(dx ** 2 + dy ** 2)
>>> dx, dy = dx / mag, dy / mag
>>> img = make_vector_field(dx, dy, scale=10, alpha=False)
>>> # xdoctest: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> kwplot.imshow(img)
>>> kwplot.show_if_requested()
"""
import warnings
warnings.warn('Deprecated, use draw_vector_field instead', DeprecationWarning)
import cv2
import kwimage
color = kwimage.Color(color).as255('rgb')
vecmask = np.zeros(dx.shape + (3,), dtype=np.uint8)
line_type_lookup = {'aa': cv2.LINE_AA}
line_type = line_type_lookup.get(line_type, line_type)
width = dx.shape[1]
height = dy.shape[0]
x_grid = np.arange(0, width, 1)
y_grid = np.arange(0, height, 1)
# Vector locations and directions
X, Y = np.meshgrid(x_grid, y_grid)
U, V = dx, dy
XYUV = [X, Y, U, V]
if isinstance(stride, float):
if stride < 0 or stride > 1:
raise ValueError('Floating point strides must be between 0 and 1')
stride = int(np.ceil(stride * min(width, height)))
# stride the points
if stride is not None and stride > 1:
XYUV = [a[::stride, ::stride] for a in XYUV]
# flatten the points
XYUV = [a.ravel() for a in XYUV]
# Filter out points with low magnitudes
if thresh is not None and thresh > 0:
M = np.sqrt((XYUV[2] ** 2) + (XYUV[3] ** 2)).ravel()
XYUV = np.array(XYUV)
flags = M > thresh
XYUV = [a[flags] for a in XYUV]
# Adjust vector magnitude for visibility
if scale is not None:
XYUV[2] *= scale
XYUV[3] *= scale
for (x, y, u, v) in zip(*XYUV):
pt1 = (int(x), int(y))
pt2 = tuple(map(int, map(np.round, (x + u, y + v))))
cv2.arrowedLine(vecmask, pt1, pt2, color=color, thickness=thickness,
tipLength=tipLength,
line_type=line_type)
vecmask = kwimage.ensure_float01(vecmask)
if isinstance(alpha, np.ndarray):
# Alpha specified as explicit numpy array
vecmask = kwimage.ensure_alpha_channel(vecmask)
vecmask[:, :, 3] = alpha
elif alpha is not False and alpha is not None:
# Alpha specified as a scale factor
vecmask = kwimage.ensure_alpha_channel(vecmask)
# vecmask[:, :, 3] = (vecmask[:, :, 0:3].sum(axis=2) > 0) * alpha
vecmask[:, :, 3] = vecmask[:, :, 0:3].sum(axis=2) * alpha
return vecmask
[docs]def draw_vector_field(image, dx, dy, stride=0.02, thresh=0.0, scale=1.0,
alpha=1.0, color='red', thickness=1, tipLength=0.1,
line_type='aa'):
"""
Create an image representing a 2D vector field.
Args:
image (ndarray): image to draw on
dx (ndarray): grid of vector x components
dy (ndarray): grid of vector y components
stride (int | float): sparsity of vectors, int specifies stride step in
pixels, a float specifies it as a percentage.
thresh (float): only plot vectors with magnitude greater than thres
scale (float): multiply magnitude for easier visualization
alpha (float): alpha value for vectors. Non-vector regions receive 0
alpha (if False, no alpha channel is used)
color (str | tuple | kwimage.Color): RGB color of the vectors
thickness (int, default=1): thickness of arrows
tipLength (float, default=0.1): fraction of line length
line_type (int): either cv2.LINE_4, cv2.LINE_8, or cv2.LINE_AA
Returns:
ndarray[float32]: The image with vectors overlaid. If image=None, then an
rgb/a image is created and returned.
Example:
>>> import kwimage
>>> width, height = 512, 512
>>> image = kwimage.grab_test_image(dsize=(width, height))
>>> x, y = np.meshgrid(np.arange(height), np.arange(width))
>>> dx, dy = x - width / 2, y - height / 2
>>> radians = np.arctan2(dx, dy)
>>> mag = np.sqrt(dx ** 2 + dy ** 2) + 1e-3
>>> dx, dy = dx / mag, dy / mag
>>> img = kwimage.draw_vector_field(image, dx, dy, scale=10, alpha=False)
>>> # xdoctest: +REQUIRES(--show)
>>> import kwplot
>>> kwplot.autompl()
>>> kwplot.imshow(img)
>>> kwplot.show_if_requested()
"""
import cv2
import kwimage
if image is None:
# Create a default image
image = np.zeros(dx.shape + (3,), dtype=np.uint8)
# image = kwimage.atleast_3channels(image)
color = kwimage.Color(color)._forimage(image)
line_type_lookup = {'aa': cv2.LINE_AA}
line_type = line_type_lookup.get(line_type, line_type)
height, width = dx.shape[0:2]
x_grid = np.arange(0, width, 1)
y_grid = np.arange(0, height, 1)
# Vector locations and directions
X, Y = np.meshgrid(x_grid, y_grid)
U, V = dx, dy
XYUV = [X, Y, U, V]
if isinstance(stride, float):
if stride < 0 or stride > 1:
raise ValueError('Floating point strides must be between 0 and 1')
stride = int(np.ceil(stride * min(width, height)))
# stride the points
if stride is not None and stride > 1:
XYUV = [a[::stride, ::stride] for a in XYUV]
# flatten the points
XYUV = [a.ravel() for a in XYUV]
# Filter out points with low magnitudes
if thresh is not None and thresh > 0:
M = np.sqrt((XYUV[2] ** 2) + (XYUV[3] ** 2)).ravel()
XYUV = np.array(XYUV)
flags = M > thresh
XYUV = [a[flags] for a in XYUV]
# Adjust vector magnitude for visibility
if scale is not None:
XYUV[2] *= scale
XYUV[3] *= scale
if alpha is not None and alpha is not False and alpha != 1:
raise NotImplementedError
for (x, y, u, v) in zip(*XYUV):
pt1 = (int(x), int(y))
pt2 = tuple(map(int, map(np.round, (x + u, y + v))))
cv2.arrowedLine(image, pt1, pt2, color=color, thickness=thickness,
tipLength=tipLength,
line_type=line_type)
if isinstance(alpha, np.ndarray):
# Alpha specified as explicit numpy array
image = kwimage.ensure_float01(image)
image = kwimage.ensure_alpha_channel(image)
image[:, :, 3] = alpha
elif alpha is not False and alpha is not None:
# Alpha specified as a scale factor
image = kwimage.ensure_float01(image)
image = kwimage.ensure_alpha_channel(image)
# image[:, :, 3] = (image[:, :, 0:3].sum(axis=2) > 0) * alpha
image[:, :, 3] = image[:, :, 0:3].sum(axis=2) * alpha
return image