import numpy as np # noqa
from edflow.util import walk # noqa
[docs]def flow2hsv(flow):
"""Given a Flowmap of shape ``[W, H, 2]`` calculates an hsv image,
showing the relative magnitude and direction of the optical flow.
Parameters
---------
flow : np.array
Optical flow with shape ``[W, H, 2]``.
Returns
-------
np.array
Containing the hsv data.
"""
import matplotlib as mpl
mpl.use("Agg")
# prepare array - value is always at max
hsv = np.zeros((flow.shape[0], flow.shape[1], 3), dtype=np.uint8)
hsv[:, :, 0] = 255
hsv[:, :, 2] = 255
# magnitude and angle
mag, ang = cart2polar(flow[..., 0], flow[..., 1])
# make it colorful
hsv[..., 0] = ang * 180 / np.pi
normalizer = mpl.colors.Normalize(mag.min(), mag.max())
hsv[..., 1] = np.int32(normalizer(mag) * 255)
return hsv
[docs]def cart2polar(x, y):
"""
Takes two array as x and y coordinates and returns the magnitude and angle.
"""
r = np.sqrt(x ** 2 + y ** 2)
phi = np.arctan2(x, y)
return r, phi
[docs]def hsv2rgb(hsv):
"""color space conversion hsv -> rgb. simple wrapper for nice name."""
import matplotlib as mpl
mpl.use("Agg")
rgb = mpl.colors.hsv_to_rgb(hsv)
return rgb
[docs]def flow2rgb(flow):
"""converts a flow field to an rgb color image.
Parameters
---------
flow : np.array
optical flow with shape ``[W, H, 2]``.
Returns
-------
np.array
Containing the rgb data. Color indicates orientation,
intensity indicates magnitude.
"""
return hsv2rgb(flow2hsv(flow))
[docs]def get_support(image):
"""
... warning: This function makes a lot of assumptions that need not be met!
Assuming that there are three categories of images and that the image_array
has been properly constructed, this function will estimate the support of
the given :attr:`image`.
Parameters
---------
image : np.ndarray
Some properly constructed image like array. No
assumptions need to be made about the shape of the image, we
simply assme each value is some color value.
Returns
-------
str
The support. Either '0->1', '-1->1' or '0->255'
"""
if image.min() < 0:
return "-1->1"
elif image.max() > 1:
return "0->255"
else:
return "0->1"
VALID_SUPPORTS = ["0->1", "-1->1", "0->255"]
[docs]def sup_str_to_num(support_str):
"""Converts a support string into usable numbers."""
mn = -1.0 if support_str == "-1->1" else 0.0
mx = 255.0 if support_str == "0->255" else 1.0
return mn, mx
[docs]def adjust_support(image, future_support, current_support=None, clip=False):
"""Will adjust the support of all color values in :attr:`image`.
Parameters
---------
image : np.ndarray
Array containing color values. Make sure this is
properly constructed.
future_support : str
The support this array is supposed to have after
the transformation. Must be one of '-1->1', '0->1', or '0->255'.
current_support : str
The support of the colors currentl in
:attr:`image`. If not given it will be estimated by
:func:`get_support`.
clip : bool
By default the return values in image are simply coming
from a linear transform, thus the actual support might be larger
than the requested interval. If set to ``True`` the returned
array will be cliped to ``future_support``.
Returns
-------
same type as image
The given :attr:`image` with transformed support.
"""
if current_support is None:
current_support = get_support(image)
else:
assert current_support in VALID_SUPPORTS
cur_min, cur_max = sup_str_to_num(current_support)
fut_min, fut_max = sup_str_to_num(future_support)
# To [0, 1]
image = image.astype(float)
image -= cur_min
image /= cur_max - cur_min
# To [fut_min, fut_max]
image *= fut_max - fut_min
image += fut_min
if clip:
image = clip_to_support(image, future_support)
if future_support == "0->255":
image = image.astype(np.uint8)
return image
[docs]def clip_to_support(image, supp_str):
vmin, vmax = sup_str_to_num(supp_str)
return np.clip(image, vmin, vmax)
[docs]def add_im_info(image, ax):
"""Adds some interesting facts about the image."""
shape = "x".join([str(s) for s in image.shape])
mn = image.min()
mx = image.max()
supp = get_support(image)
info_str = "shape: {}\nsupport: {} (min={}, max={})"
info_str = info_str.format(shape, supp, mn, mx)
ax.text(0, 0, info_str)
[docs]def im_fn(key, im, ax):
"""Plot an image. Used by :func:`plot_datum`."""
if im.shape[-1] == 1:
im = np.squeeze(im)
add_im_info(im, ax)
ax.imshow(adjust_support(im, "0->1"))
ax.set_ylabel(key, rotation=0)
[docs]def heatmap_fn(key, im, ax):
"""Assumes that heatmap shape is [H, W, N]. Used by
:func:`plot_datum`."""
im = np.mean(im, axis=-1)
im_fn(key, im, ax)
[docs]def keypoints_fn(key, keypoints, ax):
"""
Plots a list of keypoints as a dot plot.
"""
add_im_info(keypoints, ax)
x = keypoints[:, 0]
y = keypoints[:, 1]
ax.plot(x, y, "go", markersize=1)
ax.set_ylabel(key)
[docs]def flow_fn(key, im, ax):
"""Plot an flow. Used by :func:`plot_datum`."""
im = flow2rgb(im)
im_fn(key, im, ax)
[docs]def other_fn(key, obj, ax):
"""Print some text about the object. Used by :func:`plot_datum`."""
text = "{}: {} - {}".format(key, type(obj), obj)
ax.axis("off")
# ax.imshow(np.ones([10, 100]))
ax.text(0, 0, text)
PLOT_FUNCTIONS = {
"image": im_fn,
"heat": heatmap_fn,
"keypoints": keypoints_fn,
"flow": flow_fn,
"other": other_fn,
}
[docs]def default_heuristic(key, obj):
"""Determines the kind of an object. Used by :func:`plot_datum`."""
if isinstance(obj, np.ndarray):
if len(obj.shape) > 3 or len(obj.shape) < 2:
# This is no image -> Maybe later implement sequence fn
return "other"
else:
if obj.shape[-1] in [3, 4]:
return "image"
elif obj.shape[-1] == 2:
if len(obj.shape) <= 2:
return "keypoints"
else:
return "flow"
else:
return "heat"
return "other"
[docs]def plot_datum(
nested_thing,
savename="datum.png",
heuristics=default_heuristic,
plt_functions=PLOT_FUNCTIONS,
):
"""Plots all data in the nested_thing as best as can.
If heuristics is given, this determines how each leaf datum is converted
to something plottable.
Parameters
---------
nested_thing : dict or list
Some nested object.
savename : str
``Path/to/the/plot.png``.
heuristics : Callable
If given this should produce a string specifying
the kind of data of the leaf. If ``None`` determinde automatically.
See :func:`default_heuristic`.
plt_functions : dict of Callables
Maps a ``kind`` to a function which
can plot it. Each callable must be able to receive a the key, the
leaf object and the Axes to plot it in.
"""
import matplotlib.pyplot as plt # noqa
import matplotlib.gridspec as gridspec # noqa
class Plotter(object):
def __init__(self, kind_fn, savename):
self.kind_fn = kind_fn
self.savename = savename
self.buffer = []
def __call__(self, key, obj):
kind = self.kind_fn(key, obj)
self.buffer += [[kind, key, obj]]
def plot(self):
n_pl = len(self.buffer)
f = plt.figure(figsize=(5, 2 * n_pl))
gs = gridspec.GridSpec(n_pl, 1)
for i, [kind, key, obj] in enumerate(self.buffer):
ax = f.add_subplot(gs[i])
plt_functions[kind](key, obj, ax)
f.savefig(self.savename)
def __str__(self):
self.plot()
return "Saved Plot at {}".format(self.savename)
P = Plotter(heuristics, savename)
walk(nested_thing, P, pass_key=True)
print(P)