Update networks.py

networks.py

@@ -126,29 +126,33 @@ class TpsGridGen(nn.Module):
         self.out_h = out_h
         self.out_w = out_w
         self.grid_size = grid_size
-        
-        # Create grid
-        axis_coords = np.linspace(-1, 1, grid_size)
         self.N = grid_size * grid_size
+
+        # Create regular grid of control points
+        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))  # (N,1)
+        P_Y = torch.FloatTensor(P_Y.reshape(-1, 1))  # (N,1)
+        self.register_buffer('P_X', P_X)
+        self.register_buffer('P_Y', P_Y)
         
-        self.P_X_base = torch.FloatTensor(P_X.reshape(-1, 1))
-        self.P_Y_base = torch.FloatTensor(P_Y.reshape(-1, 1))
-        self.Li = self.compute_L_inverse(self.P_X_base, self.P_Y_base).unsqueeze(0)
+        # Compute inverse matrix L^-1
+        self.register_buffer('Li', self.compute_L_inverse(P_X, P_Y))
         
-        # Grid for interpolation
+        # Create sampling grid
         grid_X, grid_Y = np.meshgrid(np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h))
-        self.grid_X = torch.FloatTensor(grid_X).unsqueeze(0).unsqueeze(3)
-        self.grid_Y = torch.FloatTensor(grid_Y).unsqueeze(0).unsqueeze(3)
+        self.register_buffer('grid_X', torch.FloatTensor(grid_X).unsqueeze(0).unsqueeze(3))  # (1,H,W,1)
+        self.register_buffer('grid_Y', torch.FloatTensor(grid_Y).unsqueeze(0).unsqueeze(3))  # (1,H,W,1)
 
     def compute_L_inverse(self, X, Y):
-        N = X.size()[0]
+        N = X.size(0)
         Xmat = X.expand(N, N)
         Ymat = Y.expand(N, N)
         P_dist_squared = torch.pow(Xmat - Xmat.transpose(0, 1), 2) + torch.pow(Ymat - Ymat.transpose(0, 1), 2)
-        P_dist_squared[P_dist_squared == 0] = 1
+        P_dist_squared[P_dist_squared == 0] = 1  # Avoid log(0)
         K = torch.mul(P_dist_squared, torch.log(P_dist_squared))
         
+        # Construct L matrix
         O = torch.FloatTensor(N, 1).fill_(1)
         Z = torch.FloatTensor(3, 3).fill_(0)
         P = torch.cat((O, X, Y), 1)
@@ -156,43 +160,51 @@ class TpsGridGen(nn.Module):
         return torch.inverse(L)
 
     def forward(self, theta):
-        batch_size = theta.size()[0]
-        theta = theta.contiguous()
+        batch_size = theta.size(0)
+        device = theta.device
         
+        # Split theta into x and y components
         Q_X = theta[:, :self.N].contiguous().view(batch_size, self.N, 1)
         Q_Y = theta[:, self.N:].contiguous().view(batch_size, self.N, 1)
-        Q_X = Q_X + self.P_X_base.expand_as(Q_X)
-        Q_Y = Q_Y + self.P_Y_base.expand_as(Q_Y)
+        Q_X = Q_X + self.P_X.expand_as(Q_X)
+        Q_Y = Q_Y + self.P_Y.expand_as(Q_Y)
         
         # Compute weights
-        W_X = torch.bmm(self.Li[:, :self.N, :self.N].expand(batch_size, self.N, self.N), Q_X)
-        W_Y = torch.bmm(self.Li[:, :self.N, :self.N].expand(batch_size, self.N, self.N), Q_Y)
+        W_X = torch.bmm(self.Li[:, :self.N, :self.N].expand(batch_size, -1, -1), Q_X)
+        W_Y = torch.bmm(self.Li[:, :self.N, :self.N].expand(batch_size, -1, -1), Q_Y)
+        
+        # Repeat grid for batch processing
+        grid_X = self.grid_X.expand(batch_size, -1, -1, -1).to(device)
+        grid_Y = self.grid_Y.expand(batch_size, -1, -1, -1).to(device)
         
-        # Transform points
-        points_X = self.apply_transformation(self.grid_X, W_X, Q_X)
-        points_Y = self.apply_transformation(self.grid_Y, W_Y, Q_Y)
+        # Compute transformed coordinates
+        points_X = self.transform_points(grid_X, W_X, Q_X)
+        points_Y = self.transform_points(grid_Y, W_Y, Q_Y)
         
         return torch.cat((points_X, points_Y), 3)
 
-    def apply_transformation(self, grid, W, Q):
-        batch_size = W.size()[0]
-        P = torch.cat([
-            self.P_X_base.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(0, 4),
-            self.P_Y_base.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(0, 4)
-        ], 1)
+    def transform_points(self, grid, W, Q):
+        batch_size, h, w, _ = grid.size()
+        
+        # Compute distance between grid points and control points
+        grid_flat = grid.view(batch_size, -1, 1)
+        P = torch.cat([self.P_X, self.P_Y], 1).unsqueeze(0).expand(batch_size, -1, -1).to(grid.device)
+        delta = grid_flat - P
         
-        delta = grid.expand(batch_size, 1, self.out_h, self.out_w, 1, self.N) - P.expand(batch_size, 1, self.out_h, self.out_w, 1, self.N)
-        dist_squared = torch.pow(delta[:,0], 2) + torch.pow(delta[:,1], 2)
-        dist_squared[dist_squared == 0] = 1
+        # Compute U (radial basis function)
+        dist_squared = torch.sum(torch.pow(delta, 2), 2, keepdim=True)
+        dist_squared[dist_squared == 0] = 1  # Avoid log(0)
         U = torch.mul(dist_squared, torch.log(dist_squared))
         
-        points = torch.sum(torch.mul(W.expand(batch_size, 1, self.out_h, self.out_w, 1, self.N), U.unsqueeze(4)), 5)
-        points += torch.sum(Q.expand(batch_size, 1, self.out_h, self.out_w, 1, 3) * 
-                 torch.cat([grid.new_ones(batch_size, 1, self.out_h, self.out_w, 1),
-                           grid.expand(batch_size, 1, self.out_h, self.out_w, 1),
-                           grid.transpose(3,4).expand(batch_size, 1, self.out_h, self.out_w, 1)], 4), 5)
+        # Compute affine + non-affine transformation
+        A = torch.cat([
+            torch.ones(batch_size, h*w, 1, device=grid.device),
+            grid_flat[:, :, 0:1],
+            grid_flat[:, :, 1:2]
+        ], 2)
         
-        return points.squeeze(4)
+        points = torch.bmm(A, Q.view(batch_size, 3, -1)) + torch.bmm(U, W.view(batch_size, self.N, -1))
+        return points.view(batch_size, h, w, 1)
 
 class GMM(nn.Module):
     def __init__(self, opt=None):