图像分割 | 基于U2net的抠图项目，手把手教学

前言

兴趣是学习的动力，它激发我们去探索未知的领域，激励我们追求更深层次的知识。
当我们对某个主题或技能产生浓厚的兴趣时，学习便不再是一种负担，而是一种享受。
鲁迅曾经说过:当你对一个项目产生兴趣之后，就会像对刚追到手的妹子一样，爱不释手！！！
所以本文旨在让你对项目激发兴趣，纯动手项目，少理论，多实践，废话少，干货多。

U2net可以实现的图片处理效果：

没错！！上来就上实现效果，感兴趣了没？？？！！要是还不感兴趣的话，就不浪费家人们的时间了，点个赞就可以滚蛋了。

看到这里的人应该还没有放弃学习，功能基本就是一键抠图和生成素描画，下面分享一下抠图实现步骤。

1.下载U2net代码

代码网址：https://github.com/xuebinqin/U-2-Net
进不去的话，可以科学上网一下

1.1 方法1，直接git下载

git clone https://github.com/xuebinqin/U-2-Net.git

1.2 方法2，github上下载ZIP

1.3 方法3，微信转账200快，我手把手教你

2.下载U2net预训练模型

谷歌网盘和百度网盘下载都行，能科学上网的话最好是下谷歌网盘里的模型，不然只能用百度网盘，虽然限速很难受但总比下不下来强

这里有两个模型可选，U2net和U2netp，这是两个不同的模型，可以把U2netp看作U2net的简化版模型，这里我选择下载U2net。

3.训练环境配置

3.1 安装Anaconda

具体可以看:
保姆级conda安装教程:
https://blog.csdn.net/Q_fairy/article/details/129158178

3.2 pytorch安装

访问https://pytorch.org/

选择合适的配置安装，有没有gpu，或者跑cpu，cpu版本有没有英伟达显卡都能用，我这里选择安装cpu版本。

cd U-2-Net
conda create -n u2net python=3.6
conda activate u2net
pip install -r requirements.txt

4.抠图功能实现

4.1 PyCharm运行u2net_test.py

这里我把u2net.pth模型直接放在项目根目录下，并创建了一个test文件夹（0002-01.jpg是从test_data里面随便找的图片），并修改了image_dir、prediction_dir、model_dir,运行之后会生成0002-01.png，这个图片就是0002-01.jpg的mask图。

4.2 使用mask图实现抠图效果

原理：带Alpha通道的图片是RGBA通道原图(0002-01.jpg)是RGB图片，讲输出图(0002-01.png)当作Alpha通道，与原图组成RGBA图片(抠图完成的图片)
抠图代码:

import cv2

# 读取RGB图像和mask图像
image = cv2.imread('0002-01.jpg')  # 替换为你的RGB图片路径
mask = cv2.imread('0002-01.png', cv2.IMREAD_GRAYSCALE)  # 读取mask图，确保是灰度模式

# 确保mask是二值化的（0和255）
_, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)

# 使用mask抠图
# 将原图与mask相乘，得到抠出的图像
result = cv2.bitwise_and(image, image, mask=binary_mask)

# 创建一个新的图像，包含透明通道（4通道）
b, g, r = cv2.split(result)  # 分离抠出的颜色通道
alpha_channel = binary_mask  # 使用mask作为alpha通道

# 将各个通道合并到一个新的4通道图像中
rgba_image = cv2.merge((b, g, r, alpha_channel))

# 保存结果图像为PNG格式以保留透明度
cv2.imwrite('result_with_transparency.png', rgba_image)

# 显示结果
cv2.imshow('Original Image', image)
cv2.imshow('Mask', binary_mask)
cv2.imshow('Result with Transparency', rgba_image)

# 等待按键并关闭窗口
cv2.waitKey(0)
cv2.destroyAllWindows()

效果:

5.人像素描功能实现

5.1 下载u2net_portrait.pth预训练模型

5.2 通过预训练模型生成人像素描图

如图所示，修改一下模型的路径就可以运行了，img_dir是输入图片路径，prediction_dir是输出图片路径。
注:如果出现Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. 错误，直接按照它的提示修改加载模型代码就行，将net.load_state_dict(torch.load(model_dir))改成net.load_state_dict(torch.load(model_dir,map_location=“cpu”))

5.3 生成人像素描图展示

总结

本篇文章简单介绍了一下U2net模型的用法，基于预训练模型来做图片处理，要想训练高精准度模型还需要自己找数据集，训练自己的数据集。