BEN2:一键快速抠图!自动移除图像和视频中的背景,支持在线使用

673 阅读7分钟

❤️ 如果你也关注 AI 的发展现状,且对 AI 应用开发感兴趣,我会每日分享大模型与 AI 领域的开源项目和应用,提供运行实例和实用教程,帮助你快速上手AI技术!

🥦 微信公众号|搜一搜:蚝油菜花 🥦


大家好,我是蚝油菜花,今天跟大家分享一下 BEN2 这个专门用于从图像和视频中快速移除背景并提取前景的深度学习模型。

🚀 快速阅读

BEN2 是一款专注于从图像和视频中快速移除背景并提取前景的深度学习模型。

  1. 核心功能:自动从图像和视频中移除背景,生成高质量的前景图像,特别适用于复杂背景和细节丰富的图像。
  2. 技术原理:基于置信度引导抠图(CGM)管道,结合精细化网络处理复杂区域,如头发和边缘。

BEN2 是什么

BEN2

BEN2(Background Erase Network 2)是由 Prama LLC 开发的深度学习模型,专为从图像和视频中快速移除背景并提取前景而设计。它基于创新的置信度引导抠图(CGM)管道,通过精细化网络处理复杂区域,如头发和边缘,实现高精度的前景分割。

BEN2 在大规模数据集上训练,擅长处理高分辨率图像(如4K),并且支持 GPU 加速,确保处理速度极快。1080p 图像仅需 6 秒,4K 图像约 20 秒,适用于需要高效处理多个图像的场景。

BEN2 的主要功能

  • 背景移除与前景分割:BEN2 能自动从图像和视频中移除背景,生成高质量的前景图像。通过创新的置信度引导抠图(CGM)管道,对复杂背景和细节丰富的图像(如头发、边缘等)进行高精度分割。
  • 高分辨率处理:支持 4K 图像处理,确保高分辨率图像的分割效果,适用于需要高精度和高质量图像处理的应用场景。
  • 边缘精细化:通过精细化网络处理边缘,提升分割精度,适用于需要精细边缘处理的图像,如产品照片、肖像等。
  • 视频分割:BEN2 可以对视频中的每一帧进行前景提取,适用于动态视频编辑。
  • 简单 API 与易用性:提供简洁的 API,方便集成到各种应用中。
  • 批量图像处理:适合需要高效处理多个图像的场景。

BEN2 的技术原理

  • 置信度引导抠图(CGM):基础模型首先对图像进行初步分割,生成一个初步的前景掩码。置信度映射(confidence map)用于评估每个像素的分割置信度。对于置信度较低的像素,精细化网络会进一步处理,以提高分割精度。
  • 精细化网络:精细化网络专注于处理复杂的边缘区域,例如头发丝和透明物体的边缘。通过多尺度特征融合和上下文信息,精细化网络能生成更自然、更精确的边缘。
  • 深度学习架构:BEN2 在 DIS5k 和 Prama LLC 的 22K 专有分割数据集上进行了训练,能适应多种复杂的背景和光照条件。模型采用高效的编码器-解码器架构,结合多尺度特征提取,确保在高分辨率图像(如 4K)上的高效处理。
  • GPU 加速与 CUDA 支持:BEN2 支持 CUDA 和 GPU 加速,能提升处理速度。部署的实例可以在 6 秒内完成 1080p 图像的抠图,4K 图像处理时间约为 20 秒。

如何运行 BEN2

BEN2 是一个用于前景分割的创新模型,基于 Confidence Guided Matting (CGM) 管道,能够通过精炼网络处理低置信度像素,从而生成更精确的抠图结果。本文将详细介绍如何安装和使用 BEN2,包括单张图像处理、批量图像处理和视频分割的代码示例。

1. 安装 BEN2

1.1 克隆仓库并安装依赖

首先,您需要克隆 BEN2 的 GitHub 仓库,并安装所需的依赖项。请按照以下步骤操作:

git clone https://github.com/PramaLLC/BEN2.git
cd BEN2
pip install -r requirements.txt

1.2 安装 BEN2 模型

您可以通过以下命令直接安装 BEN2 模型:

pip install git+https://github.com/PramaLLC/BEN2.git

2. 使用 BEN2 处理单张图像

2.1 导入必要的库

import torch
from ben2 import AutoModel
from PIL import Image

