Source code for ClearMap.Visualization.Qt.widgets

import numpy as np
import pandas as pd
import pyqtgraph as pg
from PyQt5.QtGui import QColor
from matplotlib.colors import to_hex

from ClearMap.gui.gui_utils import pseudo_random_rgb_array



[docs]
class Scatter3D:
    def __init__(self, coordinates, smarties=False, colors=None, hemispheres=None, half_slice_thickness=None):
        self.__coordinates = None
        self.__has_hemispheres = hemispheres is not None
        self.__has_colours = colors is not None or smarties
        self.default_symbol = '+'
        self.alternate_symbol = 'p'
        self.half_slice_thickness = half_slice_thickness
        self.axis = 2

        if smarties and colors is None:
            n_samples = coordinates.shape[0]
            colors = pseudo_random_rgb_array(n_samples)
        if colors is not None and not isinstance(colors[0], str):  # Convert to hex if not yet
            colors = [to_hex(c) for c in colors]

        if hemispheres is not None:
            self.symbol_map = {
                0: self.default_symbol,
                255: self.alternate_symbol
            }
            # symbols = symbols.decode()

        # colors = if colors is None else np.array([QColor( * col.astype(int)) for col in colors]
        self.data = pd.DataFrame({
            'x': coordinates[:, 0],
            'y': coordinates[:, 1],
            'z': coordinates[:, 2],
            'hemisphere': hemispheres,
            'colour': colors
        })  # TODO: could use id instead of colour
        if colors is not None:
            unique_colors = np.unique(colors)
            self.point_map = pd.DataFrame({
                'colour': unique_colors,
                'pen': [pg.mkPen(c) for c in unique_colors],
                'brush': [pg.mkBrush(c) for c in unique_colors]
            })

    @property
    def coordinates(self):
        if self.__coordinates is None:
            self.__coordinates = self.data[['x', 'y', 'z']].values
        return self.__coordinates


[docs]
    def set_data(self, df):
        # print(self.data['colour'].values, df['colour'])
        if isinstance(df, dict):
            df = pd.DataFrame(df)
        if 'colour' in df.columns:
            if isinstance(df['colour'][0], np.ndarray):  # TODO: should be iterable
                df['colour'] = [to_hex(c) for c in df['colour']]  # OPTIMISE: see map
            elif isinstance(df['colour'][0], QColor):
                df['colour'] = [c.name() for c in df['colour']]  # OPTIMISE: see map
            unique_colors = np.unique(df['colour'])
            self.point_map = pd.DataFrame({
                'colour': unique_colors,
                'pen': [pg.mkPen(c) for c in unique_colors],
                'brush': [pg.mkBrush(c) for c in unique_colors]
            })
        if set(self.data.columns) >= set(df.columns):
            for col in set(self.data.columns) - set(df.columns):  # Add if missing
                df[col] = None
            self.data = df
            # print(self.data['colour'].values, df['colour'].values)
            self.__coordinates = None


    @property
    def has_colours(self):
        return self.__has_colours

    @property
    def has_hemispheres(self):
        return self.__has_hemispheres


[docs]
    def get_all_data(self, main_slice_idx, half_slice_thickness=3):  # FIXME: rename
        """
        Get surrounding markers

        Parameters
        ----------
        main_slice_idx int

        half_slice_thickness int

        Returns
        -------

        """
        if not half_slice_thickness:  # WARNING: optimisation
            output = {
                'pos': np.empty(0),
                'size': np.empty(0),
            }
            if self.has_colours:
                output['pen'] = np.empty(0)
            return output
        half_slice_thickness = self.half_slice_thickness if self.half_slice_thickness is not None else half_slice_thickness
        rows = [pd.DataFrame(columns=['x', 'y', 'colour', 'size'])]
        for i in range(main_slice_idx - half_slice_thickness, main_slice_idx + half_slice_thickness):
            current_z_data = pd.DataFrame(columns=['x', 'y', 'colour', 'size'])  # WARNING: this is x/y of the view, not the 3D image
            if i < 0:  # or i > self.coordinates[:, 2].max()
                continue
            else:
                current_slice = i
            pos = self.get_pos(current_slice)
            if not all(pos.shape):  # empty
                continue
            current_z_data[['x', 'y']] = pos
            current_z_data['colour'] = self.get_colours(current_slice).values  # Otherwise uses index from source
            current_z_data['size'] = self.get_symbol_sizes(main_slice_idx, current_slice, half_size=half_slice_thickness)
            rows.append(current_z_data)

        data_df = pd.concat(rows)
        output = {'pos': data_df[['x', 'y']].values,  # WARNING: this is x/y of the view, not the 3D image
                  'size': data_df['size'].values}
        if self.has_colours:
            output['pen'] = data_df['colour'].map(dict(self.point_map[['colour', 'pen']].values)).values
        return output



[docs]
    def get_draw_params(self, current_slice):
        colours = self.get_colours(current_slice)
        draw_params = {
            'pen': colours.map(dict(self.point_map[['colour', 'pen']].values)).values,
            'brush': colours.map(dict(self.point_map[['colour', 'brush']].values)).values
        }
        return draw_params



[docs]
    def get_symbols(self, current_slice):
        if self.has_hemispheres:
            indices = self.current_slice_indices(current_slice)
            if indices is not None:
                return self.data.loc[indices, 'hemisphere'].map(self.symbol_map).values
            else:
                return np.array([])
        else:
            return self.default_symbol



[docs]
    def get_symbol_sizes(self, main_slice_idx, slice_idx, half_size=3):
        marker_size = round(10 * ((half_size - abs(main_slice_idx - slice_idx)) / half_size))
        n_markers = self.get_n_markers(slice_idx)
        return np.full(n_markers, marker_size)



[docs]
    def get_n_markers(self, slice_idx):
        if len(self.data):
            return np.count_nonzero(self.current_slice_indices(slice_idx))
        else:
            return 0



[docs]
    def get_colours(self, current_slice):
        indices = self.current_slice_indices(current_slice)
        if indices is not None:
            return self.data.loc[indices, 'colour']
        else:
            return np.array([])



[docs]
    def current_slice_indices(self, current_slice):
        if len(self.data):
            return self.coordinates[:, self.axis] == current_slice



[docs]
    def get_pos(self, current_slice):
        indices = self.current_slice_indices(current_slice)
        if indices is not None:
            axes = [0, 1, 2]
            axes.pop(self.axis)  # coordinates in the two other axes
            return self.coordinates[np.ix_(indices, axes)]
        else:
            return np.empty((0, 2))