基于U2-Net如何训练一个一键抠图模型

1. 前言

抠图是图像编辑的基础功能之一,在抠图的基础上可以发展出很多有意思的玩法和特效。比如一键更换背景、一键任务卡通化、一键人物素描化等。正是因为这些有意思的玩法。

最近也是对此模型背后的U^2-Net网络很感兴趣,收集数据训练了人脸素描化模型,尽管受限于数据集,只能在人脸图片上转换成功,但自己仍然玩的不亦乐乎。不仅乐于玩模型的有意思的效果,更乐在训练模型过程中,以及遇到问题解决问题过程中,对模型理解的不断加深。

最近对一键扣图模型从头训练了一遍,并在训练过程中持续测试了不同阶段模型的表现,看着模型一点点的收敛,抠图效果慢慢变好。

此处记录下训练过程以及训练的效果。也可以对后来者有一个参考。

提前说一声,模型训练很耗时!

2. 代码 & 数据 & 环境准备

2.1 代码

代码是U-2-Net的开源代码,可以从Github下载:https://github.com/NathanUA/U-2-Net。这个模型本来是做显著性检测的,但是当成一键扣图模型也很好玩。

需要注意的地方是,如果是安装的最新的Pytorch,获取loss值的时候,需要将loss.data[0] 修改为loss.data.item()

笔者在训练过程中曾尝试修改Loss函数为其他的,比如改成BCESSIM的加权(参考U-2-Net作者的文章BASNet),未见明显提升。也曾修改输出通道训练其他模型,暂无好玩的结果,就当是积累经验了。

2.2 数据

数据集我们就用论文中提到的DUTS数据集,已经分好了训练集和测试集。网上搜一下直接下载即可。

当然,也可以用自己的数据集,按照DUTS的格式重新组织下数据集即可。

然后在训练代码里面把数据读取部分的路径更换为自己准备的数据的路径。

2.3 机器

然后基于Anaconda安装训练所需的Python环境,创建虚拟环境,安装pytorch, torchvision, skimage, opencv等等,直接pip install或者conda install即可。不多说。

另外多卡的话,代码还需要有一些细微的改动,在构建模型之后,将代码:

    if torch.cuda.is_available():net.cuda()

修改为

    if torch.cuda.is_available():net.cuda()net = nn.DataParallel(net)

3. 训练与测试

3.1 模型训练

以上代码、数据、机器和运行环境都已经准备好之后,就可以开始训练了。多卡训练的命令大概长下面这样:

CUDA_VISIBLE_DEVICES=0,1,2,3 nohup python3 -u u2net_train.py > log_train.log &

然后tail命令查看日志文件log_train.log,如果看到下面这样的输出,说明跑起来了:

再用命令watch -n 1 nvidia-smi查看GPU的情况,可以看到四张卡都被充分利用起来了。

模型训练将近一周,达到了接近论文的效果。

另外,由于中间保存过多,为了节省空间,笔者删掉了太多前期模型,以下展示的前期效果是另外一次训练的前期模型的效果。

3.2 各阶段模型测试

笔者微调测试代码结构,把测试转移到了Jupyter里,这样画图看效果更加直观。

笔者测试模型的时候,每张图都会画出三个图:黑色背景的抠图结果、模型输出的Mask或称Alpha,原图。这样对比来看结果一目了然。这里每张图都展示了四个阶段模型的测试效果。显然,以下图片都不在训练集里面。

四个阶段对比着看,能更加直观地感受到模型的收敛过程。

从以下四个阶段的对比图可以看出,随着训练的进行

  • 前景逐渐变亮,背景逐渐变暗,即前景收敛于1,背景收敛于0。前两幅图之间的对比最为明显。

  • 前景的轮廓从模糊到清晰细锐,轮廓处的不确定区域,越来越少。

  • 注意指缝和发梢部分的Mask的变化,细节越来越清晰。

下面这幅图请注意这个卡通人物背后背的那个是蜗牛还是啥的东西的轮廓的细化过程。以及其嘴角的一撮小胡子。这个图美中不足的是两脚之间的背景没有被识别出来。

下面这张图值得关注的应该就是其发梢的抠图细化过程、腰部的亮度变化过程。还有就是其手中的衣服了,对于要不要把一副也给抠出来,模型看起来也很纠结啊。

这个图最引人瞩目的莫过于这位美女在风中凌乱的发丝,这不是难为模型吗?说实话,如果不是看到Mask里胸前多出的东西,我都没注意到这个东西,衣服的胸结还是啥。

这大概就是训练了五天左右的效果,模型仍然在缓慢的收敛中,故事仍然在继续......

直到我实在是受不了越来越慢的收敛速度,等不及训练其他魔改的模型,终止了训练任务......

本着报喜不报忧的原则,下面再放几张测试效果还不错的图片,效果不怎么样的就不拿出来献丑了

