基于mask图+膨胀侵蚀的trimap图生成方式
前言：
trimap图在AI抠像中的用途是为了得到精准的alpha图，以便后续的合成。
trimap原意是指“三色图”，三色图的意思如下：
- 确定需要的前景区域位置——下右图的白色区域；
- 确定不需要的背景区域位置——下右图的黑色区域；
- 介于需要与不需要的待分割区域位置——下右图的灰色区域；
trimap图大多都是由人工处理得到的，而标记的过程耗时耗力。这里介绍一种基于mask图生成trimap图的方法，时间效率要比手动处理快，但是效果表现有待提高。
–-----------------------------------------------------------------------------—--------------------------------------------
–-----------------------------------------------------------------------------—--------------------------------------------
一、原图得到mask图
mask图在Matting领域中是比较常见的，是用来标记分割物预测区域的图。一般支持图像分割的算法最后的输出都有mask图，所以不论是Fast-RCNN、DeepLab、YOLO都能满足这一步的需求。这里为图简便，代码用的是DeepLab，模型文件也可在网上自行下载。
import os
import tarfile
import numpy as np
from PIL import Image
import cv2, argparse
import tensorflow as tf
class DeepLabModel(object):
"""Class to load deeplab model and run inference."""
INPUT_TENSOR_NAME = 'ImageTensor:0'
OUTPUT_TENSOR_NAME = 'SemanticPredictions:0'
INPUT_SIZE = 513
FROZEN_GRAPH_NAME = 'frozen_inference_graph'
def __init__(self, tarball_path):
self.graph = tf.Graph()
graph_def = None
pb_path = 'Loadding_model/frozen_inference_graph.pb'
graph_def = tf.GraphDef.FromString(open(pb_path, 'rb').read())
if graph_def is None:
raise RuntimeError('Cannot find inference graph in tar archive.')
with self.graph.as_default():
tf.import_graph_def(graph_def, name='')
self.sess = tf.Session(graph=self.graph)
def run(self, image):
"""Runs inference on a single image.
Args:
image: A PIL.Image object, raw input image.
Returns:
resized_image: RGB image resized from original input image.
seg_map: Segmentation map of `resized_image`.
"""
width, height = image.size
resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height)
target_size = (int(resize_ratio * width), int(resize_ratio * height))
resized_image = image.convert('RGB').resize(target_size, Image.ANTIALIAS)
batch_seg_map = self.sess.run(
self.OUTPUT_TENSOR_NAME,
feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]})
seg_map = batch_seg_map[0]
return resized_image, seg_map
def create_pascal_label_colormap():
"""Creates a label colormap used in PASCAL VOC segmentation benchmark.
Returns:
A Colormap for visualizing segmentation results.
"""
colormap = np.zeros((256, 3), dtype=int)
ind = np.arange(256, dtype=int)
for shift in reversed(range(8)):
for channel in range(3):
colormap[:, channel] |= ((ind >> channel) & 1) << shift
ind >>= 3
return colormap
def label_to_color_image(label):
"""Adds color defined by the dataset colormap to the label.
Args:
label: A 2D array with integer type, storing the segmentation label.
Returns:
result: A 2D array with floating type. The element of the array
is the color indexed by the corresponding element in the input label
to the PASCAL color map.
Raises:
ValueError: If label is not of rank 2 or its value is larger than color
map maximum entry.
"""
if label.ndim != 2:
raise ValueError('Expect 2-D input label')
colormap = create_pascal_label_colormap()
if np.max(label) >= len(colormap):
raise ValueError('label value too large.')
return colormap[label]
# # [1]:设置模型
LABEL_NAMES = np.asarray([
'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv'])
FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
# MODEL_NAME = 'xception_coco_voctrainval' # @param ['mobilenetv2_coco_voctrainaug', 'mobilenetv2_coco_voctrainval', 'xception_coco_voctrainaug', 'xception_coco_voctrainval']
# _DOWNLOAD_URL_PREFIX = 'http://download.tensorflow.org/models/'
# _MODEL_URLS = {
# 'mobilenetv2_coco_voctrainaug':
# 'deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz',
# 'mobilenetv2_coco_voctrainval':
# 'deeplabv3_mnv2_pascal_trainval_2018_01_29.tar.gz',
# 'xception_coco_voctrainaug':
# 'deeplabv3_pascal_train_aug_2018_01_04.tar.gz',
# 'xception_coco_voctrainval':
# 'deeplabv3_pascal_trainval_2018_01_04.tar.gz',}
# _TARBALL_NAME = 'deeplab_model.tar.gz'
#
# model_path = os.getcwd() + '/DeepLab_v3_model/'
# download_path = os.path.join(model_path, _TARBALL_NAME)
# if not os.path.exists(download_path):
# print('downloading model, this might take a while...')
# urllib.request.urlretrieve(_DOWNLOAD_URL_PREFIX + _MODEL_URLS[MODEL_NAME],
# download_path)
# print('download completed! loading DeepLab model...')
# # [2]:加载模型
import time
time_1 = time.time()
download_path = 'Loadding_model/frozen_inference_graph.pb'
MODEL = DeepLabModel(download_path)
print('model loaded successfully________! cost_time=', time.time() - time_1)
# # [3]:预测输出
pic_path = "input_pic/"
mask_path = "output_mask/"
if not os.path.exists(mask_path):
os.mkdir(mask_path)
for name_ in os.listdir(pic_path):
pic_data = Image.open(pic_path + name_)
res_im, seg = MODEL.run(pic_data)
seg = cv2.resize(seg.astype(np.uint8), pic_data.size)
mask_sel = (seg==15).astype(np.float32)
cv2.imwrite(mask_path + name_, (255*mask_sel).astype(np.uint8))
print('\nDone: ' + mask_path + name_)
DeepLab可预测的label包含以下20类：
LABEL_NAMES = np.asarray([
'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv'])
博主就以label=person来说明（也就是代码第149行，可以更改seg==n来重新指定预测的类别）；
下面三张图分别是：原图、deeplab得到的mask图、转成黑白的mask图。
–-----------------------------------------------------------------------------—--------------------------------------------
–-----------------------------------------------------------------------------—--------------------------------------------
二、mask图经过膨胀侵蚀得到trimap图
膨胀侵蚀的操作在OpenCV里面是比较常见的，这里就不赘述了，直接上代码：
import os
import cv2
import numpy as np
def dilate_and_erode(mask_data, struc="ELLIPSE", size=(10, 10)):
"""
膨胀侵蚀作用，得到粗略的trimap图
:param mask_data: 读取的mask图数据
:param struc: 结构方式
:param size: 核大小
:return:
"""
if struc == "RECT":
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, size)
elif struc == "CORSS":
kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, size)
else:
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, size)
msk = mask_data / 255
dilated = cv2.dilate(msk, kernel, iterations=1) * 255
eroded = cv2.erode(msk, kernel, iterations=1) * 255
res = dilated.copy()
res[((dilated == 255) & (eroded == 0))] = 128
return res
trimap_path = "data_trimap/"
mask_path = "mask.png"
size = 10
if not os.path.exists(trimap_path):
os.mkdir(trimap_path)
mask_data = cv2.imread(mask_path, 0)
trimap = dilate_and_erode(mask_data, size=(size, size))
cv2.imwrite(trimap_path + mask_path, trimap)
得到结果图如下右图：
虽然右图只比左图在边缘位置多加了一层，但如果trimap图经过传统分割算法（例如贝叶斯、KNN）处理后，边缘处的立体感就会立马感受出来巨大变化。
–-----------------------------------------------------------------------------—--------------------------------------------
–-----------------------------------------------------------------------------—--------------------------------------------
说明：
因为trimap图的用途是为了得到精准的alpha图，下面的篇幅就是拓展内容了，看一下trimap图如何得到alpha图，以及alpha图如何融合到一张背景图，生成一张不存在的假图。
–-----------------------------------------------------------------------------—--------------------------------------------
–-----------------------------------------------------------------------------—--------------------------------------------
三、用KNN算法，从trimap图得到alpha图
从trimap图得到alpha图的方式也有很多，传统的分类算法，新兴的神经网络都能完成，github上也有例子。这里就是找了一篇knn-matting的实现过程，可从trimap图得到alpha图。
'''
borrowed from https://github.com/MarcoForte/knn-matting
'''
import numpy as np
import sklearn.neighbors
import scipy.sparse
import warnings
import cv2
import os
import argparse
import pdb
import time
import matplotlib.pyplot as plt
import scipy.misc
def knn_matte(img, trimap, mylambda=100):
"""
:param img: 原图
:param trimap: trimap图
:param mylambda:
:return:
"""
[m, n, c] = img.shape
img, trimap = img / 255.0, trimap / 255.0
foreground = (trimap > 0.99).astype(int)
background = (trimap < 0.01).astype(int)
all_constraints = foreground + background
# 0.5s
print('Finding nearest neighbors')
a, b = np.unravel_index(np.arange(m * n), (m, n))
feature_vec = np.append(np.transpose(img.reshape(m * n, c)), [a, b] / np.sqrt(m * m + n * n), axis=0).T
nbrs = sklearn.neighbors.NearestNeighbors(n_neighbors=10, n_jobs=4).fit(feature_vec)
knns = nbrs.kneighbors(feature_vec)[1]
# 0.16s
# Compute Sparse A
print('Computing sparse A')
row_inds = np.repeat(np.arange(m * n), 10)
col_inds = knns.reshape(m * n * 10)
vals = 1 - np.linalg.norm(feature_vec[row_inds] - feature_vec[col_inds], axis=1) / (c + 2)
A = scipy.sparse.coo_matrix((vals, (row_inds, col_inds)), shape=(m * n, m * n))
# 0.06s
D_script = scipy.sparse.diags(np.ravel(A.sum(axis=1)))
L = D_script - A
D = scipy.sparse.diags(np.ravel(all_constraints[:, :, 0]))
v = np.ravel(foreground[:, :, 0])
c = 2 * mylambda * np.transpose(v)
H = 2 * (L + mylambda * D)
# 0.9s
print('Solving linear system for alpha')
time_2 = time.time()
warnings.filterwarnings('error')
alpha = []
try:
alpha = np.minimum(np.maximum(scipy.sparse.linalg.spsolve(H, c), 0), 1).reshape(m, n)
except Warning:
x = scipy.sparse.linalg.lsqr(H, c)
alpha = np.minimum(np.maximum(x[0], 0), 1).reshape(m, n)
print("time_1=", time.time() - time_2)
return alpha
def main():
time_1 = time.time()
img_name = "pic.png"
trimap_name = "data_trimap/mask.png"
alpha_name = "alpha.png"
img = cv2.imread(img_name)
trimap = cv2.imread(trimap_name)
alpha = knn_matte(img, trimap)
cv2.imwrite(alpha_name, alpha * 255)
print("time_all=", time.time() - time_1)
if __name__ == '__main__':
main()
结果如下：依次是deeplab_mask图、trimap图、alpha图
可以看出alpha图比mask图拥有更多的细节。
–-----------------------------------------------------------------------------—--------------------------------------------
–-----------------------------------------------------------------------------—--------------------------------------------
四、alpha图融合背景图：
最后一步呢，就是依照得到的较精准的alpha图，将原图人物融合到一张背景图中去。
import numpy as np
import sklearn.neighbors
import scipy.sparse
import warnings
import cv2
import os
import time
def get_result(pic_data, alpha_data, bg_data, center):
"""
用alpha图，贴合原图pic与背景图BG，生成贴合后的图
:param pic_data:原图信息
:param alpha_data:alpha图信息
:param bg_data:背景图信息
:param center:贴合的中心点
:return:
"""
h, w, _ = alpha_data.shape
pic_data = cv2.resize(pic_data, (w, h))
bg_data_now = bg_data.copy()
used_index = np.where(alpha_data > 10)
for n in range(len(used_index[0])):
j = used_index[1][n]
i = used_index[0][n]
bg_data_now[i + center[1]][j + center[0]] = pic_data[i][j]
return bg_data_now
pic_data = cv2.imread("pic.png")
alpha_data = cv2.imread("alpha.png")
bg_data = cv2.imread("bg_2.png")
center = (300, 45)
bg_data_now = get_result(pic_data, alpha_data, bg_data, center)
cv2.imwrite("out.png", bg_data_now)
原图：
背景图：
最后的融合图：
总结：
可以看到最后的融合图在手臂内侧处出现了少许的不完美，这是由于knn-matting预测的alpha图不精确，如果选用深度学习算法预测出一张完美的alpha图，那么在人物抠像中就有很大的发挥空间了。