shiftCorrectionWidget.py 13.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2016-2017 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/


__authors__ = ["J. Garriga"]
__license__ = "MIT"
29
__date__ = "05/07/2021"
30

31
# import os
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
32
import numpy
33

34
from silx.gui import qt
35
from silx.gui.colors import Colormap
36
37
from silx.gui.plot.StackView import StackViewMainWindow

38
import darfix
39
from darfix.core.dataset import Operation
40

Henri Payno's avatar
Henri Payno committed
41
from .operationThread import OperationThread
42
from .utils import ChooseDimensionDock
43

44

45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
class ShiftCorrectionDialog(qt.QDialog):
    """
    Dialog with `ShiftCorrectionWidget` as main window and standard buttons.
    """

    okSignal = qt.Signal()

    def __init__(self, parent=None):
        qt.QDialog.__init__(self, parent)
        self.setWindowFlags(qt.Qt.Widget)
        types = qt.QDialogButtonBox.Ok
        self._buttons = qt.QDialogButtonBox(parent=self)
        self._buttons.setStandardButtons(types)
        self._buttons.setEnabled(False)
        resetB = self._buttons.addButton(self._buttons.Reset)
        self.mainWindow = ShiftCorrectionWidget(parent=self)
        self.mainWindow.setAttribute(qt.Qt.WA_DeleteOnClose)
        self.setLayout(qt.QVBoxLayout())
        self.layout().addWidget(self.mainWindow)
        self.layout().addWidget(self._buttons)

        self._buttons.accepted.connect(self.okSignal.emit)
        resetB.clicked.connect(self.mainWindow.resetStack)
        self.mainWindow.computingSignal.connect(self._toggleButton)

    def setDataset(self, dataset, indices=None, bg_indices=None, bg_dataset=None):

        if dataset is not None:
            self._buttons.setEnabled(True)
            self.mainWindow.setDataset(dataset, indices, bg_indices, bg_dataset)

    def _toggleButton(self, state):
        self._buttons.button(qt.QDialogButtonBox.Ok).setEnabled(not state)

79

80
class ShiftCorrectionWidget(qt.QMainWindow):
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
81
82
83
    """
    A widget to apply shift correction to a stack of images
    """
84
    computingSignal = qt.Signal(bool)
85
86

    def __init__(self, parent=None):
87
88
89
        qt.QMainWindow.__init__(self, parent)

        self.setWindowFlags(qt.Qt.Widget)
90
91
        self._shift = numpy.array([0, 0])
        self._filtered_shift = None
92
        self._dimension = None
93
94
95
96
        self._update_dataset = None
        self.indices = None
        self.bg_indices = None
        self.bg_dataset = None
97

98
99
        self._inputDock = _InputDock()
        self._inputDock.widget.correctionB.setEnabled(False)
100

101
        self._sv = StackViewMainWindow()
102
        self._sv.setColormap(Colormap(name=darfix.config.DEFAULT_COLORMAP_NAME,
103
                                      normalization="linear"))
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        self.setCentralWidget(self._sv)
        self._chooseDimensionDock = ChooseDimensionDock(self)
        spacer1 = qt.QWidget(parent=self)
        spacer1.setLayout(qt.QVBoxLayout())
        spacer1.setSizePolicy(qt.QSizePolicy.Minimum, qt.QSizePolicy.Expanding)
        spacer2 = qt.QWidget(parent=self)
        spacer2.setLayout(qt.QVBoxLayout())
        spacer2.setSizePolicy(qt.QSizePolicy.Minimum, qt.QSizePolicy.Expanding)
        self._chooseDimensionDock.widget.layout().addWidget(spacer1)
        self._inputDock.widget.layout().addWidget(spacer2)
        self._chooseDimensionDock.hide()
        self.addDockWidget(qt.Qt.RightDockWidgetArea, self._chooseDimensionDock)
        self.addDockWidget(qt.Qt.RightDockWidgetArea, self._inputDock)

        self._inputDock.widget.correctionB.clicked.connect(self.correct)
119
        self._inputDock.widget.abortB.clicked.connect(self.abort)
120
121
122
        self._inputDock.widget._findShiftB.clicked.connect(self._findShift)
        self._chooseDimensionDock.widget.filterChanged.connect(self._filterStack)
        self._chooseDimensionDock.widget.stateDisabled.connect(self._wholeStack)
123

124
    def setDataset(self, dataset, indices=None, bg_indices=None, bg_dataset=None):
125
126
127
128
129
130
        """
        Dataset setter. Saves the dataset and updates the stack with the dataset
        data

        :param Dataset dataset: dataset
        """
131
        self.dataset = dataset
