U2Net模型的去背應用

Claire Chang

2023-09-27 pm 12:39

官方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()