Update networks.py

networks.py

@@ -8,46 +8,72 @@ import numpy as np
 
 class Options:
     def __init__(self):
-        # Default values
+        # Image dimensions
         self.fine_height = 256
         self.fine_width = 192
+        
+        # GMM parameters
         self.grid_size = 5
+        self.input_nc = 22  # For extractionA
+        self.input_nc_B = 1  # For extractionB
+        
+        # TOM parameters
+        self.tom_input_nc = 26  # 3(agnostic) + 3(warped) + 1(mask) + 19(features)
+        self.tom_output_nc = 4   # 3(rendered) + 1(composite mask)
+        
+        # Training settings
         self.use_dropout = False
-        self.input_nc = 22
-        self.input_nc_B = 1  
-        self.tom_input_nc = 26  
-        self.tom_output_nc = 4  
+        self.norm_layer = nn.BatchNorm2d
 
 def weights_init_normal(m):
     classname = m.__class__.__name__
     if classname.find('Conv') != -1:
         init.normal_(m.weight.data, 0.0, 0.02)
     elif classname.find('Linear') != -1:
-        init.normal(m.weight.data, 0.0, 0.02)
-    elif classname.find('BatchNorm2d') != -1:
+        init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
         init.normal_(m.weight.data, 1.0, 0.02)
         init.constant_(m.bias.data, 0.0)
 
 def init_weights(net, init_type='normal'):
-    print('initialization method [%s]' % init_type)
+    print(f'initialization method [{init_type}]')
     net.apply(weights_init_normal)
 
 class FeatureExtraction(nn.Module):
-    def __init__(self, input_nc, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_dropout=False):
+    def __init__(self, input_nc, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
         super(FeatureExtraction, self).__init__()
-        downconv = nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1)
-        model = [downconv, nn.ReLU(True), norm_layer(ngf)]
+        
+        # Build feature extraction layers
+        layers = [
+            nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(True),
+            norm_layer(ngf)
+        ]
+        
         for i in range(n_layers):
-            in_ngf = 2**i * ngf if 2**i * ngf < 512 else 512
-            out_ngf = 2**(i+1) * ngf if 2**i * ngf < 512 else 512
-            downconv = nn.Conv2d(in_ngf, out_ngf, kernel_size=4, stride=2, padding=1)
-            model += [downconv, nn.ReLU(True), norm_layer(out_ngf)]
-        model += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(True)]
-        model += [norm_layer(512)]
-        model += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(True)]
-        self.model = nn.Sequential(*model)
+            in_channels = min(2**i * ngf, 512)
+            out_channels = min(2**(i+1) * ngf, 512)
+            layers += [
+                nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1),
+                nn.ReLU(True),
+                norm_layer(out_channels)
+            ]
+        
+        # Final processing blocks
+        layers += [
+            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(True),
+            norm_layer(512),
+            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(True)
+        ]
+        
+        self.model = nn.Sequential(*layers)
         init_weights(self.model)
 
+    def forward(self, x):
+        return self.model(x)
+
 class FeatureL2Norm(nn.Module):
     def __init__(self):
         super(FeatureL2Norm, self).__init__()
@@ -83,7 +109,7 @@ class FeatureRegression(nn.Module):
             nn.ReLU(inplace=True),
             nn.Conv2d(128, 64, kernel_size=3, padding=1),
             nn.BatchNorm2d(64),
-            nn.ReLU(inplace=True),
+            nn.ReLU(inplace=True)
         )
         self.linear = nn.Linear(64 * 4 * 3, output_dim)
         self.tanh = nn.Tanh()
@@ -97,18 +123,18 @@ class FeatureRegression(nn.Module):
 class TpsGridGen(nn.Module):
     def __init__(self, out_h=256, out_w=192, grid_size=5):
         super(TpsGridGen, self).__init__()
-        self.out_h, self.out_w = out_h, out_w
+        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
         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))
-        self.P_X_base = P_X.clone()
-        self.P_Y_base = P_Y.clone()
-        self.Li = self.compute_L_inverse(P_X, P_Y).unsqueeze(0)
+        
+        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)
         
         # Grid for interpolation
         grid_X, grid_Y = np.meshgrid(np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h))
@@ -117,10 +143,12 @@ class TpsGridGen(nn.Module):
 
     def compute_L_inverse(self, X, Y):
         N = X.size()[0]
-        Xmat, Ymat = X.expand(N, N), Y.expand(N, N)
-        P_dist_squared = torch.pow(Xmat-Xmat.transpose(0, 1), 2) + torch.pow(Ymat-Ymat.transpose(0, 1), 2)
+        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
         K = torch.mul(P_dist_squared, torch.log(P_dist_squared))
+        
         O = torch.FloatTensor(N, 1).fill_(1)
         Z = torch.FloatTensor(3, 3).fill_(0)
         P = torch.cat((O, X, Y), 1)
@@ -128,22 +156,44 @@ class TpsGridGen(nn.Module):
         return torch.inverse(L)
 
     def forward(self, theta):
-        theta = theta.contiguous()
         batch_size = theta.size()[0]
+        theta = theta.contiguous()
         
