Update networks.py

networks.py

@@ -203,30 +203,43 @@ class TpsGridGen(nn.Module):
         grid_flat = grid.view(batch_size, n_points, 1)
         
         # Prepare control points
-        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)
+        P = torch.cat([self.P_X_base, self.P_Y_base], 1)  # (N,2)
+        P = P.unsqueeze(0).expand(batch_size, -1, -1)  # (B, N, 2)
         
         # Compute distance between grid points and control points
         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
+        P_expanded = P.permute(0, 2, 1).unsqueeze(1)  # (B, 1, 2, N)
+        P_expanded = P_expanded.expand(-1, n_points, -1, -1)  # (B, H*W, 2, N)
+        
+        # Reshape grid for calculation
+        grid_reshaped = grid.view(batch_size, n_points, 1, 1).expand(-1, -1, -1, self.N)  # (B, H*W, 1, N)
+        
+        # Compute delta
+        delta = grid_reshaped - P_expanded
         
         # Compute U (radial basis function)
-        dist_squared = torch.sum(torch.pow(delta, 2), dim=1, keepdim=True)  # (B, H*W, 1)
+        dist_squared = torch.sum(torch.pow(delta, 2), dim=2, keepdim=False)  # (B, H*W, N)
         dist_squared[dist_squared == 0] = 1  # Avoid log(0)
         U = torch.mul(dist_squared, torch.log(dist_squared))
         
         # Compute affine transformation
+        # Create affine matrix [1, x, y]
         A = torch.cat([
             torch.ones(batch_size, n_points, 1, device=grid.device),
             grid_flat.view(batch_size, n_points, 1)
-        ], dim=2)
+        ], dim=2)  # (B, H*W, 3)
         
-        # Combine affine and non-affine components
-        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
+        # Extract affine parameters from Q
+        affine_params = Q.view(batch_size, 3, 1)  # (B, 3, 1)
         
+        # Compute affine component
+        affine = torch.bmm(A, affine_params)  # (B, H*W, 1)
+        
+        # Compute non-affine component
+        non_affine = torch.bmm(U, W)  # (B, H*W, 1)
+        
+        # Combine components
+        points = affine + non_affine
         return points.view(batch_size, h, w, 1)
 
 class GMM(nn.Module):
@@ -367,37 +380,40 @@ def load_checkpoint(model, checkpoint_path, strict=True):
     
     state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
     
-    # 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
-        
-        # 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)
+    # Create a new state dict that matches our model architecture
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        # Handle any name changes here if needed
+        new_key = key
+        if 'gridGen' in key:
+            # Map old parameter names to new ones
+            if 'P_X' in key:
+                new_key = key.replace('P_X', 'P_X_base')
+            elif 'P_Y' in key:
+                new_key = key.replace('P_Y', 'P_Y_base')
+        
+        # Only include keys that exist in the current model
+        if new_key in model.state_dict() and value.size() == model.state_dict()[new_key].size():
+            new_state_dict[new_key] = value
+    
+    # Add missing TPS parameters if needed
+    tps_params = ['gridGen.P_X_base', 'gridGen.P_Y_base', 'gridGen.Li', 
+                 'gridGen.grid_X', 'gridGen.grid_Y']
+    for param in tps_params:
+        if param not in new_state_dict and hasattr(model, 'gridGen'):
+            print(f"Initializing missing TPS parameter: {param}")
+            # Initialize with current model's value
+            new_state_dict[param] = model.state_dict()[param]
     
-    # Load state dict
-    model.load_state_dict(state_dict, strict=strict)
+    # Load the state dict
+    model.load_state_dict(new_state_dict, strict=strict)
     
     # Print warnings
     model_keys = set(model.state_dict().keys())
-    ckpt_keys = set(state_dict.keys())
-    missing = model_keys - ckpt_keys
-    unexpected = ckpt_keys - model_keys
+    loaded_keys = set(new_state_dict.keys())
+    
+    missing = model_keys - loaded_keys
+    unexpected = set(state_dict.keys()) - set(new_state_dict.keys())
     
     if missing:
         print(f"Missing keys: {sorted(missing)}")