Added option to apply ww&wl

Thiago Franco de Moraes
1 parent 10ed1e32
Showing 2 changed files with 32 additions and 2 deletions Show diff stats
invesalius/gui/brain_seg_dialog.py
invesalius/segmentation/brain/segment.py
@@ -81,6 +81,8 @@ class BrainSegmenterDialog(wx.Dialog):
         self.txt_threshold.SetMinClientSize((w, -1))
         self.chk_new_mask = wx.CheckBox(self, wx.ID_ANY, _("Create new mask"))
         self.chk_new_mask.SetValue(True)
+        self.chk_apply_wwwl = wx.CheckBox(self, wx.ID_ANY, _("Apply WW&WL"))
+        self.chk_apply_wwwl.SetValue(False)
         self.progress = wx.Gauge(self, -1)
         self.lbl_progress_caption = wx.StaticText(self, -1, _("Elapsed time:"))
         self.lbl_time = wx.StaticText(self, -1, _("00:00:00"))
@@ -120,6 +122,7 @@ class BrainSegmenterDialog(wx.Dialog):
         )
         sizer_3.Add(self.txt_threshold, 0, wx.ALL, 5)
         main_sizer.Add(sizer_3, 0, wx.EXPAND, 0)
+        main_sizer.Add(self.chk_apply_wwwl, 0, wx.EXPAND | wx.ALL, 5)
         main_sizer.Add(self.chk_new_mask, 0, wx.EXPAND | wx.ALL, 5)
         main_sizer.Add(self.progress, 0, wx.EXPAND | wx.ALL, 5)
         time_sizer = wx.BoxSizer(wx.HORIZONTAL)
@@ -216,6 +219,7 @@ class BrainSegmenterDialog(wx.Dialog):
             device_id = self.plaidml_devices[self.cb_devices.GetValue()]
         except (KeyError, AttributeError):
             device_id = "llvm_cpu.0"
+        apply_wwwl = self.chk_apply_wwwl.GetValue()
         create_new_mask = self.chk_new_mask.GetValue()
         use_gpu = self.chk_use_gpu.GetValue()
         prob_threshold = self.sld_threshold.GetValue() / 100.0
@@ -224,8 +228,11 @@ class BrainSegmenterDialog(wx.Dialog):
         self.btn_segment.Disable()
         self.chk_new_mask.Disable()
+        window_width = slc.Slice().window_width
+        window_level = slc.Slice().window_level
+
         try:
-            self.ps = segment.SegmentProcess(image, create_new_mask, backend, device_id, use_gpu)
+            self.ps = segment.SegmentProcess(image, create_new_mask, backend, device_id, use_gpu, apply_wwwl, window_width, window_level)
             self.ps.start()
         except (multiprocessing.ProcessError, OSError, ValueError) as err:
             self.OnStop(None)
@@ -20,6 +20,17 @@ SIZE = 48
 OVERLAP = SIZE // 2 + 1
+def get_LUT_value(data, window, level):
+    shape = data.shape
+    data_ = data.ravel()
+    data = np.piecewise(data_,
+                        [data_ <= (level - 0.5 - (window-1)/2),
+                         data_ > (level - 0.5 + (window-1)/2)],
+                        [0, window, lambda data_: ((data_ - (level - 0.5))/(window-1) + 0.5)*(window)])
+    data.shape = shape
+    return data
+
+
 def gen_patches(image, patch_size, overlap):
     sz, sy, sx = image.shape
     i_cuts = list(
@@ -63,6 +74,8 @@ def brain_segment(image, probability_array, comm_array):
         model = keras.models.model_from_json(json_file.read())
     model.load_weights(str(folder.joinpath("model.h5")))
     model.compile("Adam", "binary_crossentropy")
+    
+    keras.utils.plot_model(model, "model.png", show_shapes=True, show_layer_names=True)
     image = imagedata_utils.image_normalize(image, 0.0, 1.0, output_dtype=np.float32)
     sums = np.zeros_like(image)
@@ -80,7 +93,7 @@ def brain_segment(image, probability_array, comm_array):
 ctx = multiprocessing.get_context('spawn')
 class SegmentProcess(ctx.Process):
-    def __init__(self, image, create_new_mask, backend, device_id, use_gpu):
+    def __init__(self, image, create_new_mask, backend, device_id, use_gpu, apply_wwwl=False, window_width=255, window_level=127):
         multiprocessing.Process.__init__(self)
         self._image_filename = image.filename
@@ -102,6 +115,10 @@ class SegmentProcess(ctx.Process):
         self.device_id = device_id
         self.use_gpu = use_gpu
+        self.apply_wwwl = apply_wwwl
+        self.window_width = window_width
+        self.window_level = window_level
+
         self._pconn, self._cconn = multiprocessing.Pipe()
         self._exception = None
@@ -122,6 +139,12 @@ class SegmentProcess(ctx.Process):
             shape=self._image_shape,
             mode="r",
         )
+
+        print(image.min(), image.max())
+        if self.apply_segment_threshold:
+            print("Applying window level")
+            image = get_LUT_value(image, self.window_width, self.window_level)
+
         probability_array = np.memmap(
             self._prob_array_filename,
             dtype=np.float32,