tio-boot 整合 U2Net 实现图片去背景

本文介绍如何使用 tio-boot + ONNX Runtime + U2Net 搭建一个 Java 版图片去背景 HTTP 服务。接口行为兼容常见的 rembg HTTP 上传方式：

POST /api/remove
Content-Type: multipart/form-data

file: 二进制图片

响应直接返回透明背景 PNG：

Content-Type: image/png

一、功能说明

U2Net 是一种图像显著性目标分割模型。它不是按颜色简单删除背景，而是对输入图片做前景分割，输出一张灰度 mask：

白色：主体
黑色：背景
灰色：半透明边缘

服务端处理流程如下：

客户端上传图片
   ↓
tio-boot 接收 multipart/form-data
   ↓
ImageIO 解码为 BufferedImage
   ↓
缩放到 U2Net 输入尺寸 320x320
   ↓
RGB → CHW float tensor
   ↓
ONNX Runtime 推理
   ↓
提取 mask 并归一化
   ↓
mask resize 回原图尺寸
   ↓
写入原图 alpha 通道
   ↓
输出透明 PNG

本示例不依赖 OpenCV，只使用 JDK 自带的 ImageIO 和 BufferedImage 完成图片解码、缩放和 alpha 合成。

二、项目依赖

pom.xml 示例：

<project xmlns="http://maven.apache.org/POM/4.0.0"
  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
  xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
  <modelVersion>4.0.0</modelVersion>
  <groupId>nexus.io</groupId>
  <artifactId>u2net-server</artifactId>
  <version>1.0.0</version>

  <properties>
    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    <java.version>21</java.version>
    <maven.compiler.source>${java.version}</maven.compiler.source>
    <maven.compiler.target>${java.version}</maven.compiler.target>
    <tio-boot.version>2.1.4</tio-boot.version>
    <onnxruntime.version>1.16.1</onnxruntime.version>
    <main.class>nexus.io.u2net.U2NetServerApp</main.class>
    <final.name>u2net-server</final.name>
  </properties>

  <dependencies>
    <dependency>
      <groupId>nexus.io</groupId>
      <artifactId>tio-boot</artifactId>
      <version>${tio-boot.version}</version>
    </dependency>

    <dependency>
      <groupId>com.microsoft.onnxruntime</groupId>
      <artifactId>onnxruntime</artifactId>
      <version>${onnxruntime.version}</version>
    </dependency>

    <dependency>
      <groupId>ch.qos.logback</groupId>
      <artifactId>logback-classic</artifactId>
      <version>1.3.3</version>
    </dependency>
  </dependencies>

  <build>
    <finalName>${final.name}</finalName>
    <plugins>
      <plugin>
        <groupId>org.apache.maven.plugins</groupId>
        <artifactId>maven-compiler-plugin</artifactId>
        <version>3.13.0</version>
        <configuration>
          <release>${java.version}</release>
        </configuration>
      </plugin>

      <plugin>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-maven-plugin</artifactId>
        <version>2.7.18</version>
        <configuration>
          <mainClass>${main.class}</mainClass>
        </configuration>
        <executions>
          <execution>
            <goals>
              <goal>repackage</goal>
            </goals>
          </execution>
        </executions>
      </plugin>
    </plugins>
  </build>
</project>

三、准备模型

默认模型文件：

u2net.onnx

下载地址：

https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx

推荐放到 rembg 默认模型目录：

Windows: C:\Users\你的用户名\.u2net\u2net.onnx
Linux/Mac: ~/.u2net/u2net.onnx

服务代码默认读取：

String userHome = EnvUtils.getStr("user.home");
this.modelPath = EnvUtils.get("u2net.model", userHome + "/.u2net/u2net.onnx");

也可以通过配置覆盖：

u2net.model=F:/models/u2net.onnx

或启动参数覆盖：

java -jar target/u2net-server.jar --u2net.model=F:/models/u2net.onnx

四、配置文件

src/main/resources/app.properties

server.port=7000

如需指定模型和并发数：

server.port=7000
u2net.model=F:/models/u2net.onnx
u2net.concurrent=4

u2net.concurrent 用于限制同时推理的请求数量，避免 CPU 和 native 内存被打满。

五、启动类

U2NetServerApp.java

package nexus.io.u2net;

import nexus.io.tio.boot.TioApplication;
import nexus.io.u2net.config.U2NetServerConfig;

public class U2NetServerApp {

  public static void main(String[] args) {
    TioApplication.run(U2NetServerApp.class, new U2NetServerConfig(), args);
  }
}

