Create modnet_model.py

model/modnet_model.py

@@ -0,0 +1,27 @@
+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)