import cv2
import torch
import torch.nn.functional as F
import numpy as np
import time
device = torch.device("cuda:0")
def torch_inpaint(pil_img, mask):
start = time.time()
r = mask[:, 0:1, :, :] # r 是手写体
g = mask[:, 1:2, :, :] # g 是印刷体
b = mask[:, 2:3, :, :] # b 是背景
mask = r & torch.bitwise_not(g)
# dialte 手写体 mask
dil_kernel = torch.Tensor([[[[0,1,0], [1,1,1], [0,1,0]]]]).to(device=device)
mask = F.conv2d(mask.type(torch.float32), dil_kernel, padding=1) # x1
mask = F.conv2d(mask.type(torch.float32), dil_kernel, padding=1) # x2
mask = (mask > 0)
mask = mask.repeat(1, 3, 1, 1)
mask_g = g.repeat(1, 3, 1, 1)
mask_b = b.repeat(1, 3, 1, 1)
# TODO:
### mask 膨胀一圈之后减去 mask,这部分和 背景取 & 作为平均,分区域做这步操作
tt = torch.mean(pil_img[mask_b]) # 应该选取周围的背景色而不是所有的背景色的平均
pil_img[mask] = 0
pil_img_tmp = pil_img.clone()
pil_img_tmp[mask_g] = 0
mask_g = mask_g.bool()
box_ks = 15
filters = torch.ones((3, 1, box_ks, box_ks), dtype=torch.float32, device=device)
blurred = F.conv2d(pil_img_tmp, filters, padding=(box_ks-1)//2, groups=3)
countMask = F.conv2d(((mask|mask_g) == 0).type(torch.float32), filters, padding=(box_ks-1)//2, groups=3)
# print('[filters, blurred, countMask] shape===', filters.shape, blurred.shape, countMask.shape)
eps = 1e-8
pil_img[mask] = blurred[mask] / (countMask[mask]+eps)
pil_img[(mask & (pil_img == 0))] = tt
# ori_mask = ori_mask.squeeze().cpu().numpy().transpose(1,2,0)
pil_img = pil_img.squeeze().cpu().numpy().transpose(1,2,0).astype(np.uint8)
print(time.time() - start)
return pil_img
if __name__ == '__main__':
### 最终的结果再加个这个
img = cv2.imread("./331800.jpg")
mask = cv2.imread("./331800.png", 0)
# img = img[:100, 300:500, :]
# mask = mask[:100, 300:500]
mask_one_hot = np.zeros((mask.shape[0], mask.shape[1], 3),dtype=np.uint8)
# 0-背景 1-印刷体 2-手写体
mask_one_hot[:, :, 0][mask == 2] = 1
mask_one_hot[:, :, 1][mask == 1] = 1
mask_one_hot[:, :, 2][mask == 0] = 1
im_tensor = torch.tensor([img], dtype=torch.float32)
mask_tensor = torch.tensor([mask_one_hot])
#
im_tensor = torch.transpose(im_tensor, 1,3)
im_tensor = torch.transpose(im_tensor, 2,3).cuda()
mask_tensor = torch.transpose(mask_tensor, 1,3)
mask_tensor = torch.transpose(mask_tensor, 2,3).cuda()
#
output_im = torch_inpaint(im_tensor, mask_tensor)
cv2.imwrite("output_im.png", output_im)
document image inpaint
最新推荐文章于 2023-09-06 18:18:11 发布