六、tio-boot 路由配置

U2NetServerConfig.java

package nexus.io.u2net.config;

import nexus.io.context.BootConfiguration;
import nexus.io.hook.HookCan;
import nexus.io.tio.boot.server.TioBootServer;
import nexus.io.tio.http.server.router.HttpRequestRouter;
import nexus.io.u2net.handler.HealthHandler;
import nexus.io.u2net.handler.RemoveBackgroundHandler;
import nexus.io.u2net.service.U2NetBackgroundRemovalService;

public class U2NetServerConfig implements BootConfiguration {

  @Override
  public void config() {
    HttpRequestRouter router = TioBootServer.me().getRequestRouter();
    U2NetBackgroundRemovalService removalService = new U2NetBackgroundRemovalService();

    router.add("/health", new HealthHandler());
    router.add("/api/remove", new RemoveBackgroundHandler(removalService));

    HookCan.me().addDestroyMethod(removalService::close);
  }
}

这里使用 HookCan.me().addDestroyMethod(removalService::close) 注册销毁方法。tio-boot 停止时会执行 HookCan.me().stop()，从而释放 ONNX Runtime 的 OrtSession 和 OrtEnvironment。

不要优先使用：

Runtime.getRuntime().addShutdownHook(new Thread(removalService::close));

在 tio-boot 项目中，HookCan 更符合框架生命周期。

七、HTTP Handler

RemoveBackgroundHandler.java

package nexus.io.u2net.handler;

import nexus.io.model.upload.UploadFile;
import nexus.io.tio.boot.http.TioRequestContext;
import nexus.io.tio.http.common.HttpRequest;
import nexus.io.tio.http.common.HttpResponse;
import nexus.io.tio.http.server.handler.HttpRequestHandler;
import nexus.io.tio.http.server.util.CORSUtils;
import nexus.io.u2net.service.U2NetBackgroundRemovalService;

public class RemoveBackgroundHandler implements HttpRequestHandler {

  private final U2NetBackgroundRemovalService removalService;

  public RemoveBackgroundHandler(U2NetBackgroundRemovalService removalService) {
    this.removalService = removalService;
  }

  @Override
  public HttpResponse handle(HttpRequest request) throws Exception {
    HttpResponse response = TioRequestContext.getResponse();
    CORSUtils.enableCORS(response);

    UploadFile file = request.getUploadFile("file");
    if (file == null || file.getData() == null || file.getData().length == 0) {
      return response.fail("multipart field 'file' is required");
    }

    byte[] png = removalService.removeBackground(file.getData());
    response.setContentType("image/png");
    response.setAttachmentFilename("output.png");
    response.body(png);
    return response;
  }
}

说明：

上传字段名保持为 file，兼容 rembg 的 POST /api/remove 调用习惯。
响应直接写入 PNG 字节数组，不返回 JSON。
TioRequestContext.getResponse() 从当前 tio-boot 请求上下文中获取响应对象，推荐在 tio-boot Handler 中使用。

健康检查 Handler：

package nexus.io.u2net.handler;

import java.util.Map;

import nexus.io.tio.http.common.HttpRequest;
import nexus.io.tio.http.common.HttpResponse;
import nexus.io.tio.http.server.handler.HttpRequestHandler;

public class HealthHandler implements HttpRequestHandler {

  @Override
  public HttpResponse handle(HttpRequest request) {
    return new HttpResponse(request).setJson(Map.of("ok", true));
  }
}

八、核心推理服务

U2NetBackgroundRemovalService 负责：

加载 U2Net ONNX 模型
图片预处理
创建 OnnxTensor
执行 session.run
归一化 mask
mask resize 回原图大小
合成 alpha PNG
释放 ONNX Runtime 资源

核心代码如下：

package nexus.io.u2net.service;

import java.awt.Graphics2D;
import java.awt.RenderingHints;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.nio.FloatBuffer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Collections;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

import javax.imageio.ImageIO;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtException;
import ai.onnxruntime.OrtSession;
import nexus.io.tio.utils.environment.EnvUtils;

public class U2NetBackgroundRemovalService implements AutoCloseable {

  private static final int INPUT_WIDTH = 320;
  private static final int INPUT_HEIGHT = 320;
  private static final float[] MEAN = { 0.485f, 0.456f, 0.406f };
  private static final float[] STD = { 0.229f, 0.224f, 0.225f };

  private final String modelPath;
  private final OrtEnvironment env;
  private final OrtSession session;
  private final String inputName;
  private final Semaphore semaphore;

