Update app.py

app.py

@@ -28,12 +28,14 @@ class TransformNet(nn.Module):
         self.fc = nn.Linear(1024, 512)
 
     def forward(self, x):
-        x = x.unsqueeze(-1)  # ✅ Add extra dimension for Conv2d
+        if x.dim() == 5:  
+            x = x.squeeze(-1)  
+        
         x = self.conv1(x)
         x = self.conv2(x)
-        x = x.squeeze(-1)  # ✅ Remove extra dimension before Conv1d
-        x = self.conv3(x)  # ✅ Now it correctly matches `[1024, 128, 1]`
-        return self.fc(x.max(dim=-1)[0])  # ✅ Fix pooling
+        x = self.conv3(x)
+        return self.fc(x.max(dim=-1)[0])  # ✅ Ensure correct pooling
+
 
 class PointCloudEncoder(nn.Module):
     """Point Cloud Encoder (pc_enc)"""
@@ -125,11 +127,11 @@ model.eval()
 def segment_dress(image):
     """Detect and segment the dress from the image."""
     img = Image.fromarray(image).convert("RGB")
-    img = np.array(img).transpose(2, 0, 1)  # Convert to tensor format
-    img = torch.tensor(img, dtype=torch.float32).unsqueeze(0) / 255.0  # Normalize
+    img = np.array(img).transpose(2, 0, 1)  # Convert to [C, H, W]
+    img = torch.tensor(img, dtype=torch.float32).unsqueeze(0) / 255.0  # Normalize to [1, C, H, W]
 
     with torch.no_grad():
-        mask = model(img, clothing_classes=torch.arange(18))
+        mask = model(img.squeeze(-1), clothing_classes=torch.arange(18))  # ✅ Remove extra dimension
         mask = mask.squeeze().numpy()
 
     mask = (mask > 0.5).astype(np.uint8) * 255  # Convert to binary mask