seabreeze-qt.py

# SPDX-License-Identifier: MIT

import faulthandler
import sys

import PySide6  # noqa: F401, E402
import pyqtgraph as pg
import pyqtgraph.parametertree
import numpy as np
import numpy.typing as npt

from PySide6 import QtCore, QtGui, QtWidgets


# Use the pyseabreeze backend

import seabreeze

seabreeze.use('pyseabreeze')

import seabreeze.spectrometers as sb


class MainWindow(QtWidgets.QMainWindow):  # type: ignore[misc]
    def __init__(self, device: seabreeze.spectrometers.Spectrometer) -> None:
        super().__init__()

        self._device = device
        self._wavelengths = self._device.wavelengths()

        self.setup_graph()

        self.layout_ = QtWidgets.QHBoxLayout()
        self.layout_.addWidget(self.graph)

        container = QtWidgets.QWidget()
        container.setLayout(self.layout_)

        self.setCentralWidget(container)

    def setup_graph(self) -> None:
        self.graph = pg.PlotWidget()
        self.graph.setWindowTitle('test')
        self.graph.setAutoPan(y=False)
        self.graph.setYRange(0, 1)
        self.graph.setMouseEnabled(x=False, y=False)
        self.graph_curve = self.graph.plot(
            fillLevel=-1,
            brush=self.wavelength_colormap().getBrush(span=(300, 900), orientation='horizontal'),
            pen=self.wavelength_colormap().getPen(span=(300, 900), orientation='horizontal'),
        )
        self.graph.show()

        self.graph_timer = QtCore.QTimer()
        self.graph_timer.timeout.connect(self.update_graph)
        self.graph_timer.start(50)

    def update_graph(self) -> None:
        self.graph_curve.setData(self._wavelengths, self._device.intensities())

    def wavelength_to_rgb(self, wavelength: int, gamma: float = 0.8) -> npt.NDArray[np.float64]:
        """Based on http://www.physics.sfasu.edu/astro/color/spectra.html"""
        red: float
        green: float
        blue: float
        alpha: float

        # set alpha gradient for invisible light
        if wavelength < 420:
            alpha = 0.3 + 0.7 * (wavelength - 380) / (420 - 380)
        elif wavelength > 700:
            alpha = 0.3 + 0.7 * (780 - wavelength) / (780 - 700)
        else:
            alpha = 1.
        if alpha < 0:
            alpha = 0

        # set color to the limit for invisible light ()
        if wavelength < 380:
            wavelength = 380
        if wavelength > 780:
            wavelength = 780

        if 380 <= wavelength <= 440:
            attenuation = 0.3 + 0.7 * (wavelength - 380) / (440 - 380)
            red = ((-(wavelength - 440) / (440 - 380)) * attenuation) ** gamma
            green = 0.0
            blue = (1.0 * attenuation) ** gamma
        elif 440 <= wavelength <= 490:
            red = 0.0
            green = ((wavelength - 440) / (490 - 440)) ** gamma
            blue = 1.0
        elif 490 <= wavelength <= 510:
            red = 0.0
            green = 1.0
            blue = (-(wavelength - 510) / (510 - 490)) ** gamma
        elif 510 <= wavelength <= 580:
            red = ((wavelength - 510) / (580 - 510)) ** gamma
            green = 1.0
            blue = 0.0
        elif 580 <= wavelength <= 645:
            red = 1.0
            green = (-(wavelength - 645) / (645 - 580)) ** gamma
            blue = 0.0
        elif 645 <= wavelength <= 780:
            attenuation = 0.3 + 0.7 * (780 - wavelength) / (780 - 645)
            red = (1.0 * attenuation) ** gamma
            green = 0.0
            blue = 0.0
        else:
            red, green, blue = 0.0, 0.0, 0.0

        return np.array([red, green, blue, alpha])

    def rgba_to_rgb(
        self,
        color: npt.NDArray[np.float64],
        background: npt.NDArray[np.float64] = np.array([1, 1, 1]),
    ) -> npt.NDArray[np.float64]:
        return background * (1 - color[3]) + (color[:3] * color[3])  # type: ignore[no-any-return]

    def wavelength_colormap(self) -> pyqtgraph.ColorMap:
        colors = [self.rgba_to_rgb(self.wavelength_to_rgb(x))*255 for x in np.arange(300, 900)]
        return pyqtgraph.ColorMap(None, colors)


def run() -> None:
    app = QtWidgets.QApplication(sys.argv)

    device = seabreeze.spectrometers.Spectrometer.from_first_available()
    device.integration_time_micros(20_000)

    window = MainWindow(device)
    window.show()

    app.exec()


def main() -> None:
    faulthandler.enable()

    run()


if __name__ == '__main__':
    main()