Posted on

U2Net模型的去背應用

官方GitHub: https://github.com/danielgatis/rembg/tree/main

安裝方式

CPU support

pip install rembg # for library
pip install rembg[cli] # for library + cli

GPU support

首先要去安裝onnxruntime-gpu,接著

pip install rembg[gpu] # for library
pip install rembg[gpu,cli] # for library + cli

使用範例

下面為一個讀取攝影機的簡單去背使用範例

import VideoStream
import cv2
import numpy as np
import time
from rembg import new_session, remove

videostream = VideoStream.VideoStream((1280, 720), 30, 0).start()
cam_quit = 0
total = 0
frame = 0

while cam_quit == 0:
    imageSource = videostream.read()
    imageSource = cv2.resize(imageSource, (640,360))
    aStart = time.time()
    # u2netp為model名稱
    output = remove(imageSource, session = new_session("u2netp") )
    aEnd = time.time()
    if frame <= 60:
        total = total + (aEnd - aStart)
        frame = frame + 1
    print("detect time: " + str(aEnd - aStart ))
    cv2.imshow("output", output)
    key = cv2.waitKey(1) & 0xFF
    if key == ord("q"):
        print("avg detect time: " + str(total/frame))
        cam_quit = 1

videostream.stop()
cv2.destroyAllWindows()

Stable Diffusion WebUI Rembg

Stable Diffusion WebUI所使用的去背功能也是這個唷!

相關介紹文章: https://zhuanlan.zhihu.com/p/648234420

可以參考WebUI的參數去尋找相關程式碼可設定的部分,例如Erode size、Foreground threshold、Background threshold

模型去背測試

原始影像如下圖

模型列表

u2net

  • 檔案大小: 171,873KB
  • 處理時間: 1.741s

u2netp

  • 檔案大小: 4,468KB
  • 處理時間: 0.702s

u2net_human_seg

  • 檔案大小: 171,873KB
  • 處理時間: 1.997s

u2net_cloth_seg

  • 檔案大小: 172,066KB
  • 處理時間: 5.496s

silueta

  • 檔案大小: 43,138KB
  • 處理時間: 5.496s

isnet-general-use

  • 檔案大小: 174,461KB
  • 處理時間: 2.629s

isnet-anime

  • 檔案大小: 171,944KB
  • 處理時間: 2.533s