pytorch+u2net实现眼球的虹膜瞳仁分割

一、效果展示

（图片来源网络，如有侵权请联系删除）

二、技术摘要

用到的技术如下：

1、模型训练。使用pytorch版的u2net网络（https://gitee.com/throni0/U-2-Net）。

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

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

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

三、数据集制作

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

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

四、训练

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

eyes_data用于裁切眼部图片，制作数据集。
u2net_seg则是训练用脚本。
eyes文件夹保存着本次训练使用的数据集。
运行u2net_train.py即可开始训练。
u2net有两种模型结构，u2net和u2netp。u2net训练出来的模型有173.6M，u2netp模型更小，只有4.7M。两种模型我都尝试训练了，可在saved_models看到。

五、推理

一般场景我们可能都是预测双眼的位置信息，不单独只做单个眼睛虹膜的预测。
为了实现双眼的预测，需要裁切眼部图片时，将坐标信息保存下来。模型预测后，再还原回原图片。


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

六、其他

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

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