本文主要是介绍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实现眼球的虹膜瞳仁分割
主业前端程序猿一枚。图片处理方面,作为业余爱好。如有错误,请各位大佬轻喷,谢谢!!😂
这篇关于pytorch+u2net实现眼球的虹膜瞳仁分割的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!