Realcat's picture
update: d2net lib
cb265a5
raw
history blame
5.22 kB
import matplotlib.pyplot as plt
import numpy as np
import torch
from lib.exceptions import EmptyTensorError
def preprocess_image(image, preprocessing=None):
image = image.astype(np.float32)
image = np.transpose(image, [2, 0, 1])
if preprocessing is None:
pass
elif preprocessing == 'caffe':
# RGB -> BGR
image = image[:: -1, :, :]
# Zero-center by mean pixel
mean = np.array([103.939, 116.779, 123.68])
image = image - mean.reshape([3, 1, 1])
elif preprocessing == 'torch':
image /= 255.0
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
image = (image - mean.reshape([3, 1, 1])) / std.reshape([3, 1, 1])
else:
raise ValueError('Unknown preprocessing parameter.')
return image
def imshow_image(image, preprocessing=None):
if preprocessing is None:
pass
elif preprocessing == 'caffe':
mean = np.array([103.939, 116.779, 123.68])
image = image + mean.reshape([3, 1, 1])
# RGB -> BGR
image = image[:: -1, :, :]
elif preprocessing == 'torch':
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
image = image * std.reshape([3, 1, 1]) + mean.reshape([3, 1, 1])
image *= 255.0
else:
raise ValueError('Unknown preprocessing parameter.')
image = np.transpose(image, [1, 2, 0])
image = np.round(image).astype(np.uint8)
return image
def grid_positions(h, w, device, matrix=False):
lines = torch.arange(
0, h, device=device
).view(-1, 1).float().repeat(1, w)
columns = torch.arange(
0, w, device=device
).view(1, -1).float().repeat(h, 1)
if matrix:
return torch.stack([lines, columns], dim=0)
else:
return torch.cat([lines.view(1, -1), columns.view(1, -1)], dim=0)
def upscale_positions(pos, scaling_steps=0):
for _ in range(scaling_steps):
pos = pos * 2 + 0.5
return pos
def downscale_positions(pos, scaling_steps=0):
for _ in range(scaling_steps):
pos = (pos - 0.5) / 2
return pos
def interpolate_dense_features(pos, dense_features, return_corners=False):
device = pos.device
ids = torch.arange(0, pos.size(1), device=device)
_, h, w = dense_features.size()
i = pos[0, :]
j = pos[1, :]
# Valid corners
i_top_left = torch.floor(i).long()
j_top_left = torch.floor(j).long()
valid_top_left = torch.min(i_top_left >= 0, j_top_left >= 0)
i_top_right = torch.floor(i).long()
j_top_right = torch.ceil(j).long()
valid_top_right = torch.min(i_top_right >= 0, j_top_right < w)
i_bottom_left = torch.ceil(i).long()
j_bottom_left = torch.floor(j).long()
valid_bottom_left = torch.min(i_bottom_left < h, j_bottom_left >= 0)
i_bottom_right = torch.ceil(i).long()
j_bottom_right = torch.ceil(j).long()
valid_bottom_right = torch.min(i_bottom_right < h, j_bottom_right < w)
valid_corners = torch.min(
torch.min(valid_top_left, valid_top_right),
torch.min(valid_bottom_left, valid_bottom_right)
)
i_top_left = i_top_left[valid_corners]
j_top_left = j_top_left[valid_corners]
i_top_right = i_top_right[valid_corners]
j_top_right = j_top_right[valid_corners]
i_bottom_left = i_bottom_left[valid_corners]
j_bottom_left = j_bottom_left[valid_corners]
i_bottom_right = i_bottom_right[valid_corners]
j_bottom_right = j_bottom_right[valid_corners]
ids = ids[valid_corners]
if ids.size(0) == 0:
raise EmptyTensorError
# Interpolation
i = i[ids]
j = j[ids]
dist_i_top_left = i - i_top_left.float()
dist_j_top_left = j - j_top_left.float()
w_top_left = (1 - dist_i_top_left) * (1 - dist_j_top_left)
w_top_right = (1 - dist_i_top_left) * dist_j_top_left
w_bottom_left = dist_i_top_left * (1 - dist_j_top_left)
w_bottom_right = dist_i_top_left * dist_j_top_left
descriptors = (
w_top_left * dense_features[:, i_top_left, j_top_left] +
w_top_right * dense_features[:, i_top_right, j_top_right] +
w_bottom_left * dense_features[:, i_bottom_left, j_bottom_left] +
w_bottom_right * dense_features[:, i_bottom_right, j_bottom_right]
)
pos = torch.cat([i.view(1, -1), j.view(1, -1)], dim=0)
if not return_corners:
return [descriptors, pos, ids]
else:
corners = torch.stack([
torch.stack([i_top_left, j_top_left], dim=0),
torch.stack([i_top_right, j_top_right], dim=0),
torch.stack([i_bottom_left, j_bottom_left], dim=0),
torch.stack([i_bottom_right, j_bottom_right], dim=0)
], dim=0)
return [descriptors, pos, ids, corners]
def savefig(filepath, fig=None, dpi=None):
# TomNorway - https://stackoverflow.com/a/53516034
if not fig:
fig = plt.gcf()
plt.subplots_adjust(0, 0, 1, 1, 0, 0)
for ax in fig.axes:
ax.axis('off')
ax.margins(0, 0)
ax.xaxis.set_major_locator(plt.NullLocator())
ax.yaxis.set_major_locator(plt.NullLocator())
fig.savefig(filepath, pad_inches=0, bbox_inches='tight', dpi=dpi)