from PIL import Image
import torch
from torchvision import transforms
from model.modnet import MODNet  # или как у тебя организована модель
import cv2
import numpy as np

# Предобученная модель
modnet = MODNet()
modnet.load_state_dict(torch.load('pretrained/modnet_webcam_portrait_matting.ckpt', map_location='cpu'))
modnet.eval()

def remove_background_modnet(image: Image.Image) -> Image.Image:
    image = image.convert('RGB').resize((512, 512))
    img_np = np.array(image).astype(np.uint8)
    img_tensor = transforms.ToTensor()(img_np).unsqueeze(0)

    with torch.no_grad():
        _, _, matte = modnet(img_tensor)

    matte_np = matte.squeeze().cpu().numpy()
    matte_np = (matte_np * 255).astype(np.uint8)

    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGRA)
    img_np[:, :, 3] = matte_np

    return Image.fromarray(img_np)