resolve crop-related issues

Author: torms3

deepem/models/rsunet_deprecated.py

@@ -3,7 +3,7 @@
 
 import emvision
 
-from deepem.models.layers import Conv
+from deepem.models.layers import Conv, Crop
 
 
 def create_model(opt):
@@ -19,7 +19,7 @@ def create_model(opt):
     else:
         # Batch (instance) normalization
         core = emvision.models.RSUNet(width=width[:depth])
-    return Model(core, opt.in_spec, opt.out_spec, width[0])
+    return Model(core, opt.in_spec, opt.out_spec, width[0], cropsz=opt.cropsz)
 
 
 class InputBlock(nn.Sequential):
@@ -45,7 +45,7 @@ class Model(nn.Sequential):
     """
     Residual Symmetric U-Net.
     """
-    def __init__(self, core, in_spec, out_spec, out_channels):
+    def __init__(self, core, in_spec, out_spec, out_channels, cropsz=None):
         super(Model, self).__init__()
 
         assert len(in_spec)==1, "model takes a single input"
@@ -55,3 +55,5 @@ def __init__(self, core, in_spec, out_spec, out_channels):
         self.add_module('in', InputBlock(in_channels, out_channels, io_kernel))
         self.add_module('core', core)
         self.add_module('out', OutputBlock(out_channels, out_spec, io_kernel))
+        if cropsz is not None:
+            self.add_module('crop', Crop(cropsz))

deepem/models/updown_deprecated.py

@@ -3,7 +3,7 @@
 
 import emvision
 from emvision.models.layers import BilinearUp
-from deepem.models.layers import Conv
+from deepem.models.layers import Conv, Crop
 
 
 def create_model(opt):
@@ -19,7 +19,7 @@ def create_model(opt):
     else:
         # Batch (instance) normalization
         core = emvision.models.RSUNet(width=width[:depth])
-    return Model(core, opt.in_spec, opt.out_spec, width[0])
+    return Model(core, opt.in_spec, opt.out_spec, width[0], cropsz=opt.cropsz)
 
 
 class InputBlock(nn.Sequential):
@@ -47,7 +47,7 @@ class Model(nn.Sequential):
     """
     Residual Symmetric U-Net with down/upsampling in/output.
     """
-    def __init__(self, core, in_spec, out_spec, out_channels):
+    def __init__(self, core, in_spec, out_spec, out_channels, cropsz=None):
         super(Model, self).__init__()
 
         assert len(in_spec)==1, "model takes a single input"
@@ -57,3 +57,5 @@ def __init__(self, core, in_spec, out_spec, out_channels):
         self.add_module('in', InputBlock(in_channels, out_channels, io_kernel))
         self.add_module('core', core)
         self.add_module('out', OutputBlock(out_channels, out_spec, io_kernel))
+        if cropsz is not None:
+            self.add_module('crop', Crop(cropsz))

deepem/test/model.py

@@ -18,7 +18,7 @@ def __init__(self, model, opt):
         self.in_spec = dict(opt.in_spec)
         self.scan_spec = dict(opt.scan_spec)
         self.pretrain = opt.pretrain
-        self.cropsz = opt.cropsz
+        self.force_crop = opt.force_crop
 
         # Softer softmax
         if opt.temperature is None:
@@ -62,8 +62,8 @@ def forward(self, sample):
                 outputs[k] *= self.mask[k]
 
             # Crop outputs.
-            if self.cropsz is not None:
-                outputs[k] = torch_utils.crop_border(outputs[k], self.cropsz)
+            if self.force_crop is not None:
+                outputs[k] = torch_utils.crop_border(outputs[k], self.force_crop)
 
         return outputs
 

deepem/test/option.py

@@ -2,6 +2,7 @@
 import json
 import math
 import os
+import numpy as np
 
 from deepem.utils.py_utils import vec3, vec3f
 
@@ -28,7 +29,7 @@ def initialize(self):
         self.parser.add_argument('--no_eval', action='store_true')
         self.parser.add_argument('--inputsz', type=vec3, default=None)
         self.parser.add_argument('--outputsz', type=vec3, default=None)
-        self.parser.add_argument('--cropsz', type=vec3, default=None)
+        self.parser.add_argument('--force_crop', type=vec3, default=None)
         self.parser.add_argument('--fov', type=vec3, default=None)
         self.parser.add_argument('--depth', type=int, default=4)
         self.parser.add_argument('--width', type=int, default=None, nargs='+')
@@ -115,9 +116,17 @@ def parse(self):
         opt.fov = tuple(opt.fov)
         opt.inputsz = opt.fov if opt.inputsz is None else opt.inputsz
         opt.outputsz = opt.fov if opt.outputsz is None else opt.outputsz
-        opt.cropsz = opt.cropsz
         opt.in_spec = dict(input=(1,) + opt.inputsz)
         opt.out_spec = dict()
+
+        # Crop output
+        diff = np.array(opt.fov) - np.array(opt.outputsz)
+        assert all(diff >= 0)
+        if any(diff > 0):
+            opt.cropsz = opt.outputsz
+        else:
+            opt.cropsz = None
+
         if opt.aff:
             opt.out_spec['affinity'] = (3,) + opt.outputsz
         if opt.bdr:
@@ -154,6 +163,8 @@ def parse(self):
             opt.scan_spec['blood_vessel'] = (opt.blv_num_channels,) + opt.outputsz
         if opt.glia:
             opt.scan_spec['glia'] = (1,) + opt.outputsz
+
+        # Overlap & stride
         opt.overlap = self.get_overlap(opt.outputsz, opt.overlap)
         opt.stride = tuple(int(f-o) for f,o in zip(opt.outputsz, opt.overlap))
         opt.scan_params = dict(stride=opt.stride, blend=opt.blend)

deepem/train/data.py

@@ -59,7 +59,7 @@ def build(self, opt, data, is_train, prob):
         self.modifier = lambda x: x
         if opt.modifier is not None:
             mod = imp.load_source('modifier', opt.modifier)
-            self.modifier = mod.Modifier()
+            self.modifier = mod.Modifier(**opt.mod_params)
 
         # Data loader
         size = (opt.max_iter - opt.chkpt_num) * opt.batch_size