CV-UNet Universal Matting代码实例：自定义抠图功能开发

1. 引言

1.1 背景与需求

在图像处理和计算机视觉领域，图像抠图（Image Matting）是一项关键任务，广泛应用于电商、广告设计、影视后期和AI生成内容（AIGC）等场景。传统抠图方法依赖人工标注或简单阈值分割，效率低且难以应对复杂边缘（如发丝、透明物体）。随着深度学习的发展，基于语义分割和Alpha预测的自动抠图技术逐渐成熟。

CV-UNet Universal Matting 正是在这一背景下推出的高效解决方案。它基于经典的U-Net 架构进行改进，结合现代注意力机制与轻量化设计，实现了高精度、快速响应的通用抠图能力。该项目由开发者“科哥”进行二次开发并封装为 WebUI 工具，支持一键部署、批量处理和本地化运行，极大降低了使用门槛。

1.2 技术价值与应用场景

CV-UNet 的核心优势在于其通用性（Universal）和实用性（Practical）：

支持人物、产品、动物等多种主体类型
输出带 Alpha 通道的 PNG 图像，适用于设计软件（Photoshop、Figma）
提供单图实时预览与批量自动化处理模式
可集成至企业内部系统，用于商品图自动化背景移除

本文将围绕该系统的实际应用展开，重点介绍其架构原理、核心代码实现以及如何在此基础上进行二次开发，构建可定制的智能抠图服务。

2. 系统架构与工作流程解析

2.1 整体架构概览

CV-UNet Universal Matting 采用典型的前后端分离架构：

[用户界面 WebUI] ↓ (HTTP 请求) [Flask 后端服务] ↓ (模型推理) [PyTorch 模型: CV-UNet] ↓ (输出结果) [文件系统 outputs/]

前端基于 HTML + JavaScript 实现交互式界面，后端使用 Flask 框架接收请求并调用 PyTorch 模型完成推理，最终返回处理后的图像数据。

2.2 核心组件说明

组件	功能
CV-UNet 模型	基于 U-Net 改进的编码器-解码器结构，输出四通道 RGBA 结果
Flask API 服务	接收图片上传、触发推理、返回结果路径
WebUI 界面	提供可视化操作入口，支持拖拽上传、结果预览
批量处理器	遍历指定目录中的所有图片并依次调用模型

2.3 工作流程详解

用户通过 WebUI 上传图片或输入文件夹路径
Flask 接收请求，保存图片至临时目录
调用inference.py中的模型推理函数
模型前向传播，生成带有透明通道的 Alpha mask
合并原图与 Alpha 通道，输出 RGBA 格式的 PNG 文件
返回结果链接并在前端展示

整个过程对用户完全透明，仅需点击按钮即可完成。

3. 核心代码实现分析

3.1 模型定义：CV-UNet 结构

以下是简化版的 CV-UNet 模型定义代码，基于 PyTorch 实现：

# models/cv_unet.py import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): """双卷积块：conv + bn + relu ×2""" def __init__(self, in_channels, out_channels): super().__init__() self.double_conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.double_conv(x) class Down(nn.Module): """下采样模块""" def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_channels, out_channels) ) def forward(self, x): return self.maxpool_conv(x) class Up(nn.Module): """上采样模块""" def __init__(self, in_channels, out_channels): super().__init__() self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv = DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 = self.up(x1) diffY = x2.size()[2] - x1.size()[2] diffX = x2.size()[3] - x1.size()[3] x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x = torch.cat([x2, x1], dim=1) return self.conv(x) class OutConv(nn.Module): """输出层：生成4通道RGBA""" def __init__(self, in_channels, out_channels=4): super(OutConv, self).__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) def forward(self, x): return torch.sigmoid(self.conv(x)) class CVUNet(nn.Module): def __init__(self, n_channels=3): super(CVUNet, self).__init__() self.inc = DoubleConv(n_channels, 64) self.down1 = Down(64, 128) self.down2 = Down(128, 256) self.down3 = Down(256, 512) self.down4 = Down(512, 1024) self.up1 = Up(1024 + 512, 512) self.up2 = Up(512 + 256, 256) self.up3 = Up(256 + 128, 128) self.up4 = Up(128 + 64, 64) self.outc = OutConv(64) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) logits = self.outc(x) return logits

