forked from IPKM/nmreval
95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
import numpy as np
|
|
import pyqtgraph as pg
|
|
|
|
from scipy.interpolate import splrep, splev
|
|
|
|
from ...Qt import QtCore, QtWidgets
|
|
from ..._py.baseline_dialog import Ui_SignalEdit
|
|
|
|
|
|
class QBaselineDialog(QtWidgets.QDialog, Ui_SignalEdit):
|
|
finished = QtCore.pyqtSignal(str, tuple)
|
|
|
|
def __init__(self, parent=None):
|
|
super().__init__(parent=parent)
|
|
self.setupUi(self)
|
|
|
|
self.data = None
|
|
|
|
self.graph = pg.PlotDataItem(x=[], y=[], pen=pg.mkPen({'color': 'b'}), name='Original')
|
|
self.graph_corr = pg.PlotDataItem(x=[], y=[], pen=pg.mkPen({'color': 'r'}), name='Corrected')
|
|
self.baseline = pg.PlotDataItem(x=[], y=[], name='Baseline')
|
|
|
|
self.anchors = []
|
|
self.anchor_lines = []
|
|
self.spline = None
|
|
|
|
self.legend = self.graphicsView.addLegend()
|
|
|
|
self.graphicsView.scene().sigMouseClicked.connect(self.add_node)
|
|
self.graphicsView.addItem(self.graph_corr)
|
|
self.graphicsView.addItem(self.graph)
|
|
self.graphicsView.addItem(self.baseline)
|
|
|
|
def add_data(self, x, y):
|
|
if self.data is not None:
|
|
QtWidgets.QMessageBox().information(self, 'Invalid number of datasets',
|
|
'Baseline correction is only working on one set at a time.')
|
|
self.close()
|
|
self.anchors.extend([np.min(x), np.max(x)])
|
|
self.data = (x, y)
|
|
self.graph.setData(x=x, y=y.real)
|
|
self.graph_corr.setData(x=x, y=y.real)
|
|
|
|
def accept(self):
|
|
self.finished.emit('bls', (splev(self.data[0], self.spline),))
|
|
self.close()
|
|
|
|
def add_node(self, evt):
|
|
vb = self.graphicsView.plotItem.vb
|
|
|
|
if self.graphicsView.plotItem.sceneBoundingRect().contains(evt.scenePos()) and evt.button() == 1:
|
|
pos = vb.mapSceneToView(evt.scenePos())
|
|
x = pos.x()
|
|
|
|
self.anchors.append(x)
|
|
self.anchors.sort()
|
|
row = self.anchors.index(x)
|
|
self.listWidget.insertItem(row-1, QtWidgets.QListWidgetItem(str(x)))
|
|
|
|
inf_line = pg.InfiniteLine(pos=x)
|
|
self.anchor_lines.insert(row-1, inf_line)
|
|
self.graphicsView.addItem(inf_line)
|
|
|
|
self.change_baseline()
|
|
|
|
def change_baseline(self):
|
|
if self.data:
|
|
x, y = self.data
|
|
|
|
def mean(xx):
|
|
return np.mean(y[max(0, np.argmin(abs(x-xx))-5):min(len(x), np.argmin(abs(x-xx))+6)].real)
|
|
|
|
y_node = [mean(x_node) for x_node in self.anchors]
|
|
try:
|
|
self.spline = splrep(self.anchors, y_node, per=False)
|
|
except TypeError:
|
|
self.spline = splrep(self.anchors, y_node, per=False, k=1)
|
|
|
|
bl = splev(x, self.spline)
|
|
|
|
self.baseline.setData(x=x, y=bl)
|
|
self.graph_corr.setData(x=x, y=y.real-bl)
|
|
|
|
def keyPressEvent(self, evt):
|
|
if self.listWidget.hasFocus() and evt.key() == QtCore.Qt.Key_Delete:
|
|
r = self.listWidget.currentRow()
|
|
self.anchors.pop(r+1)
|
|
listitem = self.listWidget.takeItem(r)
|
|
del listitem
|
|
self.graphicsView.removeItem(self.anchor_lines.pop(r))
|
|
self.change_baseline()
|
|
|
|
else:
|
|
super().keyPressEvent(evt)
|