fix: input to the CNN could only be 256x256

javiribera · Aug 10, 2019 · f860f7c · f860f7c
1 parent 8316f72
commit f860f7c
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/object-locator/models/unet_model.py b/object-locator/models/unet_model.py
@@ -59,10 +59,15 @@ def __init__(self, n_channels, n_classes,
 
         self.known_n_points = known_n_points
         if known_n_points is None:
-            self.branch_1 = nn.Sequential(nn.Linear(512, 64),
+            height_mid_features = height//(2**8)
+            width_mid_features = width//(2**8)
+            self.branch_1 = nn.Sequential(nn.Linear(height_mid_features*\
+                                                    width_mid_features*\
+                                                    512,
+                                                    64),
                                           nn.ReLU(inplace=True),
                                           nn.Dropout(p=0.5))
-            self.branch_2 = nn.Sequential(nn.Linear(256*256, 64),
+            self.branch_2 = nn.Sequential(nn.Linear(height*width, 64),
                                           nn.ReLU(inplace=True),
                                           nn.Dropout(p=0.5))
             self.regressor = nn.Sequential(nn.Linear(64 + 64, 1),