132
133
134
135
        self._update_dataset = dataset
        self.indices = indices
        self.bg_indices = bg_indices
        self.bg_dataset = bg_dataset
136
        self._inputDock.widget.correctionB.setEnabled(True)
137
        if len(self.dataset.data.shape) > 3:
138
            self._chooseDimensionDock.show()
139
            self._chooseDimensionDock.widget.setDimensions(self._update_dataset.dims)
140
        if not self._chooseDimensionDock.widget._checkbox.isChecked():
141
            self._wholeStack()
142
143

    def getDataset(self):
144
        return self._update_dataset, self.indices, self.bg_indices, self.bg_dataset
145

146
    def correct(self):
147
        """
Julia Garriga Ferrer's avatar
Docs    
Julia Garriga Ferrer committed
148
        Function that starts the thread to compute the shift given
149
150
        at the input widget
        """
151
152
        dx = self._inputDock.widget.getDx()
        dy = self._inputDock.widget.getDy()
153
154
155
156
157
158
159
160
        self.shift = numpy.array([dy, dx])
        if self._filtered_shift is None or self._inputDock.widget.checkbox.isChecked():
            frames = numpy.arange(self._update_dataset.get_data(indices=self.indices, dimension=self._dimension).shape[0])
            self.thread_correction = OperationThread(self, self._update_dataset.apply_shift)
            self.thread_correction.setArgs(numpy.outer(self.shift, frames), self._dimension, indices=self.indices)
        else:
            self.thread_correction = OperationThread(self, self._update_dataset.apply_shift_along_dimension)
            self.thread_correction.setArgs(self._filtered_shift, self._dimension[0], indices=self.indices)
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
161
        self.thread_correction.finished.connect(self._updateData)
162
        self._inputDock.widget.correctionB.setEnabled(False)
163
        self._inputDock.widget.abortB.show()
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
164
        self.thread_correction.start()
165
        self.computingSignal.emit(True)
166

167
168
169
170
    def abort(self):
        self._inputDock.widget.abortB.setEnabled(False)
        self._update_dataset.stop_operation(Operation.SHIFT)

171
172
173
    def resetStack(self):
        self._update_dataset = self.dataset
        self.setStack()
174

175
    def updateProgress(self, progress):
176
        self.sigProgressChanged.emit(progress)
177

178
    def _findShift(self):
179
180
181
182
183
184
        if self._filtered_shift is not None:
            self.thread_detection = OperationThread(self, self._update_dataset.find_shift_along_dimension)
            self.thread_detection.setArgs(self._dimension[0], indices=self.indices)
        else:
            self.thread_detection = OperationThread(self, self._update_dataset.find_shift)
            self.thread_detection.setArgs(self._dimension, indices=self.indices)
185
        self._inputDock.widget._findShiftB.setEnabled(False)
186
187
        self.thread_detection.finished.connect(self._updateShift)
        self.thread_detection.start()
188
        self.computingSignal.emit(True)
189

Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
190
    def _updateShift(self):
191
        self._inputDock.widget._findShiftB.setEnabled(True)
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
192
        self.thread_detection.finished.disconnect(self._updateShift)
193
194
195
196
197
        if self._filtered_shift is None:
            self.shift = numpy.round(self.thread_detection.data[:, 1], 5)
        else:
            self._filtered_shift = numpy.round(self.thread_detection.data[:, :, 1], 5)
            self.shift = self._filtered_shift[self._dimension[1][0]]
198
        self.computingSignal.emit(False)
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
199

200
    def _updateData(self):
201
202
203
        """
        Updates the stack with the data computed in the thread
        """
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
204
        self.thread_correction.finished.disconnect(self._updateData)
205
206
        self._inputDock.widget.abortB.hide()
        self._inputDock.widget.abortB.setEnabled(True)
207
        self._inputDock.widget.correctionB.setEnabled(True)
208
        self.computingSignal.emit(False)
209
210
211
        if self.thread_correction.data:
            self._update_dataset = self.thread_correction.data
            assert self._update_dataset is not None
212
            self.setStack(self._update_dataset)
213
214
        else:
            print("\nCorrection aborted")
215

216
    def setStack(self, dataset=None):
217
        """
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
218
        Sets new data to the stack.
219
        Mantains the current frame showed in the view.
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
220
221

        :param Dataset dataset: if not None, data set to the stack will be from the given dataset.
222
        """
223
224
        if dataset is None:
            dataset = self.dataset
225
        nframe = self._sv.getFrameNumber()
226
227
228
229
        if self.indices is None:
            self._sv.setStack(dataset.get_data() if dataset is not None else None)
        else:
            self._sv.setStack(dataset.get_data(self.indices) if dataset is not None else None)
