Update networks.py

networks.py

@@ -182,12 +182,12 @@ class TpsGridGen(nn.Module):
         Li_block = self.Li[:self.N, :self.N]
         
         # Compute weights
-        W_X = torch.bmm(Li_block.expand(batch_size, self.N, self.N), Q_X)
-        W_Y = torch.bmm(Li_block.expand(batch_size, self.N, self.N), Q_Y)
+        W_X = torch.bmm(Li_block.expand(batch_size, -1, -1), Q_X)
+        W_Y = torch.bmm(Li_block.expand(batch_size, -1, -1), Q_Y)
         
         # Prepare grid tensors
-        grid_X = self.grid_X.expand(batch_size, self.out_h, self.out_w, 1).to(device)
-        grid_Y = self.grid_Y.expand(batch_size, self.out_h, self.out_w, 1).to(device)
+        grid_X = self.grid_X.expand(batch_size, -1, -1, -1)
+        grid_Y = self.grid_Y.expand(batch_size, -1, -1, -1)
         
         # Compute transformed coordinates
         points_X = self.transform_points(grid_X, W_X, Q_X)
@@ -197,30 +197,36 @@ class TpsGridGen(nn.Module):
 
     def transform_points(self, grid, W, Q):
         batch_size, h, w, _ = grid.size()
+        n_points = h * w
         
-        # Flatten grid to (batch_size, H*W, 2)
-        grid_flat = grid.view(batch_size, -1, 1)
+        # Flatten grid to (batch_size, H*W, 1)
+        grid_flat = grid.view(batch_size, n_points, 1)
         
         # Prepare control points
-        P = torch.cat([self.P_X_base, self.P_Y_base], 1).expand(batch_size, -1, -1).to(grid.device)
+        P = torch.cat([self.P_X_base, self.P_Y_base], 1).t().unsqueeze(0)  # (1, 2, N)
+        P = P.expand(batch_size, -1, -1)  # (B, 2, N)
         
         # Compute distance between grid points and control points
-        delta = grid_flat - P
+        grid_expanded = grid_flat.expand(-1, -1, self.N)  # (B, H*W, N)
+        P_expanded = P.expand(n_points, -1, -1).permute(1, 0, 2)  # (B, H*W, N)
+        delta = grid_expanded - P_expanded
         
         # Compute U (radial basis function)
-        dist_squared = torch.sum(torch.pow(delta, 2), 2, keepdim=True)
+        dist_squared = torch.sum(torch.pow(delta, 2), dim=1, keepdim=True)  # (B, H*W, 1)
         dist_squared[dist_squared == 0] = 1  # Avoid log(0)
         U = torch.mul(dist_squared, torch.log(dist_squared))
         
         # Compute affine transformation
         A = torch.cat([
-            torch.ones(batch_size, h*w, 1, device=grid.device),
-            grid_flat[:, :, 0:1],
-            grid_flat[:, :, 1:2]
-        ], 2)
+            torch.ones(batch_size, n_points, 1, device=grid.device),
+            grid_flat.view(batch_size, n_points, 1)
+        ], dim=2)
         
         # Combine affine and non-affine components
-        points = torch.bmm(A, Q.view(batch_size, 3, 1)) + torch.bmm(U, W)
+        affine = torch.bmm(A, Q.view(batch_size, 1, 3).permute(0, 2, 1))
+        non_affine = torch.bmm(U.permute(0, 2, 1), W).permute(0, 2, 1)
+        points = affine + non_affine
+        
         return points.view(batch_size, h, w, 1)
 
 class GMM(nn.Module):
@@ -361,36 +367,37 @@ def load_checkpoint(model, checkpoint_path, strict=True):
     
     state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
     
-    # Create mapping for old buffer names to new names
-    buffer_mapping = {
-        'gridGen.P_X': 'gridGen.P_X_base',
-        'gridGen.P_Y': 'gridGen.P_Y_base',
-        'gridGen.Li': 'gridGen.Li',
-        'gridGen.grid_X': 'gridGen.grid_X',
-        'gridGen.grid_Y': 'gridGen.grid_Y'
-    }
-    
-    # Update state_dict keys
-    updated_state_dict = {}
-    for key, value in state_dict.items():
-        # Handle buffer name changes
-        for old_name, new_name in buffer_mapping.items():
-            if key.startswith(old_name):
-                key = key.replace(old_name, new_name)
+    # Initialize TPS grid parameters if missing
+    if 'gridGen.P_X_base' not in state_dict:
+        print("Initializing TPS grid parameters...")
+        grid_size = model.gridGen.grid_size
+        axis_coords = np.linspace(-1, 1, grid_size)
+        P_Y, P_X = np.meshgrid(axis_coords, axis_coords)
+        P_X = torch.FloatTensor(P_X.reshape(-1, 1))
+        P_Y = torch.FloatTensor(P_Y.reshape(-1, 1))
+        state_dict['gridGen.P_X_base'] = P_X
+        state_dict['gridGen.P_Y_base'] = P_Y
         
-        # Only keep keys that match current model
-        if key in model.state_dict() and value.size() == model.state_dict()[key].size():
-            updated_state_dict[key] = value
+        # Compute Li
+        Li = model.gridGen.compute_L_inverse(P_X, P_Y)
+        state_dict['gridGen.Li'] = Li
+        
+        # Create grid
+        grid_X, grid_Y = np.meshgrid(
+            np.linspace(-1, 1, model.gridGen.out_w),
+            np.linspace(-1, 1, model.gridGen.out_h)
+        )
+        state_dict['gridGen.grid_X'] = torch.FloatTensor(grid_X).unsqueeze(0).unsqueeze(3)
+        state_dict['gridGen.grid_Y'] = torch.FloatTensor(grid_Y).unsqueeze(0).unsqueeze(3)
     
-    # Load the updated state dict
-    model.load_state_dict(updated_state_dict, strict=strict)
+    # Load state dict
+    model.load_state_dict(state_dict, strict=strict)
     
-    # Print warnings about missing keys
+    # Print warnings
     model_keys = set(model.state_dict().keys())
-    checkpoint_keys = set(state_dict.keys())
-    
-    missing = model_keys - set(updated_state_dict.keys())
-    unexpected = checkpoint_keys - set(updated_state_dict.keys())
+    ckpt_keys = set(state_dict.keys())
+    missing = model_keys - ckpt_keys
+    unexpected = ckpt_keys - model_keys
     
     if missing:
         print(f"Missing keys: {sorted(missing)}")