-        # Split theta into point coordinates
         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)
         
         # Compute weights
-        W_X, W_Y = self.apply_theta(Q_X, Q_Y)
+        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)
+        
+        # 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)
         
-        # Calculate transformed grid
-        points_X, points_Y = self.transform_points(W_X, W_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)
+        
+        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
+        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)
+        
+        return points.squeeze(4)
+
 class GMM(nn.Module):
     def __init__(self, opt=None):
         super(GMM, self).__init__()
@@ -167,57 +217,49 @@ class GMM(nn.Module):
         grid = self.gridGen(theta)
         return grid, theta
 
-class UnetGenerator(nn.Module):
-    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.InstanceNorm2d):
-        super(UnetGenerator, self).__init__()
-        unet_block = UnetSkipConnectionBlock(
-            ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
-        
-        for _ in range(num_downs - 5):
-            unet_block = UnetSkipConnectionBlock(
-                ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-                
-        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        
-        self.model = UnetSkipConnectionBlock(
-            output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
-
-    def forward(self, input):
-        return self.model(input)
-
 class UnetSkipConnectionBlock(nn.Module):
-    def __init__(self, outer_nc, inner_nc, input_nc=None, submodule=None, 
-                 outermost=False, innermost=False, norm_layer=nn.InstanceNorm2d):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, 
+                 norm_layer=nn.InstanceNorm2d, use_dropout=False):
         super(UnetSkipConnectionBlock, self).__init__()
         self.outermost = outermost
         use_bias = norm_layer == nn.InstanceNorm2d
-        
+
         if input_nc is None:
             input_nc = outer_nc
             
-        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
         downrelu = nn.LeakyReLU(0.2, True)
         downnorm = norm_layer(inner_nc)
         uprelu = nn.ReLU(True)
         upnorm = norm_layer(outer_nc)
 
         if outermost:
-            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1)
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
             down = [downconv]
             up = [uprelu, upconv, nn.Tanh()]
             model = down + [submodule] + up
         elif innermost:
-            upconv = nn.ConvTranspose2d(inner_nc, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
             down = [downrelu, downconv]
             up = [uprelu, upconv, upnorm]
             model = down + up
         else:
-            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
             down = [downrelu, downconv, downnorm]
             up = [uprelu, upconv, upnorm]
-            model = down + [submodule] + up
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
 
         self.model = nn.Sequential(*model)
 
@@ -227,17 +269,47 @@ class UnetSkipConnectionBlock(nn.Module):
         else:
             return torch.cat([x, self.model(x)], 1)
 
+class UnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.InstanceNorm2d, use_dropout=False):
+        super(UnetGenerator, self).__init__()
+        
+        # Build UNet structure
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 8, ngf * 8, input_nc=None, submodule=None,
+            norm_layer=norm_layer, innermost=True)
+        
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(
+                ngf * 8, ngf * 8, input_nc=None, submodule=unet_block,
+                norm_layer=norm_layer, use_dropout=use_dropout)
+                
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 4, ngf * 8, input_nc=None, submodule=unet_block,
+            norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(
+            ngf * 2, ngf * 4, input_nc=None, submodule=unet_block,
+            norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(
+            ngf, ngf * 2, input_nc=None, submodule=unet_block,
+            norm_layer=norm_layer)
+            
+        self.model = UnetSkipConnectionBlock(
+            output_nc, ngf, input_nc=input_nc, submodule=unet_block,
+            outermost=True, norm_layer=norm_layer)
+
+    def forward(self, input):
+        return self.model(input)
+
 class TOM(nn.Module):
-    """ Try-On Module """
     def __init__(self, opt=None):
         super(TOM, self).__init__()
         if opt is None:
             opt = Options()
         
-        # Input: [agnostic(3) + warped_design(3) + warped_mask(1) + features(19)] = 26 channels
         self.unet = UnetGenerator(
             input_nc=opt.tom_input_nc,
-            output_nc=opt.tom_output_nc,  # [rendered(3) + mask(1)]
+            output_nc=opt.tom_output_nc,
             num_downs=6,
             norm_layer=nn.InstanceNorm2d
         )
@@ -259,4 +331,24 @@ def load_checkpoint(model, checkpoint_path, strict=True):
         raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
     
     state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
-    model.load_state_dict(state_dict, strict=strict)
+    
+    # Filter out unexpected keys
+    model_state_dict = model.state_dict()
+    filtered_state_dict = {k: v for k, v in state_dict.items()
+                         if k in model_state_dict and v.size() == model_state_dict[k].size()}
+    
+    # Load filtered state dict
+    model.load_state_dict(filtered_state_dict, strict=strict)
+    
+    # Print warnings
+    missing = [k for k in model_state_dict if k not in state_dict]
+    unexpected = [k for k in state_dict if k not in model_state_dict]
+    size_mismatch = [k for k in state_dict 
+                    if k in model_state_dict and state_dict[k].size() != model_state_dict[k].size()]
+    
+    if missing:
+        print(f"Missing keys: {missing}")
+    if unexpected:
+        print(f"Unexpected keys: {unexpected}")
+    if size_mismatch:
+        print(f"Size mismatch: {size_mismatch}")