230
        self._sv.setFrameNumber(nframe)
231

232
233
    def clearStack(self):
        self._sv.setStack(None)
234
        self._inputDock.widget.correctionB.setEnabled(False)
235

236
    def _filterStack(self, dim=0, val=0):
237
        self._dimension = [dim, val]
238

239
        data = self._update_dataset.get_data(self.indices, self._dimension)
240
241
242
        if self.dataset.dims.ndim == 2:
            stack_size = self.dataset.dims.get(dim[0]).size
            reset_shift = self._filtered_shift is None or self._filtered_shift.shape[0] != stack_size
243
            self._inputDock.widget.checkbox.show()
244
245
            self._filtered_shift = numpy.zeros((stack_size, 2)) if reset_shift else self._filtered_shift
            self.shift = self._filtered_shift[val[0]]
246
        if data.shape[0]:
247
            self._sv.setStack(data)
248
        else:
249
            self._sv.setStack(None)
250
251

    def _wholeStack(self):
252
        self._dimension = None
253
254
        self._filtered_shift = None
        self.shift = numpy.array([0, 0])
255
        self._inputDock.widget.checkbox.hide()
256
        self.setStack(self._update_dataset)
257

258
    def getStack(self):
259
260
261
262
263
        """
        Stack getter

        :returns: StackViewMainWindow:
        """
264
        return self._sv
265

266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
    def getStackViewColormap(self):
        """
        Returns the colormap from the stackView

        :rtype: silx.gui.colors.Colormap
        """
        return self._sv.getColormap()

    def setStackViewColormap(self, colormap):
        """
        Sets the stackView colormap

        :param colormap: Colormap to set
        :type colormap: silx.gui.colors.Colormap
        """
        self._sv.setColormap(colormap)

283
284
285
286
287
288
289
    @property
    def shift(self):
        return self._shift

    @shift.setter
    def shift(self, shift):
        self._shift = shift
290
291
        self._inputDock.widget.setDx(shift[1])
        self._inputDock.widget.setDy(shift[0])
292
293
294
295
296
297
298
299


class _InputDock(qt.QDockWidget):

    def __init__(self, parent=None):
        qt.QDockWidget.__init__(self, parent)
        self.widget = _InputWidget()
        self.setWidget(self.widget)
300

301
302
303
304
305
306
307
308

class _InputWidget(qt.QWidget):
    """
    Widget used to obtain the double parameters for the shift correction.
    """
    def __init__(self, parent=None):
        super(_InputWidget, self).__init__(parent)

309
310
311
        self._findShiftB = qt.QPushButton("Find shift")
        labelx = qt.QLabel("Horizontal shift:")
        labely = qt.QLabel("Vertical shift:")
312
313
        self.dxLE = qt.QLineEdit("0.0")
        self.dyLE = qt.QLineEdit("0.0")
314
        self.correctionB = qt.QPushButton("Correct")
315
316
        self.abortB = qt.QPushButton("Abort")
        self.abortB.hide()
317
318
        self.checkbox = qt.QCheckBox("Apply only to selected value")
        self.checkbox.setChecked(False)
319
        self.checkbox.hide()
320
321
322

        self.dxLE.setValidator(qt.QDoubleValidator())
        self.dyLE.setValidator(qt.QDoubleValidator())
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
323

324
        layout = qt.QGridLayout()
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
325

326
327
328
329
330
331
        layout.addWidget(self._findShiftB, 0, 0, 1, 2)
        layout.addWidget(labelx, 1, 0)
        layout.addWidget(labely, 2, 0)
        layout.addWidget(self.dxLE, 1, 1)
        layout.addWidget(self.dyLE, 2, 1)
        layout.addWidget(self.correctionB, 4, 0, 1, 2)
332
        layout.addWidget(self.abortB, 4, 0, 1, 2)
333
        layout.addWidget(self.checkbox, 3, 1)
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
334

335
336
        self.setLayout(layout)

337
338
339
340
341
342
    def setDx(self, dx):
        """
        Set the shift in the x axis
        """
        self.dxLE.setText(str(dx))

343
    def getDx(self):
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
344
345
        """
        Get the shift in the x axis
346

Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
347
348
        :return float:
        """
349
350
        return float(self.dxLE.text())

351
352
353
354
355
356
    def setDy(self, dy):
        """
        Set the shift in the x axis
        """
        self.dyLE.setText(str(dy))

357
    def getDy(self):
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
358
359
        """
        Get the shift in the y axis
360

Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
361
362
        :return float:
        """
363
        return float(self.dyLE.text())