说明：该模型输出为 4 通道张量，分别对应 R、G、B 和 A（Alpha）通道，其中 Alpha 表示前景透明度。

3.2 推理逻辑实现

# inference.py import cv2 import numpy as np from PIL import Image import torch from models.cv_unet import CVUNet def load_model(model_path, device): model = CVUNet().to(device) model.load_state_dict(torch.load(model_path, map_location=device)) model.eval() return model def preprocess_image(image_path, target_size=(512, 512)): img = cv2.imread(image_path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w = img.shape[:2] img_resized = cv2.resize(img, target_size) img_tensor = torch.from_numpy(img_resized).float() / 255.0 img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0) # (C, H, W) -> (1, C, H, W) return img_tensor, (h, w) def postprocess_output(output_tensor, original_size): output = output_tensor.squeeze().cpu().detach().numpy() # (4, H, W) alpha = (output[3] * 255).astype(np.uint8) # Alpha 通道 rgb = (output[:3] * 255).astype(np.uint8).transpose(1, 2, 0) # 调整大小回原始尺寸 alpha = cv2.resize(alpha, original_size[::-1], interpolation=cv2.INTER_CUBIC) rgb = cv2.resize(rgb, original_size[::-1], interpolation=cv2.INTER_CUBIC) # 合成 RGBA 图像 rgba = np.dstack((rgb, alpha)) return Image.fromarray(rgba, 'RGBA') def run_inference(input_path, model, device, output_dir): input_tensor, orig_size = preprocess_image(input_path) with torch.no_grad(): output_tensor = model(input_tensor.to(device)) result_img = postprocess_output(output_tensor, orig_size) # 保存结果 filename = os.path.basename(input_path).rsplit('.', 1)[0] + '.png' save_path = os.path.join(output_dir, filename) result_img.save(save_path) return save_path

上述代码完成了从图像读取、预处理、模型推理到后处理输出的完整流程。

4. 自定义功能扩展实践

4.1 添加批量处理功能

为了支持批量处理，我们编写一个简单的批处理脚本：

# batch_processor.py import os from pathlib import Path from inference import run_inference, load_model def process_folder(folder_path, model, device, output_base="outputs"): folder = Path(folder_path) image_files = list(folder.glob("*.jpg")) + list(folder.glob("*.png")) + list(folder.glob("*.webp")) if not image_files: print("未找到支持的图片文件") return # 创建输出目录 timestamp = datetime.now().strftime("%Y%m%d%H%M%S") output_dir = os.path.join(output_base, f"outputs_{timestamp}") os.makedirs(output_dir, exist_ok=True) success_count = 0 for img_file in image_files: try: save_path = run_inference(str(img_file), model, device, output_dir) print(f"✅ 处理成功: {img_file.name} -> {save_path}") success_count += 1 except Exception as e: print(f"❌ 处理失败 {img_file.name}: {str(e)}") print(f"\n批量处理完成！成功 {success_count}/{len(image_files)} 张")

此脚本可用于 WebUI 后端的批量处理接口。

4.2 Flask 接口集成

# app.py from flask import Flask, request, jsonify, send_from_directory import os from batch_processor import process_folder from inference import load_model, run_inference app = Flask(__name__) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") MODEL = load_model("checkpoints/cvunet.pth", DEVICE) @app.route('/single', methods=['POST']) def single_inference(): file = request.files['image'] temp_path = os.path.join('temp', file.filename) file.save(temp_path) output_dir = 'outputs/latest_single' os.makedirs(output_dir, exist_ok=True) result_path = run_inference(temp_path, MODEL, DEVICE, output_dir) result_url = f"/result/{os.path.relpath(result_path, '.')}" return jsonify({'result_url': result_url}) @app.route('/batch', methods=['POST']) def batch_inference(): data = request.json folder_path = data.get('folder_path') if not os.path.exists(folder_path): return jsonify({'error': '文件夹不存在'}), 400 process_folder(folder_path, MODEL, DEVICE) return jsonify({'status': 'success', 'message': '批量处理已启动'}) @app.route('/result/<path:filename>') def serve_result(filename): return send_from_directory('.', filename) if __name__ == '__main__': os.makedirs('temp', exist_ok=True) app.run(host='0.0.0.0', port=7860)

该接口支持/single单图上传和/batch批量处理请求，便于前端调用。