resolve crop-related issues

Author: torms3

deepem/models/layers.py

@@ -39,5 +39,6 @@ def __init__(self, cropsz):
     def forward(self, x):
         if self.cropsz is not None:
             for k, v in x.items():
-                x[k] = torch_utils.crop_center(v, self.cropsz)
+                cropsz = [int(v.shape[i]*self.cropsz[i]) for i in [-3,-2,-1]]
+                x[k] = torch_utils.crop_center(v, cropsz)
         return x

deepem/test/option.py

@@ -124,7 +124,8 @@ def parse(self):
         diff = np.array(opt.fov) - np.array(opt.outputsz)
         assert all(diff >= 0)
         if any(diff > 0):
-            opt.cropsz = opt.outputsz
+            # opt.cropsz = opt.outputsz
+            opt.cropsz = [o/float(f) for f,o in zip(opt.fov,opt.outputsz)]
         else:
             opt.cropsz = None
 

deepem/train/data.py

@@ -32,7 +32,8 @@ def __init__(self, opt, data, is_train=True, prob=None):
 
     def __call__(self):
         sample = next(self.dataiter)
-        sample = self.modifier(sample)
+        if self.is_train:
+            sample = self.modifier(sample)
         for k in sample:
             is_input = k in self.inputs
             sample[k].requires_grad_(is_input)

deepem/train/logger.py

@@ -2,6 +2,7 @@
 import sys
 import datetime
 from collections import OrderedDict
+import numpy as np
 
 import torch
 from torchvision.utils import make_grid
@@ -87,15 +88,22 @@ def flush(self):
             return ret
 
     def log_images(self, phase, iter_num, preds, sample):
+        # Peep output size
+        key = sorted(self.out_spec)[0]
+        cropsz = sample[key].shape[-3:]
+        for k in sorted(self.out_spec):
+            outsz = sample[k].shape[-3:]
+            assert np.array_equal(outsz, cropsz)
+
         # Inputs
         for k in sorted(self.in_spec):
             tag = '{}/images/{}'.format(phase, k)
             tensor = sample[k][0,...].cpu()
-            tensor = torch_utils.crop_center(tensor, self.outputsz)
+            tensor = torch_utils.crop_center(tensor, cropsz)
             self.log_image(tag, tensor, iter_num)
 
         # Outputs
-        for k  in sorted(self.out_spec):
+        for k in sorted(self.out_spec):
             if k == 'affinity':
                 # Prediction
                 tag = '{}/images/{}'.format(phase, k)

deepem/train/option.py

@@ -107,6 +107,9 @@ def initialize(self):
         self.parser.add_argument('--glia', type=float, default=0)  # Glia
         self.parser.add_argument('--glia_mask', action='store_true')
 
+        # Test training
+        self.parser.add_argument('--test', action='store_true')
+
         self.initialized = True
 
     def parse(self):
@@ -119,6 +122,8 @@ def parse(self):
             opt.exp_dir = opt.exp_name
         else:
             opt.exp_dir = 'experiments/{}'.format(opt.exp_name)
+        if opt.test:
+            opt.exp_dir = 'test/' + opt.exp_dir
         opt.log_dir = os.path.join(opt.exp_dir, 'logs')
         opt.model_dir = os.path.join(opt.exp_dir, 'models')
 
@@ -162,7 +167,8 @@ def parse(self):
         diff = np.array(opt.fov) - np.array(opt.outputsz)
         assert all(diff >= 0)
         if any(diff > 0):
-            opt.cropsz = opt.outputsz
+            # opt.cropsz = opt.outputsz
+            opt.cropsz = [o/float(f) for f,o in zip(opt.fov,opt.outputsz)]
         else:
             opt.cropsz = None
 
@@ -181,6 +187,13 @@ def parse(self):
             'glia':  ('glia', 1),
         }
 
+        # Test training
+        if opt.test:
+            opt.eval_intv = 100
+            opt.eval_iter = 10
+            opt.avgs_intv = 10
+            opt.imgs_intv = 100
+
         for k, v in class_dict.items():
             loss_w = args[k]
             if loss_w > 0:

deepem/train/run.py

@@ -35,9 +35,6 @@ def train(opt):
             # Load training samples.
             sample = train_loader()
 
-            # Crop sample.
-
-
             # Optimizer step
             optimizer.zero_grad()
             losses, nmasks, preds = forward(model, sample, opt)

deepem/utils/torch_utils.py

@@ -39,9 +39,6 @@ def crop_border(v, size):
 
 
 def crop_center(v, size):
-    # TODO: hack
-    if all([a <= b for a, b in zip(v.shape[-3:], size[-3:])]):
-        return v
     assert all([a >= b for a, b in zip(v.shape[-3:], size[-3:])])
     z, y, x = size[-3:]
     sx = (v.shape[-1] - x) // 2