2.2 加载模型并进行推理

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 打开图像
image = Image.open("image.jpg")  # 替换为您的图像路径

# 加载预训练模型
model = AutoModel.from_pretrained("PramaLLC/BEN2")
model.to(device).eval()

# 进行推理
foreground = model.inference(image)

# 保存前景图像
foreground.save("foreground.png")

2.3 代码解释

  • torch.device('cuda' if torch.cuda.is_available() else 'cpu'): 检查是否有可用的 GPU,如果有则使用 GPU,否则使用 CPU。
  • Image.open("image.jpg"): 打开输入图像,路径需要根据实际情况修改。
  • AutoModel.from_pretrained("PramaLLC/BEN2"): 从 Hugging Face 加载预训练的 BEN2 模型。
  • model.inference(image): 对输入图像进行推理,生成前景图像。
  • foreground.save("foreground.png"): 将生成的前景图像保存为 PNG 文件。

3. 批量图像处理

3.1 导入必要的库

import BEN2
from PIL import Image
import torch

3.2 初始化模型并加载权重

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 初始化模型
model = BEN2.BEN_Base().to(device).eval()

# 加载模型权重
model.loadcheckpoints("./BEN2_Base.pth")

3.3 批量处理多张图像

# 打开多张图像
file1 = "./image1.png"  # 输入图像1
file2 = "./image2.png"  # 输入图像2
image1 = Image.open(file1)
image2 = Image.open(file2)

# 批量推理
foregrounds = model.inference([image1, image2])

# 保存结果
foregrounds[0].save("./foreground1.png")
foregrounds[1].save("./foreground2.png")

3.4 代码解释

  • model.inference([image1, image2]): 对多张图像进行批量推理,返回多个前景图像。

  • foregrounds[0].save("./foreground1.png"): 保存第 1 张图像的前景结果。

  • foregrounds[1].save("./foreground2.png"): 保存第 2 张图像的前景结果。

  • 提示: 对于消费级 GPU,建议批量大小不超过 3,因为更大的批量并不会显著提高推理速度。

4. 视频分割

4.1 安装 FFmpeg

在进行视频分割之前,您需要安装 FFmpeg。可以通过以下命令安装:

sudo apt update
sudo apt install ffmpeg

4.2 视频分割代码

import BEN2
from PIL import Image
import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 初始化模型
model = BEN2.BEN_Base().to(device).eval()

# 加载模型权重
model.loadcheckpoints("./BEN2_Base.pth")

# 视频分割
model.segment_video(
    video_path="/path_to_your_video.mp4",  # 输入视频路径
    output_path="./",  # 输出路径,默认为当前目录
    fps=0,  # 如果设置为 0,CV2 将自动检测原始视频的帧率
    refine_foreground=False,  # 是否启用前景精炼,默认为 False
    batch=1,  # 批量大小,默认为 1
    print_frames_processed=True,  # 是否打印处理的帧数,默认为 True
    webm=False,  # 是否输出 WebM 格式的视频,默认为 False
    rgb_value=(0, 255, 0)  # 如果不使用 WebM,背景颜色将为绿色,默认为 (0, 255, 0)
)

4.3 代码解释

  • video_path="/path_to_your_video.mp4": 指定输入视频的路径。
  • output_path="./": 指定输出视频的保存路径,默认为当前目录。
  • fps=0: 如果设置为 0,CV2 将自动检测原始视频的帧率。
  • refine_foreground=False: 是否启用前景精炼,默认为 False。启用后可以提高边缘的准确性,但会增加推理时间。
  • batch=1: 批量大小,默认为 1。对于消费级 GPU,建议不超过 3。
  • print_frames_processed=True: 是否打印处理的帧数,默认为 True。
  • webm=False: 是否输出 WebM 格式的视频,默认为 False。如果不使用 WebM,背景颜色将为绿色。
  • rgb_value=(0, 255, 0): 如果不使用 WebM,背景颜色将为绿色,默认为 (0, 255, 0)。

资源


❤️ 如果你也关注 AI 的发展现状,且对 AI 应用开发感兴趣,我会每日分享大模型与 AI 领域的开源项目和应用,提供运行实例和实用教程,帮助你快速上手AI技术!

🥦 微信公众号|搜一搜:蚝油菜花 🥦