image-matching-webui / extra_utils /visualize_util.py
Vincentqyw
add: files
9223079
raw
history blame
22.3 kB
""" Organize some frequently used visualization functions. """
import cv2
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import copy
import seaborn as sns
# Plot junctions onto the image (return a separate copy)
def plot_junctions(input_image, junctions, junc_size=3, color=None):
"""
input_image: can be 0~1 float or 0~255 uint8.
junctions: Nx2 or 2xN np array.
junc_size: the size of the plotted circles.
"""
# Create image copy
image = copy.copy(input_image)
# Make sure the image is converted to 255 uint8
if image.dtype == np.uint8:
pass
# A float type image ranging from 0~1
elif image.dtype in [np.float32, np.float64, np.float] and image.max() <= 2.0:
image = (image * 255.0).astype(np.uint8)
# A float type image ranging from 0.~255.
elif image.dtype in [np.float32, np.float64, np.float] and image.mean() > 10.0:
image = image.astype(np.uint8)
else:
raise ValueError(
"[Error] Unknown image data type. Expect 0~1 float or 0~255 uint8."
)
# Check whether the image is single channel
if len(image.shape) == 2 or ((len(image.shape) == 3) and (image.shape[-1] == 1)):
# Squeeze to H*W first
image = image.squeeze()
# Stack to channle 3
image = np.concatenate([image[..., None] for _ in range(3)], axis=-1)
# Junction dimensions should be N*2
if not len(junctions.shape) == 2:
raise ValueError("[Error] junctions should be 2-dim array.")
# Always convert to N*2
if junctions.shape[-1] != 2:
if junctions.shape[0] == 2:
junctions = junctions.T
else:
raise ValueError("[Error] At least one of the two dims should be 2.")
# Round and convert junctions to int (and check the boundary)
H, W = image.shape[:2]
junctions = (np.round(junctions)).astype(np.int)
junctions[junctions < 0] = 0
junctions[junctions[:, 0] >= H, 0] = H - 1 # (first dim) max bounded by H-1
junctions[junctions[:, 1] >= W, 1] = W - 1 # (second dim) max bounded by W-1
# Iterate through all the junctions
num_junc = junctions.shape[0]
if color is None:
color = (0, 255.0, 0)
for idx in range(num_junc):
# Fetch one junction
junc = junctions[idx, :]
cv2.circle(
image, tuple(np.flip(junc)), radius=junc_size, color=color, thickness=3
)
return image
# Plot line segements given junctions and line adjecent map
def plot_line_segments(
input_image,
junctions,
line_map,
junc_size=3,
color=(0, 255.0, 0),
line_width=1,
plot_survived_junc=True,
):
"""
input_image: can be 0~1 float or 0~255 uint8.
junctions: Nx2 or 2xN np array.
line_map: NxN np array
junc_size: the size of the plotted circles.
color: color of the line segments (can be string "random")
line_width: width of the drawn segments.
plot_survived_junc: whether we only plot the survived junctions.
"""
# Create image copy
image = copy.copy(input_image)
# Make sure the image is converted to 255 uint8
if image.dtype == np.uint8:
pass
# A float type image ranging from 0~1
elif image.dtype in [np.float32, np.float64, np.float] and image.max() <= 2.0:
image = (image * 255.0).astype(np.uint8)
# A float type image ranging from 0.~255.
elif image.dtype in [np.float32, np.float64, np.float] and image.mean() > 10.0:
image = image.astype(np.uint8)
else:
raise ValueError(
"[Error] Unknown image data type. Expect 0~1 float or 0~255 uint8."
)
# Check whether the image is single channel
if len(image.shape) == 2 or ((len(image.shape) == 3) and (image.shape[-1] == 1)):
# Squeeze to H*W first
image = image.squeeze()
# Stack to channle 3
image = np.concatenate([image[..., None] for _ in range(3)], axis=-1)
# Junction dimensions should be 2
if not len(junctions.shape) == 2:
raise ValueError("[Error] junctions should be 2-dim array.")
# Always convert to N*2
if junctions.shape[-1] != 2:
if junctions.shape[0] == 2:
junctions = junctions.T
else:
raise ValueError("[Error] At least one of the two dims should be 2.")
# line_map dimension should be 2
if not len(line_map.shape) == 2:
raise ValueError("[Error] line_map should be 2-dim array.")
# Color should be "random" or a list or tuple with length 3
if color != "random":
if not (isinstance(color, tuple) or isinstance(color, list)):
raise ValueError("[Error] color should have type list or tuple.")
else:
if len(color) != 3:
raise ValueError(
"[Error] color should be a list or tuple with length 3."
)
# Make a copy of the line_map
line_map_tmp = copy.copy(line_map)
# Parse line_map back to segment pairs
segments = np.zeros([0, 4])
for idx in range(junctions.shape[0]):
# if no connectivity, just skip it
if line_map_tmp[idx, :].sum() == 0:
continue
# record the line segment
else:
for idx2 in np.where(line_map_tmp[idx, :] == 1)[0]:
p1 = np.flip(junctions[idx, :]) # Convert to xy format
p2 = np.flip(junctions[idx2, :]) # Convert to xy format
segments = np.concatenate(
(segments, np.array([p1[0], p1[1], p2[0], p2[1]])[None, ...]),
axis=0,
)
# Update line_map
line_map_tmp[idx, idx2] = 0
line_map_tmp[idx2, idx] = 0
# Draw segment pairs
for idx in range(segments.shape[0]):
seg = np.round(segments[idx, :]).astype(np.int)
# Decide the color
if color != "random":
color = tuple(color)
else:
color = tuple(
np.random.rand(
3,
)
)
cv2.line(
image, tuple(seg[:2]), tuple(seg[2:]), color=color, thickness=line_width
)
# Also draw the junctions
if not plot_survived_junc:
num_junc = junctions.shape[0]
for idx in range(num_junc):
# Fetch one junction
junc = junctions[idx, :]
cv2.circle(
image,
tuple(np.flip(junc)),
radius=junc_size,
color=(0, 255.0, 0),
thickness=3,
)
# Only plot the junctions which are part of a line segment
else:
for idx in range(segments.shape[0]):
seg = np.round(segments[idx, :]).astype(np.int) # Already in HW format.
cv2.circle(
image,
tuple(seg[:2]),
radius=junc_size,
color=(0, 255.0, 0),
thickness=3,
)
cv2.circle(
image,
tuple(seg[2:]),
radius=junc_size,
color=(0, 255.0, 0),
thickness=3,
)
return image
# Plot line segments given Nx4 or Nx2x2 line segments
def plot_line_segments_from_segments(
input_image, line_segments, junc_size=3, color=(0, 255.0, 0), line_width=1
):
# Create image copy
image = copy.copy(input_image)
# Make sure the image is converted to 255 uint8
if image.dtype == np.uint8:
pass
# A float type image ranging from 0~1
elif image.dtype in [np.float32, np.float64, np.float] and image.max() <= 2.0:
image = (image * 255.0).astype(np.uint8)
# A float type image ranging from 0.~255.
elif image.dtype in [np.float32, np.float64, np.float] and image.mean() > 10.0:
image = image.astype(np.uint8)
else:
raise ValueError(
"[Error] Unknown image data type. Expect 0~1 float or 0~255 uint8."
)
# Check whether the image is single channel
if len(image.shape) == 2 or ((len(image.shape) == 3) and (image.shape[-1] == 1)):
# Squeeze to H*W first
image = image.squeeze()
# Stack to channle 3
image = np.concatenate([image[..., None] for _ in range(3)], axis=-1)
# Check the if line_segments are in (1) Nx4, or (2) Nx2x2.
H, W, _ = image.shape
# (1) Nx4 format
if len(line_segments.shape) == 2 and line_segments.shape[-1] == 4:
# Round to int32
line_segments = line_segments.astype(np.int32)
# Clip H dimension
line_segments[:, 0] = np.clip(line_segments[:, 0], a_min=0, a_max=H - 1)
line_segments[:, 2] = np.clip(line_segments[:, 2], a_min=0, a_max=H - 1)
# Clip W dimension
line_segments[:, 1] = np.clip(line_segments[:, 1], a_min=0, a_max=W - 1)
line_segments[:, 3] = np.clip(line_segments[:, 3], a_min=0, a_max=W - 1)
# Convert to Nx2x2 format
line_segments = np.concatenate(
[
np.expand_dims(line_segments[:, :2], axis=1),
np.expand_dims(line_segments[:, 2:], axis=1),
],
axis=1,
)
# (2) Nx2x2 format
elif len(line_segments.shape) == 3 and line_segments.shape[-1] == 2:
# Round to int32
line_segments = line_segments.astype(np.int32)
# Clip H dimension
line_segments[:, :, 0] = np.clip(line_segments[:, :, 0], a_min=0, a_max=H - 1)
line_segments[:, :, 1] = np.clip(line_segments[:, :, 1], a_min=0, a_max=W - 1)
else:
raise ValueError(
"[Error] line_segments should be either Nx4 or Nx2x2 in HW format."
)
# Draw segment pairs (all segments should be in HW format)
image = image.copy()
for idx in range(line_segments.shape[0]):
seg = np.round(line_segments[idx, :, :]).astype(np.int32)
# Decide the color
if color != "random":
color = tuple(color)
else:
color = tuple(
np.random.rand(
3,
)
)
cv2.line(
image,
tuple(np.flip(seg[0, :])),
tuple(np.flip(seg[1, :])),
color=color,
thickness=line_width,
)
# Also draw the junctions
cv2.circle(
image,
tuple(np.flip(seg[0, :])),
radius=junc_size,
color=(0, 255.0, 0),
thickness=3,
)
cv2.circle(
image,
tuple(np.flip(seg[1, :])),
radius=junc_size,
color=(0, 255.0, 0),
thickness=3,
)
return image
# Additional functions to visualize multiple images at the same time,
# e.g. for line matching
def plot_images(imgs, titles=None, cmaps="gray", dpi=100, size=5, pad=0.5):
"""Plot a set of images horizontally.
Args:
imgs: a list of NumPy or PyTorch images, RGB (H, W, 3) or mono (H, W).
titles: a list of strings, as titles for each image.
cmaps: colormaps for monochrome images.
"""
n = len(imgs)
if not isinstance(cmaps, (list, tuple)):
cmaps = [cmaps] * n
# figsize = (size*n, size*3/4) if size is not None else None
figsize = (size * n, size * 6 / 5) if size is not None else None
fig, ax = plt.subplots(1, n, figsize=figsize, dpi=dpi)
if n == 1:
ax = [ax]
for i in range(n):
ax[i].imshow(imgs[i], cmap=plt.get_cmap(cmaps[i]))
ax[i].get_yaxis().set_ticks([])
ax[i].get_xaxis().set_ticks([])
ax[i].set_axis_off()
for spine in ax[i].spines.values(): # remove frame
spine.set_visible(False)
if titles:
ax[i].set_title(titles[i])
fig.tight_layout(pad=pad)
return fig
def plot_keypoints(kpts, colors="lime", ps=4):
"""Plot keypoints for existing images.
Args:
kpts: list of ndarrays of size (N, 2).
colors: string, or list of list of tuples (one for each keypoints).
ps: size of the keypoints as float.
"""
if not isinstance(colors, list):
colors = [colors] * len(kpts)
axes = plt.gcf().axes
for a, k, c in zip(axes, kpts, colors):
a.scatter(k[:, 0], k[:, 1], c=c, s=ps, linewidths=0)
def plot_matches(kpts0, kpts1, color=None, lw=1.5, ps=4, indices=(0, 1), a=1.0):
"""Plot matches for a pair of existing images.
Args:
kpts0, kpts1: corresponding keypoints of size (N, 2).
color: color of each match, string or RGB tuple. Random if not given.
lw: width of the lines.
ps: size of the end points (no endpoint if ps=0)
indices: indices of the images to draw the matches on.
a: alpha opacity of the match lines.
"""
fig = plt.gcf()
ax = fig.axes
assert len(ax) > max(indices)
ax0, ax1 = ax[indices[0]], ax[indices[1]]
fig.canvas.draw()
assert len(kpts0) == len(kpts1)
if color is None:
color = matplotlib.cm.hsv(np.random.rand(len(kpts0))).tolist()
elif len(color) > 0 and not isinstance(color[0], (tuple, list)):
color = [color] * len(kpts0)
if lw > 0:
# transform the points into the figure coordinate system
transFigure = fig.transFigure.inverted()
fkpts0 = transFigure.transform(ax0.transData.transform(kpts0))
fkpts1 = transFigure.transform(ax1.transData.transform(kpts1))
fig.lines += [
matplotlib.lines.Line2D(
(fkpts0[i, 0], fkpts1[i, 0]),
(fkpts0[i, 1], fkpts1[i, 1]),
zorder=1,
transform=fig.transFigure,
c=color[i],
linewidth=lw,
alpha=a,
)
for i in range(len(kpts0))
]
# freeze the axes to prevent the transform to change
ax0.autoscale(enable=False)
ax1.autoscale(enable=False)
if ps > 0:
ax0.scatter(kpts0[:, 0], kpts0[:, 1], c=color, s=ps, zorder=2)
ax1.scatter(kpts1[:, 0], kpts1[:, 1], c=color, s=ps, zorder=2)
def plot_lines(
lines, line_colors="orange", point_colors="cyan", ps=4, lw=2, indices=(0, 1)
):
"""Plot lines and endpoints for existing images.
Args:
lines: list of ndarrays of size (N, 2, 2).
colors: string, or list of list of tuples (one for each keypoints).
ps: size of the keypoints as float pixels.
lw: line width as float pixels.
indices: indices of the images to draw the matches on.
"""
if not isinstance(line_colors, list):
line_colors = [line_colors] * len(lines)
if not isinstance(point_colors, list):
point_colors = [point_colors] * len(lines)
fig = plt.gcf()
ax = fig.axes
assert len(ax) > max(indices)
axes = [ax[i] for i in indices]
fig.canvas.draw()
# Plot the lines and junctions
for a, l, lc, pc in zip(axes, lines, line_colors, point_colors):
for i in range(len(l)):
line = matplotlib.lines.Line2D(
(l[i, 0, 0], l[i, 1, 0]),
(l[i, 0, 1], l[i, 1, 1]),
zorder=1,
c=lc,
linewidth=lw,
)
a.add_line(line)
pts = l.reshape(-1, 2)
a.scatter(pts[:, 0], pts[:, 1], c=pc, s=ps, linewidths=0, zorder=2)
return fig
def plot_line_matches(kpts0, kpts1, color=None, lw=1.5, indices=(0, 1), a=1.0):
"""Plot matches for a pair of existing images, parametrized by their middle point.
Args:
kpts0, kpts1: corresponding middle points of the lines of size (N, 2).
color: color of each match, string or RGB tuple. Random if not given.
lw: width of the lines.
indices: indices of the images to draw the matches on.
a: alpha opacity of the match lines.
"""
fig = plt.gcf()
ax = fig.axes
assert len(ax) > max(indices)
ax0, ax1 = ax[indices[0]], ax[indices[1]]
fig.canvas.draw()
assert len(kpts0) == len(kpts1)
if color is None:
color = matplotlib.cm.hsv(np.random.rand(len(kpts0))).tolist()
elif len(color) > 0 and not isinstance(color[0], (tuple, list)):
color = [color] * len(kpts0)
if lw > 0:
# transform the points into the figure coordinate system
transFigure = fig.transFigure.inverted()
fkpts0 = transFigure.transform(ax0.transData.transform(kpts0))
fkpts1 = transFigure.transform(ax1.transData.transform(kpts1))
fig.lines += [
matplotlib.lines.Line2D(
(fkpts0[i, 0], fkpts1[i, 0]),
(fkpts0[i, 1], fkpts1[i, 1]),
zorder=1,
transform=fig.transFigure,
c=color[i],
linewidth=lw,
alpha=a,
)
for i in range(len(kpts0))
]
# freeze the axes to prevent the transform to change
ax0.autoscale(enable=False)
ax1.autoscale(enable=False)
def plot_color_line_matches(lines, correct_matches=None, lw=2, indices=(0, 1)):
"""Plot line matches for existing images with multiple colors.
Args:
lines: list of ndarrays of size (N, 2, 2).
correct_matches: bool array of size (N,) indicating correct matches.
lw: line width as float pixels.
indices: indices of the images to draw the matches on.
"""
n_lines = len(lines[0])
colors = sns.color_palette("husl", n_colors=n_lines)
np.random.shuffle(colors)
alphas = np.ones(n_lines)
# If correct_matches is not None, display wrong matches with a low alpha
if correct_matches is not None:
alphas[~np.array(correct_matches)] = 0.2
fig = plt.gcf()
ax = fig.axes
assert len(ax) > max(indices)
axes = [ax[i] for i in indices]
fig.canvas.draw()
# Plot the lines
for a, l in zip(axes, lines):
# Transform the points into the figure coordinate system
transFigure = fig.transFigure.inverted()
endpoint0 = transFigure.transform(a.transData.transform(l[:, 0]))
endpoint1 = transFigure.transform(a.transData.transform(l[:, 1]))
fig.lines += [
matplotlib.lines.Line2D(
(endpoint0[i, 0], endpoint1[i, 0]),
(endpoint0[i, 1], endpoint1[i, 1]),
zorder=1,
transform=fig.transFigure,
c=colors[i],
alpha=alphas[i],
linewidth=lw,
)
for i in range(n_lines)
]
return fig
def plot_color_lines(lines, correct_matches, wrong_matches, lw=2, indices=(0, 1)):
"""Plot line matches for existing images with multiple colors:
green for correct matches, red for wrong ones, and blue for the rest.
Args:
lines: list of ndarrays of size (N, 2, 2).
correct_matches: list of bool arrays of size N with correct matches.
wrong_matches: list of bool arrays of size (N,) with correct matches.
lw: line width as float pixels.
indices: indices of the images to draw the matches on.
"""
# palette = sns.color_palette()
palette = sns.color_palette("hls", 8)
blue = palette[5] # palette[0]
red = palette[0] # palette[3]
green = palette[2] # palette[2]
colors = [np.array([blue] * len(l)) for l in lines]
for i, c in enumerate(colors):
c[np.array(correct_matches[i])] = green
c[np.array(wrong_matches[i])] = red
fig = plt.gcf()
ax = fig.axes
assert len(ax) > max(indices)
axes = [ax[i] for i in indices]
fig.canvas.draw()
# Plot the lines
for a, l, c in zip(axes, lines, colors):
# Transform the points into the figure coordinate system
transFigure = fig.transFigure.inverted()
endpoint0 = transFigure.transform(a.transData.transform(l[:, 0]))
endpoint1 = transFigure.transform(a.transData.transform(l[:, 1]))
fig.lines += [
matplotlib.lines.Line2D(
(endpoint0[i, 0], endpoint1[i, 0]),
(endpoint0[i, 1], endpoint1[i, 1]),
zorder=1,
transform=fig.transFigure,
c=c[i],
linewidth=lw,
)
for i in range(len(l))
]
def plot_subsegment_matches(lines, subsegments, lw=2, indices=(0, 1)):
"""Plot line matches for existing images with multiple colors and
highlight the actually matched subsegments.
Args:
lines: list of ndarrays of size (N, 2, 2).
subsegments: list of ndarrays of size (N, 2, 2).
lw: line width as float pixels.
indices: indices of the images to draw the matches on.
"""
n_lines = len(lines[0])
colors = sns.cubehelix_palette(
start=2, rot=-0.2, dark=0.3, light=0.7, gamma=1.3, hue=1, n_colors=n_lines
)
fig = plt.gcf()
ax = fig.axes
assert len(ax) > max(indices)
axes = [ax[i] for i in indices]
fig.canvas.draw()
# Plot the lines
for a, l, ss in zip(axes, lines, subsegments):
# Transform the points into the figure coordinate system
transFigure = fig.transFigure.inverted()
# Draw full line
endpoint0 = transFigure.transform(a.transData.transform(l[:, 0]))
endpoint1 = transFigure.transform(a.transData.transform(l[:, 1]))
fig.lines += [
matplotlib.lines.Line2D(
(endpoint0[i, 0], endpoint1[i, 0]),
(endpoint0[i, 1], endpoint1[i, 1]),
zorder=1,
transform=fig.transFigure,
c="red",
alpha=0.7,
linewidth=lw,
)
for i in range(n_lines)
]
# Draw matched subsegment
endpoint0 = transFigure.transform(a.transData.transform(ss[:, 0]))
endpoint1 = transFigure.transform(a.transData.transform(ss[:, 1]))
fig.lines += [
matplotlib.lines.Line2D(
(endpoint0[i, 0], endpoint1[i, 0]),
(endpoint0[i, 1], endpoint1[i, 1]),
zorder=1,
transform=fig.transFigure,
c=colors[i],
alpha=1,
linewidth=lw,
)
for i in range(n_lines)
]