Spaces:
Running
Running
File size: 5,220 Bytes
2c8b554 1ccca05 2c8b554 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
import matplotlib.pyplot as plt
import numpy as np
import torch
from .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)
|