Spaces:
Sleeping
Sleeping
File size: 4,906 Bytes
06f26d7 |
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 |
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2022-06-09 14:59:55
import torch
import random
import numpy as np
from einops import rearrange
def batch_inpainging_from_grad(im_in, mask, gradx, grady):
'''
Recovering from gradient for batch data (torch tensro).
Input:
im_in: N x c x h x w, torch tensor, masked image
mask: N x 1 x h x w, torch tensor
gradx, grady: N x c x h x w, torch tensor, image gradient
'''
im_out = torch.zeros_like(im_in.data)
for ii in range(im_in.shape[0]):
im_current, gradx_current, grady_current = [rearrange(x[ii,].cpu().numpy(), 'c h w -> h w c')
for x in [im_in, gradx, grady]]
mask_current = mask[ii, 0,].cpu().numpy()
out_current = inpainting_from_grad(im_current, mask_current, gradx_current, grady_current)
im_out[ii,] = torch.from_numpy(rearrange(out_current, 'h w c -> c h w')).to(
device=im_in.device,
dtype=im_in.dtype
)
return im_out
def inpainting_from_grad(im_in, mask, gradx, grady):
'''
Input:
im_in: h x w x c, masked image, numpy array
mask: h x w, image mask, 1 represents missing value
gradx: h x w x c, gradient along x-axis, numpy array
grady: h x w x c, gradient along y-axis, numpy array
Output:
im_out: recoverd image
'''
h, w = im_in.shape[:2]
counts_h = np.sum(1-mask, axis=0, keepdims=False)
counts_w = np.sum(1-mask, axis=1, keepdims=False)
if np.any(counts_h[1:-1,] == h):
idx = find_first_index(counts_h[1:-1,], h) + 1
im_out = fill_image_from_gradx(im_in, mask, gradx, idx)
elif np.any(counts_w[1:-1,] == w):
idx = find_first_index(counts_w[1:-1,], w) + 1
im_out = inpainting_from_grad(im_in.T, mask.T, gradx.T, idx)
else:
idx = random.choices(list(range(1,w-1)), k=1, weights=counts_h[1:-1])[0]
line = fill_line(im_in[:, idx, ], mask[:, idx,], grady[:, idx,])
im_in[:, idx,] = line
im_out = fill_image_from_gradx(im_in, mask, gradx, idx)
if im_in.ndim > mask.ndim:
mask = mask[:, :, None]
im_out = im_in + im_out * mask
return im_out
def fill_image_from_gradx(im_in, mask, gradx, idx):
init = np.zeros_like(im_in)
init[:, idx,] = im_in[:, idx,]
right = np.cumsum(init[:, idx:-1, ] + gradx[:, idx+1:, ], axis=1)
left = np.cumsum(
init[:, idx:0:-1, ] - gradx[:, idx:0:-1, ],
axis=1
)[:, ::-1]
center = im_in[:, idx, ][:, None] # h x 1 x 3
im_out = np.concatenate((left, center, right), axis=1)
return im_out
def fill_line(xx, mm, grad):
'''
Fill one line from grad.
Input:
xx: n x c array, masked vector
mm: (n,) array, mask, 1 represent missing value
grad: (n,) array
'''
n = xx.shape[0]
assert mm.sum() < n
if mm.sum() == 0:
return xx
else:
idx1 = find_first_index(mm, 1)
if idx1 == 0:
idx2 = find_first_index(mm, 0)
subx = xx[idx2::-1,].copy()
subgrad = grad[idx2::-1, ].copy()
subx -= subgrad
xx[:idx2,] = np.cumsum(subx, axis=0)[idx2-1::-1,]
mm[idx1:idx2,] = 0
else:
idx2 = find_first_index(mm[idx1:,], 0) + idx1
subx = xx[idx1-1:idx2-1,].copy()
subgrad = grad[idx1:idx2,].copy()
subx += subgrad
xx[idx1:idx2,] = np.cumsum(subx, axis=0)
mm[idx1:idx2,] = 0
return fill_line(xx, mm, grad)
def find_first_index(mm, value):
'''
Input:
mm: (n, ) array
value: scalar
'''
try:
out = next((idx for idx, val in np.ndenumerate(mm) if val == value))[0]
except StopIteration:
out = mm.shape[0]
return out
if __name__ == '__main__':
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).resolve().parents[1]))
from utils import util_image
from datapipe.masks.train import process_mask
# mask_file_names = [x for x in Path('../lama/LaMa_test_images').glob('*mask*.png')]
mask_file_names = [x for x in Path('./testdata/inpainting/val/places/').glob('*mask*.png')]
file_names = [x.parents[0]/(x.stem.rsplit('_mask',1)[0]+'.png') for x in mask_file_names]
for im_path, mask_path in zip(file_names, mask_file_names):
im = util_image.imread(im_path, chn='rgb', dtype='float32')
mask = process_mask(util_image.imread(mask_path, chn='rgb', dtype='float32')[:, :, 0])
grad_dict = util_image.imgrad(im)
im_masked = im * (1 - mask[:, :, None])
im_recover = inpainting_from_grad(im_masked, mask, grad_dict['gradx'], grad_dict['grady'])
error_max = np.abs(im_recover -im).max()
print('Error Max: {:.2e}'.format(error_max))
|