"""Visualize image."""
import numpy as np
import mxnet as mx
[docs]def plot_image(img, ax=None, reverse_rgb=False):
"""Visualize image.
Parameters
----------
img : numpy.ndarray or mxnet.nd.NDArray
Image with shape `H, W, 3`.
ax : matplotlib axes, optional
You can reuse previous axes if provided.
reverse_rgb : bool, optional
Reverse RGB<->BGR orders if `True`.
Returns
-------
matplotlib axes
The ploted axes.
Examples
--------
from matplotlib import pyplot as plt
ax = plot_image(img)
plt.show()
"""
from matplotlib import pyplot as plt
if ax is None:
# create new axes
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
if isinstance(img, mx.nd.NDArray):
img = img.asnumpy()
img = img.copy()
if reverse_rgb:
img[:, :, (0, 1, 2)] = img[:, :, (2, 1, 0)]
ax.imshow(img.astype(np.uint8))
return ax