  public U2NetBackgroundRemovalService() {
    String userHome = EnvUtils.getStr("user.home");

    this.modelPath = EnvUtils.get("u2net.model", userHome + "/.u2net/u2net.onnx");
    this.semaphore = new Semaphore(EnvUtils.getInt("u2net.concurrent", Runtime.getRuntime().availableProcessors()));

    if (!Files.isRegularFile(Path.of(modelPath))) {
      throw new IllegalStateException("U2Net model not found: " + modelPath);
    }

    try {
      this.env = OrtEnvironment.getEnvironment();
      this.session = env.createSession(modelPath, new OrtSession.SessionOptions());
      this.inputName = session.getInputInfo().keySet().iterator().next();
      System.out.println("U2Net model loaded: " + modelPath + ", input=" + inputName);
    } catch (OrtException e) {
      throw new IllegalStateException("Failed to load U2Net model: " + modelPath, e);
    }
  }

  public byte[] removeBackground(byte[] imageBytes) throws Exception {
    if (!semaphore.tryAcquire(30, TimeUnit.SECONDS)) {
      throw new IllegalStateException("Server busy, please retry later");
    }

    try {
      BufferedImage source = ImageIO.read(new ByteArrayInputStream(imageBytes));
      if (source == null) {
        throw new IllegalArgumentException("Unsupported or invalid image");
      }

      BufferedImage rgbInput = resizeToRgb(source, INPUT_WIDTH, INPUT_HEIGHT);
      float[] input = toNormalizedChw(rgbInput);
      float[][] mask = inferMask(input);
      BufferedImage alphaMask = resizeMask(mask, source.getWidth(), source.getHeight());
      BufferedImage result = applyAlpha(source, alphaMask);

      ByteArrayOutputStream output = new ByteArrayOutputStream();
      ImageIO.write(result, "png", output);
      return output.toByteArray();
    } finally {
      semaphore.release();
    }
  }

  private float[][] inferMask(float[] input) throws OrtException {
    long[] shape = { 1L, 3L, INPUT_HEIGHT, INPUT_WIDTH };
    try (OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(input), shape);
        OrtSession.Result result = session.run(Collections.singletonMap(inputName, tensor))) {
      Object value = result.get(0).getValue();
      float[][] rawMask = extractFirstMask(value);
      return normalizeMask(rawMask);
    }
  }

  private float[][] extractFirstMask(Object value) {
    if (value instanceof float[][][][] output4d) {
      return output4d[0][0];
    }
    if (value instanceof float[][][] output3d) {
      return output3d[0];
    }
    if (value instanceof float[][] output2d) {
      return output2d;
    }
    throw new IllegalStateException("Unsupported U2Net output type: " + value.getClass().getName());
  }

  private float[][] normalizeMask(float[][] mask) {
    float min = Float.MAX_VALUE;
    float max = -Float.MAX_VALUE;
    for (float[] row : mask) {
      for (float v : row) {
        min = Math.min(min, v);
        max = Math.max(max, v);
      }
    }

    float range = max - min;
    if (range < 1e-6f) {
      return mask;
    }

    float[][] normalized = new float[mask.length][mask[0].length];
    for (int y = 0; y < mask.length; y++) {
      for (int x = 0; x < mask[y].length; x++) {
        normalized[y][x] = clamp((mask[y][x] - min) / range);
      }
    }
    return normalized;
  }

  private BufferedImage resizeToRgb(BufferedImage source, int width, int height) {
    BufferedImage resized = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
    Graphics2D g = resized.createGraphics();
    g.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BILINEAR);
    g.setRenderingHint(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);
    g.drawImage(source, 0, 0, width, height, null);
    g.dispose();
    return resized;
  }

  private float[] toNormalizedChw(BufferedImage image) {
    int width = image.getWidth();
    int height = image.getHeight();
    float[] chw = new float[3 * width * height];

    for (int y = 0; y < height; y++) {
      for (int x = 0; x < width; x++) {
        int rgb = image.getRGB(x, y);
        int offset = y * width + x;
        float r = ((rgb >> 16) & 0xff) / 255.0f;
        float g = ((rgb >> 8) & 0xff) / 255.0f;
        float b = (rgb & 0xff) / 255.0f;

        chw[offset] = (r - MEAN[0]) / STD[0];
        chw[width * height + offset] = (g - MEAN[1]) / STD[1];
        chw[2 * width * height + offset] = (b - MEAN[2]) / STD[2];
      }
    }
    return chw;
  }

  private BufferedImage resizeMask(float[][] mask, int width, int height) {
    BufferedImage small = new BufferedImage(mask[0].length, mask.length, BufferedImage.TYPE_BYTE_GRAY);
    for (int y = 0; y < mask.length; y++) {
      for (int x = 0; x < mask[y].length; x++) {
        int v = Math.round(clamp(mask[y][x]) * 255.0f);
        small.getRaster().setSample(x, y, 0, v);
      }
    }

    BufferedImage resized = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_GRAY);
    Graphics2D g = resized.createGraphics();
    g.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BILINEAR);
    g.drawImage(small, 0, 0, width, height, null);
    g.dispose();
    return resized;
  }

  private BufferedImage applyAlpha(BufferedImage source, BufferedImage alphaMask) {
    BufferedImage result = new BufferedImage(source.getWidth(), source.getHeight(), BufferedImage.TYPE_INT_ARGB);
    for (int y = 0; y < source.getHeight(); y++) {
      for (int x = 0; x < source.getWidth(); x++) {
        int rgb = source.getRGB(x, y);
        int alpha = alphaMask.getRaster().getSample(x, y, 0);
        result.setRGB(x, y, (alpha << 24) | (rgb & 0x00ffffff));
      }
    }
    return result;
  }

  private float clamp(float value) {
    return Math.max(0.0f, Math.min(1.0f, value));
  }

  @Override
  public void close() {
    try {
      session.close();
    } catch (OrtException e) {
      throw new RuntimeException(e);
    }
    env.close();
  }
}

