Source code for SDF.GUI.browser.preview

from functools import wraps
from typing import Union, Optional

import matplotlib.pyplot as plt
import numpy as np
from PIL.Image import Image
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QWidget, QCheckBox, QHBoxLayout, QVBoxLayout
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar

from SDF.data_model import Dataset, SDFObject, ImageDataset, ArrayDataset2D, ArrayDataset1D


def update_ax(plot_func):
    """
    Decorator for plot functions used in this widget. It performs all the pre- and
    post-plotting work so that the plot functions just need to update the axis `self.ax`
    """
    @wraps(plot_func)
    def wrapper(self: "PlotPreviewWidget", *args, **kwargs):
        # clear axis
        self.ax.clear()

        # call the actual plot function
        result = plot_func(self, *args, **kwargs)

        # update axis
        self.ax.relim()
        self.ax.autoscale()
        self.canvas.draw()

        return result

    return wrapper


class PlotPreviewWidget(QWidget):
    def __init__(self, parent):
        super().__init__(parent)
        self.figure = plt.figure()
        self.ax = self.figure.add_subplot(1, 1, 1)
        self.canvas = FigureCanvas(self.figure)
        self.init_ui()

    def init_ui(self):
        """
        Initializes all widgets (buttons, labels, etc.) displayed
        on this widget.
        """
        self.figure.subplots_adjust(top=0.96, bottom=0.08, left=0.08, right=0.96)

        # navigation toolbar
        toolbar = NavigationToolbar(self.canvas, self)

        # layout
        vbox = QVBoxLayout()
        vbox.addWidget(toolbar)
        vbox.addWidget(self.canvas)
        self.setLayout(vbox)

    def clear(self):
        self.ax.clear()
        self.canvas.draw()

    @update_ax
    def plot_image(self, img: Union[Image, np.ndarray]):
        if isinstance(img, Image):
            # plt.imshow displays grayscale PIL Images much brighter than np.ndarrays of it
            img = np.asarray(img)
        self.ax.imshow(img, cmap="gray")

    @update_ax
    def plot_array(self, data: np.ndarray, x_axis: int = None):
        if data.ndim not in (1, 2):
            raise ValueError(f"Cannot handle {data.ndim} dimensional arrays")

        self.ax.set_aspect('auto')

        # 1D
        if data.ndim == 1:
            self.ax.plot(data)
            return

        # 2D
        y_indices = np.arange(data.shape[1])

        if x_axis is None:
            x = np.arange(data.shape[0])
        else:
            x = data[:, x_axis]
            y_indices = y_indices[y_indices != x_axis]

        if len(y_indices) > 10:
            msg = f'Too many columns ({data.shape[1]}) to plot a preview'
            self.ax.text(0.1, 0.5, msg, {'color': 'red', 'size': 16})
            return

        for y_index in y_indices:
            self.ax.plot(x, data[:, y_index])


class DataPreviewWidget(QWidget):

    def __init__(self):
        super().__init__()
        self.show_preview = True
        self.sdf_obj: Optional[SDFObject] = None

        self.plot_preview_widget = PlotPreviewWidget(self)
        self.init_ui()

    def init_ui(self):
        """Initializes all widgets (buttons, labels, etc.) displayed on this widget"""
        self.setLayout(QVBoxLayout())

        checkbox_data_preview = QCheckBox("Data preview")
        checkbox_data_preview.toggle()
        checkbox_data_preview.stateChanged.connect(self.toggle_show_preview)

        hbox_checkbox = QHBoxLayout()
        hbox_checkbox.addWidget(checkbox_data_preview)
        hbox_checkbox.addStretch(1)
        self.layout().addLayout(hbox_checkbox)

        self.layout().addWidget(self.plot_preview_widget)

    def display_sdf_object(self, obj: SDFObject) -> None:
        """
        Function to update shwon information to match given sdf_object.
        """
        self.sdf_obj = obj
        if not self.show_preview:
            return
        if not isinstance(obj, Dataset):
            self.plot_preview_widget.clear()
            return
        if isinstance(obj, ImageDataset):
            self.plot_preview_widget.plot_image(obj.data)
        elif isinstance(obj, (ArrayDataset1D, ArrayDataset2D)):
            self.plot_preview_widget.plot_array(obj.data)
        else:
            self.plot_preview_widget.clear()

    def toggle_show_preview(self, state: int) -> None:
        if state == Qt.Checked:
            self.show_preview = True
            if self.sdf_obj is None:
                return
            self.display_sdf_object(self.sdf_obj)
        else:
            self.show_preview = False
            self.plot_preview_widget.clear()