def test_center_crop_with_incorrectly_large_crop_size():
img = np.ones((4, 4), dtype=np.uint8)
with pytest.raises(ValueError) as exc_info:
A.center_crop(img, 8, 8)
assert str(
exc_info.value
) == "Requested crop size (8, 8) is larger than the image size (4, 4)"
def test_center_crop(target):
img = np.array([[1, 1, 1, 1], [0, 1, 1, 1], [0, 0, 1, 1], [0, 0, 0, 1]],
dtype=np.uint8)
expected = np.array([[1, 1], [0, 1]], dtype=np.uint8)
img, expected = convert_2d_to_target_format([img, expected], target=target)
cropped_img = A.center_crop(img, 2, 2)
assert np.array_equal(cropped_img, expected)
def plot_predictions(images, predictions, paths):
plt.rcParams['figure.figsize'] = [12, 8]
plt.rcParams['figure.dpi'] = 200
_, predictions = torch.max(predictions, 1)
device = predictions.device
canvas = torch.zeros(predictions.shape[:-1],
dtype=torch.long,
device=device)
for pred in predictions.permute(3, 0, 1, 2):
canvas = torch.where(canvas == 0.0, pred, canvas)
# plot_raw_surfaces(images, predictions.unsqueeze(1))
predictions = canvas.cpu().numpy()
for img, pred, path in zip(images, predictions, paths):
rgb = np.array([label_colors[c % 100]
for c in pred]).astype(np.float32) / 255
m = max(img.shape[:-1])
rgb = A.resize(rgb, width=m, height=m, interpolation=cv2.INTER_NEAREST)
rgb = A.center_crop(rgb, *img.shape[:-1])
fig, (ax1, ax2) = plt.subplots(1, 2)
fig.suptitle(path)
ax1.axis('off')
ax1.imshow(img)
ax2.axis('off')
ax2.imshow(rgb)
plt.show()
def test_center_crop_float(target):
img = np.array([[0.4, 0.4, 0.4, 0.4], [0.0, 0.4, 0.4, 0.4],
[0.0, 0.0, 0.4, 0.4], [0.0, 0.0, 0.0, 0.4]],
dtype=np.float32)
expected = np.array([[0.4, 0.4], [0.0, 0.4]], dtype=np.float32)
img, expected = convert_2d_to_target_format([img, expected], target=target)
cropped_img = A.center_crop(img, 2, 2)
assert_array_almost_equal_nulp(cropped_img, expected)
def save_predictions(images, predictions, paths):
plt.rcParams['figure.figsize'] = [12, 8]
plt.rcParams['figure.dpi'] = 200
normal_ps = predictions
normal_ps = (normal_ps.cpu().numpy() + 1) / 2
for img, normal_p, path in zip(images, normal_ps, paths):
normal = normal_p.transpose(1, 2, 0)
m = max(img.shape[:-1])
normal = A.resize(normal, width=m, height=m, interpolation=cv2.INTER_NEAREST)
normal = A.center_crop(normal, *img.shape[:-1])
normal_path = str(Path(path).with_suffix(".exr"))
cv2.imwrite(normal_path, normal)
plt.axis('off')
plt.imshow(normal)
plt.savefig(str(Path(path).with_suffix(".png")))
plt.close()
def plot_predictions(images, predictions, paths):
plt.rcParams['figure.figsize'] = [12, 8]
plt.rcParams['figure.dpi'] = 200
normal_ps = predictions
normal_ps = (normal_ps.cpu().numpy() + 1) / 2
for img, normal_p, path in zip(images, normal_ps, paths):
normal = normal_p.transpose(1, 2, 0)
m = max(img.shape[:-1])
normal = A.resize(normal, width=m, height=m, interpolation=cv2.INTER_NEAREST)
normal = A.center_crop(normal, *img.shape[:-1])
fig, (ax1, ax2) = plt.subplots(1, 2)
fig.suptitle(path)
ax1.axis('off')
ax1.imshow(img)
ax2.axis('off')
ax2.imshow(normal)
plt.show()
def save_predictions(images, predictions, paths):
plt.rcParams['figure.figsize'] = [12, 8]
plt.rcParams['figure.dpi'] = 200
_, predictions = torch.max(predictions, 1)
canvas = layers_to_canvas(predictions)
predictions = canvas.cpu().numpy()
for img, pred, path in zip(images, predictions, paths):
rgb = np.array([label_colors[c % 100]
for c in pred]).astype(np.float32) / 255
m = max(img.shape[:-1])
rgb = A.resize(rgb, width=m, height=m, interpolation=cv2.INTER_NEAREST)
rgb = A.center_crop(rgb, *img.shape[:-1])
pred_path = str(Path(path).with_suffix(".exr"))
cv2.imwrite(pred_path, rgb)
plt.axis('off')
plt.imshow(rgb)
plt.savefig(str(Path(path).with_suffix(".png")))
plt.close()
