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("✅ 배경제거 완료 및 저장됨")