简介

图像抠背景算子：用 u2net 显著性模型把主体从背景中抠出，输出带 alpha 通道的 PNG。直接用 onnxruntime 跑 u2net ONNX，不引入 rembg 依赖，显著性推理与抠图合成逻辑都实现在算子内。

功能描述

• 前处理：图 LANCZOS resize 到 320×320，先整体除以像素最大值，再按 mean=(0.485,0.456,0.406) / std=(0.229,0.224,0.225) 归一化，NCHW float32 送模型
• 后处理：取模型第一个输出做 min-max 归一化得到显著性 mask（输出全常量时 mask 全 0），LANCZOS resize 回原图尺寸，作为 alpha 通道合成 RGBA
• bgcolor 非空时把抠图 alpha_composite 到该 RGBA 纯色底上，输出不再透明；bgcolor 必须是 4 元组，否则构造期报错
• alpha_matting=True 走 pymatting 的 estimate_alpha_cf + estimate_foreground_ml（trimap 由前景/背景两个阈值裁出）；镜像未装 pymatting 时构造期直接抛 ValueError，不静默降级
• 不用 rembg 的原因：rembg >= 2.0.55 都要求 pillow>=12.1，而镜像里是 Pillow 11.3，装它会顶掉几乎所有图像/视频算子共用的 PIL；u2net 本身就是一个 ONNX 模型，onnxruntime 已在镜像内
• 权重只从本地 {model_path}/{model_name}.onnx 读，缺失时构造期抛 FileNotFoundError，不联网下载
• 每行输入先落到 /tmp 下的临时目录再处理，行处理结束自动清理；image_url 且是本地路径时不复制，直接读原文件
• 输出统一是 PNG，文件名 <本地文件 stem>_nobg.png：本地路径输入取原图 stem，bos:// / s3:// 输入取 image_<行内下标>_0，base64 / binary 输入取 <unix 时间戳>_<行内下标>_0
• GPU 需要 num_gpus > 0 且 onnxruntime 报告 CUDAExecutionProvider 可用，否则用 CPU（u2net 在 CPU 上可用，回归测试即 CPU 跑）
• 行级失败（输入为 null、下载或解码失败等）返回 None 并计入内部 failed 计数，不中断整列

算子参数

输入

输出

参数

注意事项

• model_name 不要带.onnx后缀：传成 rembg/u2net.onnx 会多拼一次后缀而找不到权重文件。
• alpha_matting_erode_size 只是保留的参数位，当前实现的 trimap 只用前景/背景两个阈值，没有做腐蚀，改这个值不影响结果。
• 输出文件名只由输入文件 stem 决定：不同目录下同名的本地图片会写到同一个输出路径而互相覆盖，批量场景建议按批设置不同的 output_dir。
• 需要不透明结果时传 bgcolor（例如 [255, 255, 255, 255]），否则下游若按 RGB 打开 PNG 会得到黑色背景。

调用示例

1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.image.image_remove_background import ImageRemoveBackground
10
11if __name__ == "__main__":
12    if os.getenv("DAFT_RUNNER", "native") == "ray":
13        import ray
14        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15        daft.set_runner_ray()
16    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18samples = {"image": ["bos://your-bucket/sample.jpg"]}
19    ds = daft.from_pydict(samples)
20    ds = ds.with_column(
21        "cutout",
22        aihc_udf(
23            ImageRemoveBackground,
24            construct_args={
25                "image_src_type": "image_url",
26                "model_path": "/path/to/models",
27                "model_name": "rembg/u2net",
28                "output_dir": "/tmp/aihc_image_remove_background",
29            },
30            num_cpus=1,
31            concurrency=1,
32            batch_size=1,
33        )(col("image")),
34    )
35    ds.show()