九、关于 OnnxTensor 创建

每次推理都需要根据当前图片创建输入 tensor：

OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(input), shape);

U2Net 输入大小为：

1 * 3 * 320 * 320 float
= 307200 float
= 约 1.2 MB

这行代码有一定开销，但通常不是整次请求中最大的开销。更耗时的步骤通常是：

图片解码
resize
Java 循环做 RGB → CHW → normalize
session.run
mask resize
PNG 编码

推荐写法是每次请求创建一次 OnnxTensor，用完立即关闭：

try (OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(input), shape);
    OrtSession.Result result = session.run(Collections.singletonMap(inputName, tensor))) {
  // process output
}

不要把同一个 OnnxTensor 跨请求复用。每张图片输入数据不同，复用 tensor 容易引入并发和数据污染问题。真正需要优化时，优先考虑：

使用 ThreadLocal<float[]> 复用输入数组
批量读取像素，减少 getRGB(x, y) 的逐点调用
控制 u2net.concurrent
增加分阶段耗时日志，确认瓶颈后再优化

十、运行与测试

编译：

mvn -q -DskipTests compile

打包：

mvn -q -DskipTests package

启动：

java -jar target/u2net-server.jar

指定模型路径启动：

java -jar target/u2net-server.jar --u2net.model=F:/models/u2net.onnx

健康检查：

curl.exe "http://localhost:7000/health"

上传图片去背景：

curl.exe -F "[email protected]" "http://localhost:7000/api/remove" -o output.png

输出必须使用 PNG，因为透明背景需要 alpha 通道，JPG 不支持透明背景。

十一、常见问题

1. 模型文件不存在

错误示例：

U2Net model not found: C:/Users/xxx/.u2net/u2net.onnx

解决方式：

下载 u2net.onnx
放到 ~/.u2net/u2net.onnx
或通过 u2net.model 指定模型路径

2. 返回图片没有透明背景

确认客户端保存为 PNG：

curl.exe -F "[email protected]" "http://localhost:7000/api/remove" -o output.png

不要保存为 JPG。

3. 并发请求时 CPU 很高

U2Net CPU 推理比较消耗计算资源。可以降低并发：

u2net.concurrent=2

也可以换更小的模型，例如 u2netp.onnx，但效果会弱一些。

4. 边缘效果不如 rembg

rembg 支持 alpha matting，可以改善头发、毛发和半透明边缘。本示例实现的是基础 mask → alpha 合成流程，没有额外实现 alpha matting。需要更好边缘时，可以在 mask 后处理阶段增加：

mask blur
erode / dilate
guided filter
alpha matting

十二、总结

通过上述实现，可以把原来 Python rembg 的 HTTP 去背景能力迁移为 Java 服务：

tio-boot 接口
  + ONNX Runtime 推理
  + U2Net mask
  + Java ImageIO 合成透明 PNG

接口保持为：

POST /api/remove
[email protected]

因此前端、curl、网关或已有调用方基本不需要改造，只需要把后端服务从 rembg 切换到 Java 版 U2Net 服务即可。