持续创作,加速成长!这是我参与「掘金日新计划 · 10 月更文挑战」的第4天,github.com/lizhengwei1…
数据集:pan.baidu.com/share/init?… 密码是:dzsn
数据集介绍
北京玩星汇聚科技有限公司所高质量标注并开源的一份人像抠图(matting) 数据集,是目前已知最大的人像matting数据集,包含 34427 张图像和对应的matting结果图. 数据集中的原始图片来源于Flickr、百度、淘宝. 基于该数据集所训练的人像软分割模型已商用。
数据集下载
数据集分享在百度网盘(国内):
提取码:dzsn
MEGA(国外):
数据集文件为:matting_human_half.zip
,解压后其目录结构为:
- matting_human_half/
- clip_img #人像图片(半身),.jpg 格式
- matting #人像标注,.png 格式
clip_img文件夹:原图;
matting文件夹:原图对应的mask图,四个维度,第四个维度是图片对应的mask;
该数据集中,图片是经过人脸检测和区域裁剪后生成了600×800的半身人像. 标注的人像 matting 图片为 png 格式,可以从 png 图片中提取人像的 mask图,代码如下:
#!/usr/bin/python3
#!--*-- coding:utf-8 --*--
import cv2
import matplotlib.pyplot as plt
png = cv2.imread('./data/matting/1803151818/matting_00000000/1803151818-00000006.png', cv2.IMREAD_UNCHANGED)
print(png.shape)
#(800, 600, 4)
png_img = cv2.cvtColor(png[:,:,:3], cv2.COLOR_BGR2RGB)
alpha = png[:,:,3]
plt.figure(figsize=(10, 8))
plt.subplot(1, 2, 1)
plt.imshow(png_img)
plt.title("Matting PNG img")
plt.axis("off")
plt.subplot(1, 2, 2)
plt.imshow(alpha)
plt.title("Matting Alpha img")
plt.axis("off")
plt.show()
运行结果:
计算alpha图
下载代码,并解压,将clip_img文件夹和matting文件夹。如下图:
生成train.txt
在data文件夹下,新建get_train_txt.py生成train.txt
import os
pic_path = "matting/"
with open("train.txt", "w", encoding="UTF-8") as ff:
for name_0 in os.listdir(pic_path):
for name_1 in os.listdir(pic_path + "/" + name_0):
for name_2 in os.listdir(pic_path + "/" + name_0 + "/" + name_1):
pic_input_path = name_0 + "/" + name_1 + "/" + name_2
ff.write(pic_input_path + "n")
ff.close()
print("well done____________!")
遍历matting,写入train.txt,运行结果:
1803151818/matting_00000000/1803151818-00000003.png
1803151818/matting_00000000/1803151818-00000004.png
1803151818/matting_00000000/1803151818-00000006.png
1803151818/matting_00000000/1803151818-00000007.png
1803151818/matting_00000000/1803151818-00000008.png
1803151818/matting_00000000/1803151818-00000009.png
1803151818/matting_00000000/1803151818-00000010.png
1803151818/matting_00000000/1803151818-00000012.png
1803151818/matting_00000000/1803151818-00000015.png
1803151818/matting_00000000/1803151818-00000017.png
1803151818/matting_00000000/1803151818-00000020.png
1803151818/matting_00000000/1803151818-00000021.png
生成mask
在data新建get_mask.py生成mask。
import os
import cv2
matting_path = "matting/"
mask_path = "mask/"
for name_0 in os.listdir(matting_path):
if not os.path.exists(mask_path + "/" + name_0):
os.makedirs(mask_path + "/" + name_0)
for name_1 in os.listdir(matting_path + "/" + name_0):
if not os.path.exists(mask_path + name_0 + "/" + name_1):
os.mkdir(mask_path + name_0 + "/" + name_1)
for name_2 in os.listdir(matting_path + "/" + name_0 + "/" + name_1):
pic_input_path = matting_path + "/" + name_0 + "/" + name_1 + "/" + name_2
pic_output_path = mask_path + "/" + name_0 + "/" + name_1 + "/" + name_2
print("pic_input_path=", pic_input_path)
in_image = cv2.imread(pic_input_path, cv2.IMREAD_UNCHANGED)
alpha = in_image[:, :, 3]
cv2.imwrite(pic_output_path, alpha)
运行结果:
生成Trimap图
源代码有这个脚本文件gen_trimap.py,需要做一些修改。
首先对参数做修改:
def get_args():
parser = argparse.ArgumentParser(description='Trimap')
parser.add_argument('--mskDir', type=str, default='mask', help="masks directory")
parser.add_argument('--saveDir', type=str, default='trimap', help="where trimap result save to")
parser.add_argument('--list', type=str, default='train.txt', help="list of images id")
parser.add_argument('--size', type=int, default=10, help="kernel size")
args = parser.parse_args()
print(args)
return args
注释掉gen_trimap.py第37/43/49行的断言语句
# assert(cnt1 == cnt2 + cnt3)
导入os包
import os
在67行,插入:
trimap_name_1 = trimap_name.split("/")[:-1]
trimap_path = "/".join(trimap_name_1)
if not os.path.exists(trimap_path):
os.makedirs(trimap_path)
如果trimap路径不存在则创建路径。
修改完成后,运行gen_trimap.py,生成trimap图。运行结果:
单张图:
生成alpha图
修改./data/knn_matting.py脚本,主要修改该main方法。
def main():
with open('train.txt') as fs:
list_path = fs.read().splitlines()
for image_path in list_path:
img_name = os.path.join(rgbDir, image_path)
trimap_name = os.path.join(trimapDir, image_path[:-4] + '.png')
alpha_name = os.path.join(alphaDir, image_path[:-4] + '.png')
img_name = img_name.replace('\', '/')
trimap_name = trimap_name.replace('\', '/')
alpha_name = alpha_name.replace('\', '/')
index_path = alpha_name.rindex('/')
alpha_dir = alpha_name[:index_path]
if not os.path.exists(alpha_dir):
os.makedirs(alpha_dir)
img_name = './' + img_name
img = cv2.imread(img_name)
trimap = cv2.imread(trimap_name)
alpha = knn_matte(img, trimap)
cv2.imwrite(alpha_name, alpha * 255)
运行结果:
可以看出,我们己经得到比较精细的Alpha图了。不过计算过程比较慢,几分钟一张图。如果时间太长可以直接使用mask图。
训练
修改dataset.py文件
将dataset.py文件第17/18/19行修改为:
image_name = os.path.join(data_dir, 'clip_img', file_name['image'].replace("matting", "clip").replace("png", "jpg"))
trimap_name = os.path.join(data_dir, 'trimap', file_name['trimap'].replace("clip", "matting"))
alpha_name = os.path.join(data_dir, 'alpha', file_name['alpha'].replace("clip", "matting"))
./data/clip_img/1803201916/clip_00000000/1803201916-00000117.png,这张图片的格式有问题,需要修改为 .jpg
格式
将dataset.py文件第101/102/103行,修改为
trimap[trimap == 0] = 0
trimap[trimap >= 250] = 2
trimap[np.where(~((trimap == 0) | (trimap == 2)))] = 1
原因:trimap图是三色图,但是它的“三色”并不像上图中0/128/255只有这三色,它是在[0, 255]这个区间范围内。所以新改的代码,将这“三色”用区间区分,作为三种不同的label传入训练。
训练T-Net
python train.py --dataDir='./data' --saveDir='./ckpt' --trainData='human_matting_data' --trainList='./data/train.txt' --lrdecayType='keep' --nEpochs=200 --save_epoch=1 --load='human_matting' --patch_size=320 --lr=1e-5 --nThreads=4 --train_batch=4 --train_phase='pre_train_t_net'
训练M-Net:
python train.py --dataDir='./data' --saveDir='./ckpt' --trainData='human_matting_data' --trainList='./data/train.txt' --lrdecayType='keep' --nEpochs=400 --save_epoch=1 --load='human_matting' --patch_size=320 --lr=5e-6 --gpus='0' --nThreads=4 --train_batch=4 --train_phase='end_to_end'
测试
打开test_camera.py,修改如下参数:
parser = argparse.ArgumentParser(description='human matting')
parser.add_argument('--model', default='./ckpt/human_matting/model/model_obj.pth', help='preTrained model')
parser.add_argument('--size', type=int, default=320,help='input size')
parser.add_argument('--without_gpu', action='store_true', default=False, help='no use gpu')
然后,运行test_camera.py就可以得到运行结果,由于目前没有训练好,只是调通了算法
等调通算法后再展示详细的结果。
代码
download.csdn.net/download/hh…
参考文章:
【SHM】Semantic Human Matting抠图算法调试_zzZ_CMing的博客-CSDN博客_human matting