Update model/modnet.py

model/modnet.py

@@ -6,14 +6,32 @@ from torchvision.models.mobilenetv2 import mobilenet_v2
 import torch.nn as nn
 
 
+def clean_state_dict(state_dict):
+    """Remove 'module.' prefix if present in keys."""
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        if k.startswith('module.'):
+            new_state_dict[k[7:]] = v
+        else:
+            new_state_dict[k] = v
+    return new_state_dict
+
+
 class MODNet(nn.Module):
     def __init__(self, in_channels=3, hr_channels=32, backbone_arch='mobilenetv2', backbone_pretrained=True):
         super(MODNet, self).__init__()
-        self.backbone = nn.Identity()  # Replaced with identity for fast test
+
+        self.backbone = mobilenet_v2(pretrained=backbone_pretrained).features
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(1280, 320)  # Example dimensions, adjust as needed
 
     def forward(self, x, inference=False):
-        # Fast dummy forward to avoid GPU/memory bottlenecks
-        return x, x, x  # dummy semantic, detail, matte
+        features = self.backbone(x)
+        pooled = self.avgpool(features)
+        flattened = torch.flatten(pooled, 1)
+        semantic = self.fc(flattened)
+        return semantic, semantic, semantic  # Dummy triple output for compatibility
+
 
 def preprocess_image(image: Image.Image, device: torch.device) -> torch.Tensor:
     img = np.array(image.convert("RGB"))
@@ -23,22 +41,16 @@ def preprocess_image(image: Image.Image, device: torch.device) -> torch.Tensor:
     return img_input
 
 
-def clean_state_dict(state_dict):
-    """Remove 'module.' prefix if present in keys."""
-    new_state_dict = {}
-    for k, v in state_dict.items():
-        if k.startswith('module.'):
-            new_state_dict[k[7:]] = v
-        else:
-            new_state_dict[k] = v
-    return new_state_dict
-
-
 def remove_background_modnet(image: Image.Image) -> Image.Image:
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
     modnet = MODNet()
     modnet.to(device)
+
+    # Load weights
+    state_dict = torch.load('pretrained/modnet_webcam_portrait_matting.ckpt', map_location=device)
+    modnet.load_state_dict(clean_state_dict(state_dict), strict=False)
+
     modnet.eval()
 
     img_input = preprocess_image(image, device)