Update app.py

app.py

@@ -7,50 +7,82 @@ from PIL import Image
 import cv2
 
 class TransformNet(nn.Module):
-    def __init__(self, input_dim=6):  
+    """Transformation Network for PointCloud Encoding"""
+    def __init__(self, input_dim=6):  # ✅ Ensure input has 6 channels
         super(TransformNet, self).__init__()
         self.conv1 = nn.Sequential(
-            nn.Conv2d(input_dim, 64, kernel_size=1, stride=1),  
+            nn.Conv2d(input_dim, 64, kernel_size=(1, 1)),  # ✅ Conv2d (Matches checkpoint)
             nn.BatchNorm2d(64),
             nn.ReLU()
         )
         self.conv2 = nn.Sequential(
-            nn.Conv2d(64, 128, kernel_size=1, stride=1),  
+            nn.Conv2d(64, 128, kernel_size=(1, 1)),  # ✅ Conv2d (Matches checkpoint)
             nn.BatchNorm2d(128),
             nn.ReLU()
         )
         self.conv3 = nn.Sequential(
-            nn.Conv2d(128, 1024, kernel_size=1, stride=1),  
-            nn.BatchNorm2d(1024),
+            nn.Conv1d(128, 1024, kernel_size=1),  # ✅ Conv1d to match `[1024, 128, 1]`
+            nn.BatchNorm1d(1024),
             nn.ReLU()
         )
         self.fc = nn.Linear(1024, 512)
 
     def forward(self, x):
-        x = x.unsqueeze(-1)  
+        x = x.unsqueeze(-1)  # ✅ Add extra dimension for Conv2d
         x = self.conv1(x)
         x = self.conv2(x)
-        x = self.conv3(x)
-        return self.fc(x.max(dim=-1)[0].squeeze(-1))
+        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
+
+class TransformNet(nn.Module):
+    """Transformation Network for PointCloud Encoding"""
+    def __init__(self, input_dim=6):  # ✅ Ensure input has 6 channels
+        super(TransformNet, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(input_dim, 64, kernel_size=(1, 1)),  # ✅ Conv2d (Matches checkpoint)
+            nn.BatchNorm2d(64),
+            nn.ReLU()
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(64, 128, kernel_size=(1, 1)),  # ✅ Conv2d (Matches checkpoint)
+            nn.BatchNorm2d(128),
+            nn.ReLU()
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv1d(128, 1024, kernel_size=1),  # ✅ Conv1d to match `[1024, 128, 1]`
+            nn.BatchNorm1d(1024),
+            nn.ReLU()
+        )
+        self.fc = nn.Linear(1024, 512)
+
+    def forward(self, x):
+        x = x.unsqueeze(-1)  # ✅ Add extra dimension for Conv2d
+        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
 
 class PointCloudEncoder(nn.Module):
+    """Point Cloud Encoder (pc_enc)"""
     def __init__(self):
         super(PointCloudEncoder, self).__init__()
         self.transform_net = TransformNet()
         self.convs = nn.ModuleList([
             nn.Sequential(
-                nn.Conv2d(512, 256, kernel_size=1, stride=1),  
+                nn.Conv2d(512, 256, kernel_size=(1, 1)),  # ✅ Conv2d (Matches checkpoint)
                 nn.BatchNorm2d(256),
                 nn.ReLU()
             ),
             nn.Sequential(
-                nn.Conv2d(256, 128, kernel_size=1, stride=1),  
+                nn.Conv2d(256, 128, kernel_size=(1, 1)),  # ✅ Conv2d (Matches checkpoint)
                 nn.BatchNorm2d(128),
                 nn.ReLU()
             ),
             nn.Sequential(
-                nn.Conv2d(128, 64, kernel_size=1, stride=1),  
-                nn.BatchNorm2d(64),
+                nn.Conv1d(128, 64, kernel_size=1),  # ✅ Conv1d to match `[64, 128, 1]`
+                nn.BatchNorm1d(64),
                 nn.ReLU()
             )
         ])
@@ -58,9 +90,13 @@ class PointCloudEncoder(nn.Module):
 
     def forward(self, x):
         x = self.transform_net(x)
-        for conv in self.convs:
-            x = conv(x)
-        return self.lin_global(x.max(dim=-1)[0].squeeze(-1))  
+        for i, conv in enumerate(self.convs):
+            if i < 2:
+                x = conv(x)  # ✅ Conv2d
+            else:
+                x = x.squeeze(-1)  # ✅ Remove extra dimension before Conv1d
+                x = conv(x)  # ✅ Conv1d (Matches `[64, 128, 1]`)
+        return self.lin_global(x.max(dim=-1)[0])  # ✅ Fix pooling
 
 class GarmentEncoder(nn.Module):
     """Garment Feature Encoder (garm_enc)"""
@@ -110,7 +146,7 @@ class CloseNet(nn.Module):
         return self.segm_dec(features)
 
 # Load Pretrained Model
-model_path = "model_arch/closenet.pth"
+model_path = "/content/closenet.pth"
 model = CloseNet()
 model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")), strict=False)
 model.eval()