Update networks.py

networks.py

@@ -6,6 +6,15 @@ from torchvision import models
 import os
 import numpy as np
 
+# Configuration class to hold all parameters
+class Options:
+    def __init__(self):
+        # Default values
+        self.fine_height = 256
+        self.fine_width = 192
+        self.grid_size = 3
+        self.use_dropout = False
+
 def weights_init_normal(m):
     classname = m.__class__.__name__
     if classname.find('Conv') != -1:
@@ -143,6 +152,7 @@ class TpsGridGen(nn.Module):
         super(TpsGridGen, self).__init__()
         self.out_h, self.out_w = out_h, out_w
         self.reg_factor = reg_factor
+        self.grid_size = grid_size
 
         # create grid in numpy
         self.grid = np.zeros([self.out_h, self.out_w, 3], dtype=np.float32)
@@ -173,7 +183,6 @@ class TpsGridGen(nn.Module):
     def forward(self, theta):
         warped_grid = self.apply_transformation(
             theta, torch.cat((self.grid_X, self.grid_Y), 3))
-
         return warped_grid
 
     def compute_L_inverse(self, X, Y):
@@ -285,10 +294,6 @@ class TpsGridGen(nn.Module):
 
         return torch.cat((points_X_prime, points_Y_prime), 3)
 
-# Defines the Unet generator.
-# |num_downs|: number of downsamplings in UNet. For example,
-# if |num_downs| == 7, image of size 128x128 will become of size 1x1
-# at the bottleneck
 class UnetGenerator(nn.Module):
     def __init__(self, input_nc, output_nc, num_downs, ngf=64,
                  norm_layer=nn.BatchNorm2d, use_dropout=False):
@@ -313,9 +318,6 @@ class UnetGenerator(nn.Module):
     def forward(self, input):
         return self.model(input)
 
-# Defines the submodule with skip connection.
-# X -------------------identity---------------------- X
-#   |-- downsampling -- |submodule| -- upsampling --|
 class UnetSkipConnectionBlock(nn.Module):
     def __init__(self, outer_nc, inner_nc, input_nc=None,
                  submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
@@ -464,9 +466,12 @@ class GicLoss(nn.Module):
 class GMM(nn.Module):
     """ Geometric Matching Module
     """
-
-    def __init__(self, opt):
+    def __init__(self, opt=None):
         super(GMM, self).__init__()
+        # Initialize default options if none provided
+        if opt is None:
+            opt = Options()
+            
         self.extractionA = FeatureExtraction(
             22, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d)
         self.extractionB = FeatureExtraction(