Update networks.py

networks.py

@@ -195,60 +195,38 @@ class TpsGridGen(nn.Module):
         
         return torch.cat((points_X, points_Y), 3)
 
+    # In TpsGridGen class, replace transform_points method with this:
     def transform_points(self, grid, W, Q):
         batch_size, h, w, _ = grid.size()
         n_points = h * w
         
-        # 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)  # (N,2)
+        # Control points P (N, 2)
+        P = torch.cat([self.P_X_base, self.P_Y_base], 1)
         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.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=2, keepdim=False)  # (B, H*W, N)
+        # Compute U = r^2 * log(r^2)
+        grid_flat = grid.view(batch_size, n_points, 2)  # (B, H*W, 2)
+        dist = grid_flat.unsqueeze(2) - P.unsqueeze(1)  # (B, H*W, N, 2)
+        dist_squared = torch.sum(dist**2, dim=3)  # (B, H*W, N)
         dist_squared[dist_squared == 0] = 1  # Avoid log(0)
-        U = torch.mul(dist_squared, torch.log(dist_squared))
+        U = 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)  # (B, H*W, 3)
+        # Compute affine part [1, x, y]
+        ones = torch.ones(batch_size, n_points, 1, device=grid.device)
+        A = torch.cat([ones, grid_flat], dim=2)  # (B, H*W, 3)
         
-        # Reshape Q to include affine parameters
-        # Q has shape (B, N, 1) - we need to extract affine parameters
-        # Instead, we'll use the full Li matrix for the affine part
-        # This is a simplified approach that works for the forward pass
+        # Warp coefficients
+        W = W.view(batch_size, self.N, 1)
+        Q = Q.view(batch_size, self.N, 1)
         
-        # Compute affine component directly from Q
-        affine_x = Q[:, :, 0].mean(dim=1, keepdim=True)  # Simplified affine X
-        affine_y = Q[:, :, 0].mean(dim=1, keepdim=True)  # Simplified affine Y
-        affine = torch.cat([
-            torch.ones(batch_size, n_points, 1, device=grid.device),
-            grid_flat.view(batch_size, n_points, 1) * affine_x,
-            grid_flat.view(batch_size, n_points, 1) * affine_y
-        ], dim=2)
-        
-        # Compute non-affine component
+        # Non-affine part
         non_affine = torch.bmm(U, W)  # (B, H*W, 1)
         
-        # Combine components
-        points = affine[:, :, :1] + non_affine  # Only use the affine bias for X/Y
+        # Affine part
+        affine = torch.bmm(A, Q)  # (B, H*W, 1)
         
+        # Combine components
+        points = affine + non_affine
         return points.view(batch_size, h, w, 1)
 
 class GMM(nn.Module):