一、效果展示


pytorch 图像分割 数据增强 pytorch分割框架_opencv


pytorch 图像分割 数据增强 pytorch分割框架_pytorch 图像分割 数据增强_02

二、技术摘要

用到的技术如下:

1、模型训练。使用pytorch版的u2net网络

2、裁切眼睛图片,制作数据集。用到dlib68点模型,识别出眼部位置。后面u2net识别的时候,也需要先用dlib裁切,再送入模型。

3、图像操作。使用Pillow,简单方便。

环境的配置,可自行查找资料。

三、数据集制作

使用dlib将人物图片的眼睛位置裁切出来。

pytorch 图像分割 数据增强 pytorch分割框架_图像识别_03


眼睛裁切出来后,需要一张张的把虹膜位置标识出来,这步很繁琐。我这里制作了235张图片的label。

pytorch 图像分割 数据增强 pytorch分割框架_深度学习_04

四、训练

后面会放上本次训练使用的数据集和训练脚本,解压后会得到这样的目录结构

pytorch 图像分割 数据增强 pytorch分割框架_pytorch_05


eyes_data用于裁切眼部图片,制作数据集。

u2net_seg则是训练用脚本。

eyes文件夹保存着本次训练使用的数据集。

运行u2net_train.py即可开始训练。

u2net有两种模型结构,u2net和u2netp。u2net训练出来的模型有173.6M,u2netp模型更小,只有4.7M。两种模型我都尝试训练了,可在saved_models看到。

五、推理

一般场景我们可能都是预测双眼的位置信息,不单独只做单个眼睛虹膜的预测。

为了实现双眼的预测,需要裁切眼部图片时,将坐标信息保存下来。模型预测后,再还原回原图片。

pytorch 图像分割 数据增强 pytorch分割框架_pytorch_06


pytorch 图像分割 数据增强 pytorch分割框架_opencv_07


运行u2net_seg/u2net_eyes_mask.py脚本即可。

六、其他

最后给出资源,包括数据集、u2net和u2net已训练好的权重文件、双眼虹膜预测脚本,配置环境后可直接运行。如果你对结果不满意,也可以使用数据集,对模型继续进行训练。
人物数据集可以到这个网站下载:G-Lab人脸生成实验 资源路径:pytorch+u2net实现眼球的虹膜瞳仁分割

主业前端程序猿一枚。图片处理方面,作为业余爱好。如有错误,请各位大佬轻喷,谢谢!!😂