Update modnet_utils.py

modnet_utils.py

@@ -1,63 +1,5 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-import cv2
 from PIL import Image
-from torchvision import transforms
 
-
-class MODNet(nn.Module):
-    def __init__(self, backbone):
-        super(MODNet, self).__init__()
-        self.backbone = backbone
-        self.seg_head = nn.Sequential(
-            nn.Conv2d(1280, 64, kernel_size=3, padding=1),  # changed from 320 to 1280
-            nn.ReLU(),
-            nn.Conv2d(64, 1, kernel_size=1),
-            nn.Sigmoid()
-        )
-
-    def forward(self, x):
-        features = self.backbone(x)
-        pred_matte = self.seg_head(features)
-        return pred_matte
-
-
-def preprocess_image(image: Image.Image, device: torch.device) -> torch.Tensor:
-    transform = transforms.Compose([
-        transforms.Resize((512, 512)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                             std=[0.229, 0.224, 0.225])
-    ])
-    img_tensor = transform(image.convert("RGB")).unsqueeze(0).to(device)
-    return img_tensor
-
-
-def remove_background_modnet(image: Image.Image) -> Image.Image:
-    from torchvision.models.mobilenet import mobilenet_v2
-
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-    backbone = mobilenet_v2(pretrained=True).features
-    modnet = MODNet(backbone)
-    modnet.to(device)
-
-    state_dict = torch.load('pretrained/modnet_webcam_portrait_matting.ckpt', map_location=device)
-    modnet.load_state_dict(state_dict, strict=False)
-    modnet.eval()
-
-    img_tensor = preprocess_image(image, device)
-
-    with torch.no_grad():
-        pred_matte = modnet(img_tensor)
-
-    matte = pred_matte[0][0].cpu().numpy()
-    matte = cv2.resize(matte, image.size)
-    matte = np.uint8(matte * 255)
-
-    image = image.convert("RGBA")
-    image_np = np.array(image)
-    image_np[:, :, 3] = matte
-    return Image.fromarray(image_np)
+def remove_background(image: Image.Image) -> Image.Image:
+    from modnet_model import remove_background_modnet
+    return remove_background_modnet(image)