上面的抠图效果还是有待提高,比如头发等边缘处,还是可见部分背景未分离。前几天刚转发了动物抠图的新论文,边缘和毛发的抠图效果很赞。其单开一条支路专门做轮廓边缘处的抠图的思路值得参考。

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.hqwc.cn/news/228459.html

如若内容造成侵权/违法违规/事实不符,请联系编程知识网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

MySQL基础进阶篇

进阶篇 存储引擎 MySQL体系结构: 存储引擎就是存储数据、建立索引、更新/查询数据等技术的实现方式。存储引擎是基于表而不是基于库的,所以存储引擎也可以被称为表引擎。 默认存储引擎是InnoDB。 相关操作: -- 查询建表语句 show create …

鸿蒙4.0开发笔记之ArkTS语法基础的UI描述、基础组件的使用与如何查看组件是否有参数(八)

文章目录 一、声明式UI描述1、无/有参数组件2、如何查看组件是否有参数 二、Image组件的使用三、组件的属性设置四、补充1、使用组件的成员函数配置组件的事件方法2、配置子组件3、多组件嵌套 一、声明式UI描述 在HarmonyOS的ArkTS语法中,万物皆组件。ArkTS以声明方…

多个nginx共享值、缓存问题

背景 目前我在集成登录认证功能(cas),使用的架构是nginxlua,由于我们有多个系统(全是前端项目),每套系统都采用nginxlua的方式进行部署(即每个系统都是一个nginx)&#…

【2023.11.28】关于Servlet路径的学习

创建Servlet 这是Tomcat配置的初始路径,在web项目内,该路径代表了webapp下index.html所在的页面。 WebServlet(name "login", value "/login",loadOnStartup 1) public class LoginServlet extends HttpServlet { 使用注解的方…

Spring Cloud 版本升级记:OpenFeignClient与Gateway的爱恨交织

Spring Cloud 版本升级记:OpenFeignClient与Gateway的爱恨交织 近日,在负责的项目中,我对 Spring Boot、Spring Cloud 以及 Spring Cloud Alibaba 进行了版本升级。原以为会一切顺利,没想到却遭遇了 Spring Cloud Gateway 无法正…

Apache Airflow (十四) :Airflow分布式集群搭建及测试

🏡 个人主页:IT贫道_大数据OLAP体系技术栈,Apache Doris,Clickhouse 技术-CSDN博客 🚩 私聊博主:加入大数据技术讨论群聊,获取更多大数据资料。 🔔 博主个人B栈地址:豹哥教你大数据的个人空间-豹…

LeetCode算法题解(动态规划)|LeetCode198. 打家劫舍、LeetCode213. 打家劫舍 II、LeetCode337. 打家劫舍 III

一、LeetCode198. 打家劫舍 题目链接:198. 打家劫舍 题目描述: 你是一个专业的小偷,计划偷窃沿街的房屋。每间房内都藏有一定的现金,影响你偷窃的唯一制约因素就是相邻的房屋装有相互连通的防盗系统,如果两间相邻的…

java多线程-扩展知识三:乐观锁与悲观锁

1、悲观锁 悲观锁有点像是一位比较悲观(也可以说是未雨绸缪)的人,总是会假设最坏的情况,避免出现问题。 悲观锁总是假设最坏的情况,认为共享资源每次被访问的时候就会出现问题(比如共享数据被修改),所以每次…

Ansys Lumerical|带 1D-2D 光栅的出瞳扩展器

附件下载 联系工作人员获取附件 此示例显示了设置和模拟出瞳扩展器 (EPE) 的工作流程,EPE 是波导型增强现实 (AR) 设备的重要组成部分。该工作流程将利用 Lumerical 和 Zemax OpticStudio 之间的动态链接功能 。为了…

使用 watch+$nextTick 解决Vue引入组件无法使用问题

问题描述: 很多时候我们都需要使用第三方组件库,比如Element-UI,Swiper 等等。 如果我们想要在这些结构中传入自己从服务器请求中获取的数据就会出现无法显示的问题。 比如我们在下面的Swiper例子中,我们需要new Swiper 才能让…

让@RefreshScope注解来帮助我们实现动态刷新

文章目录 前言举例作用参考文章总结 前言 在实际开发当中我们常常会看到有些类上会加一个注解:RefreshScope,有没有对应的小伙伴去思考过这个东西,这个注解有什么作用?为什么要加?下面我们就来看看这个 RefreshScope …

前端入门(四)Ajax、Promise异步、Axios通信、vue-router路由

文章目录 AjaxAjax特点 Promise 异步编程(缺)Promise基本使用状态 - PromiseState结果 - PromiseResult Axios基本使用 Vue路由 - vue-router单页面Web应用(single page web application,SPA)vue-router基本使用路由使…