44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
import fastdeploy as fd
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
class FastDeployMattingModule:
|
|
def __init__(self, encoder_onnx, decoder_onnx, use_gpu=False, device_id=0):
|
|
rt_opt = fd.RuntimeOption()
|
|
if use_gpu:
|
|
rt_opt.use_gpu(device_id)
|
|
else:
|
|
rt_opt.use_cpu()
|
|
self.model = fd.vision.matting.SAM(
|
|
model_file=encoder_onnx,
|
|
params_file=decoder_onnx,
|
|
runtime_option=rt_opt,
|
|
model_format=fd.ModelFormat.ONNX,
|
|
)
|
|
|
|
def remove_background(self, img_path):
|
|
img = cv2.imread(img_path)
|
|
if img is None:
|
|
raise FileNotFoundError("Failed to load:", img_path)
|
|
|
|
res = self.model.predict(img)
|
|
alpha = (res.alpha * 255).astype(np.uint8)
|
|
bgr = res.foreground
|
|
|
|
bgra = cv2.cvtColor(bgr, cv2.COLOR_BGR2BGRA)
|
|
bgra[..., 3] = alpha
|
|
return Image.fromarray(cv2.cvtColor(bgra, cv2.COLOR_BGRA2RGBA))
|
|
|
|
|
|
# 사용 예시
|
|
if __name__ == "__main__":
|
|
fm = FastDeployMattingModule(
|
|
encoder_onnx="sam_vit_b_01ec64.encoder.onnx",
|
|
decoder_onnx="sam_vit_b_01ec64.decoder.onnx",
|
|
use_gpu=False
|
|
)
|
|
img_rgba = fm.remove_background("8.jpg")
|
|
img_rgba.save("8_sam_matting.png")
|
|
print("✅ 배경제거 완료 및 저장됨")
|