pytme 0.1.5__cp311-cp311-macosx_14_0_arm64.whl

scripts/preprocessor_gui.py

@@ -0,0 +1,729 @@
+#!python3
+""" Simplify picking adequate filtering and masking parameters using a GUI.
+    Exposes tme.preprocessor.Preprocessor and tme.fitter_utils member functions
+    to achieve this aim.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+import inspect
+from typing import Tuple, Callable, List
+from typing_extensions import Annotated
+
+import numpy as np
+import pandas as pd
+import napari
+from napari.layers import Image
+from napari.utils.events import EventedList
+
+from magicgui import widgets
+from qtpy.QtWidgets import QFileDialog
+from numpy.typing import NDArray
+
+from tme import Preprocessor, Density
+from tme.matching_utils import create_mask
+
+preprocessor = Preprocessor()
+SLIDER_MIN, SLIDER_MAX = 0, 25
+
+
+def gaussian_filter(template, sigma: float, **kwargs: dict):
+    return preprocessor.gaussian_filter(template=template, sigma=sigma, **kwargs)
+
+
+def bandpass_filter(
+    template,
+    minimum_frequency: float,
+    maximum_frequency: float,
+    gaussian_sigma: float,
+    **kwargs: dict,
+):
+    return preprocessor.bandpass_filter(
+        template=template,
+        minimum_frequency=minimum_frequency,
+        maximum_frequency=maximum_frequency,
+        sampling_rate=1,
+        gaussian_sigma=gaussian_sigma,
+        **kwargs,
+    )
+
+
+def difference_of_gaussian_filter(
+    template, sigmas: Tuple[float, float], **kwargs: dict
+):
+    low_sigma, high_sigma = sigmas
+    return preprocessor.difference_of_gaussian_filter(
+        template=template, low_sigma=low_sigma, high_sigma=high_sigma, **kwargs
+    )
+
+
+def edge_gaussian_filter(
+    template,
+    sigma: float,
+    edge_algorithm: Annotated[
+        str,
+        {"choices": ["sobel", "prewitt", "laplace", "gaussian", "gaussian_laplace"]},
+    ],
+    reverse: bool = False,
+    **kwargs: dict,
+):
+    return preprocessor.edge_gaussian_filter(
+        template=template,
+        sigma=sigma,
+        reverse=reverse,
+        edge_algorithm=edge_algorithm,
+    )
+
+
+def local_gaussian_filter(
+    template,
+    lbd: float,
+    sigma_range: Tuple[float, float],
+    gaussian_sigma: float,
+    reverse: bool = False,
+    **kwargs: dict,
+):
+    return preprocessor.local_gaussian_filter(
+        template=template,
+        lbd=lbd,
+        sigma_range=sigma_range,
+        gaussian_sigma=gaussian_sigma,
+    )
+
+
+def ntree(
+    template,
+    sigma_range: Tuple[float, float],
+    **kwargs: dict,
+):
+    return preprocessor.ntree_filter(template=template, sigma_range=sigma_range)
+
+
+def mean(
+    template,
+    width: int,
+    **kwargs: dict,
+):
+    return preprocessor.mean_filter(template=template, width=width)
+
+
+def wedge(
+    template: NDArray,
+    tilt_start: float,
+    tilt_stop: float,
+    gaussian_sigma: float,
+    tilt_axis: int = 1,
+    infinite_plane: bool = True,
+    extrude_plane: bool = True,
+):
+    template_ft = np.fft.rfftn(template)
+    wedge_mask = preprocessor.continuous_wedge_mask(
+        start_tilt=tilt_start,
+        stop_tilt=tilt_stop,
+        tilt_axis=tilt_axis,
+        shape=template.shape,
+        sigma=gaussian_sigma,
+        omit_negative_frequencies=True,
+        infinite_plane=infinite_plane,
+        extrude_plane=extrude_plane,
+    )
+    np.multiply(template_ft, wedge_mask, out=template_ft)
+    template = np.real(np.fft.irfftn(template_ft))
+    return template
+
+
+def widgets_from_function(function: Callable, exclude_params: List = ["self"]):
+    """
+    Creates list of magicui widgets by inspecting function typing ann
+    """
+    ret = []
+    for name, param in inspect.signature(function).parameters.items():
+        if name in exclude_params:
+            continue
+
+        if param.annotation is float:
+            widget = widgets.FloatSpinBox(
+                name=name,
+                value=param.default if param.default != inspect._empty else 0,
+                min=SLIDER_MIN,
+                step=0.5,
+            )
+        elif param.annotation == Tuple[float, float]:
+            widget = widgets.FloatRangeSlider(
+                name=param.name,
+                value=param.default
+                if param.default != inspect._empty
+                else (0.0, SLIDER_MAX / 2),
+                min=SLIDER_MIN,
+                max=SLIDER_MAX,
+            )
+        elif param.annotation is int:
+            widget = widgets.SpinBox(
+                name=name,
+                value=param.default if param.default != inspect._empty else 0,
+            )
+        elif param.annotation is bool:
+            widget = widgets.CheckBox(
+                name=name,
+                value=param.default if param.default != inspect._empty else False,
+            )
+        elif hasattr(param.annotation, "__metadata__"):
+            metadata = param.annotation.__metadata__[0]
+            if "choices" in metadata:
+                widget = widgets.ComboBox(
+                    name=param.name,
+                    choices=metadata["choices"],
+                    value=param.default
+                    if param.default != inspect._empty
+                    else metadata["choices"][0],
+                )
+        else:
+            continue
+        ret.append(widget)
+    return ret
+
+
+WRAPPED_FUNCTIONS = {
+    "gaussian_filter": gaussian_filter,
+    "bandpass_filter": bandpass_filter,
+    "edge_gaussian_filter": edge_gaussian_filter,
+    "ntree_filter": ntree,
+    "local_gaussian_filter": local_gaussian_filter,
+    "difference_of_gaussian_filter": difference_of_gaussian_filter,
+    "mean_filter": mean,
+    "continuous_wedge_mask": wedge,
+}
+
+EXCLUDED_FUNCTIONS = [
+    "apply_method",
+    "method_to_id",
+    "wedge_mask",
+    "fourier_crop",
+    "fourier_uncrop",
+    "interpolate_box",
+    "molmap",
+    "local_gaussian_alignment_filter",
+    # "continuous_wedge_mask",
+    "wedge_mask",
+    "bandpass_mask",
+]
+
+
+class FilterWidget(widgets.Container):
+    def __init__(self, preprocessor, viewer):
+        super().__init__(layout="vertical")
+
+        self.preprocessor = preprocessor
+        self.viewer = viewer
+        self.name_mapping = {}
+        self.action_widgets = []
+
+        self.layer_dropdown = widgets.ComboBox(
+            name="Target Layer", choices=self._get_layer_names()
+        )
+        self.append(self.layer_dropdown)
+        self.viewer.layers.events.inserted.connect(self._update_layer_dropdown)
+        self.viewer.layers.events.removed.connect(self._update_layer_dropdown)
+
+        self.method_dropdown = widgets.ComboBox(
+            name="Choose Filter", choices=self._get_method_names()
+        )
+        self.method_dropdown.changed.connect(self._on_method_changed)
+        self.append(self.method_dropdown)
+
+        self.apply_btn = widgets.PushButton(text="Apply Filter", enabled=False)
+        self.apply_btn.changed.connect(self._action)
+        self.append(self.apply_btn)
+
+        # Create GUI for initially selected filtering method
+        self._on_method_changed(None)
+
+    def _get_method_names(self):
+        method_names = [
+            name
+            for name, member in inspect.getmembers(self.preprocessor, inspect.ismethod)
+            if not name.startswith("_") and name not in EXCLUDED_FUNCTIONS
+        ]
+
+        sanitized_names = [self._sanitize_name(name) for name in method_names]
+        self.name_mapping.update(dict(zip(sanitized_names, method_names)))
+
+        return sanitized_names
+
+    def _sanitize_name(self, name: str) -> str:
+        # Replace underscores with spaces and capitalize each word
+        removes = ["blur", "filter"]
+        for remove in removes:
+            name = name.replace(remove, "")
+        return name.strip().replace("_", " ").title()
+
+    def _desanitize_name(self, name: str) -> str:
+        name = name.lower().strip()
+        for function_name, _ in inspect.getmembers(self.preprocessor, inspect.ismethod):
+            if function_name.startswith(name):
+                return function_name
+        return name
+
+    def _get_function(self, name: str):
+        function = WRAPPED_FUNCTIONS.get(name, None)
+        if not function:
+            function = getattr(self.preprocessor, name, None)
+        return function
+
+    def _on_method_changed(self, event=None):
+        # Clear previous parameter widgets
+        for widget in self.action_widgets:
+            self.remove(widget)
+        self.action_widgets.clear()
+
+        function_name = self.name_mapping.get(self.method_dropdown.value)
+        function = self._get_function(function_name)
+
+        widgets = widgets_from_function(function, exclude_params=["self", "template"])
+        for widget in widgets:
+            self.action_widgets.append(widget)
+            self.insert(-1, widget)
+
+    def _update_layer_dropdown(self, event: EventedList):
+        """Update the dropdown menu when layers change."""
+        self.layer_dropdown.choices = self._get_layer_names()
+        self.apply_btn.enabled = bool(self.viewer.layers)
+
+    def _get_layer_names(self):
+        """Return list of layer names in the viewer."""
+        return sorted([layer.name for layer in self.viewer.layers])
+
+    def _action(self, event):
+        selected_layer = self.viewer.layers[self.layer_dropdown.value]
+        selected_layer_metadata = selected_layer.metadata.copy()
+        kwargs = {widget.name: widget.value for widget in self.action_widgets}
+
+        function_name = self.name_mapping.get(self.method_dropdown.value)
+        function = self._get_function(function_name)
+
+        processed_data = function(selected_layer.data, **kwargs)
+
+        new_layer_name = f"{selected_layer.name} ({self.method_dropdown.value})"
+
+        if new_layer_name in self.viewer.layers:
+            selected_layer = self.viewer.layers[new_layer_name]
+
+        filter_name = self._desanitize_name(self.method_dropdown.value)
+        used_filter = selected_layer.metadata.get("used_filter", False)
+        if used_filter == filter_name:
+            selected_layer.data = processed_data
+        else:
+            new_layer = self.viewer.add_image(
+                data=processed_data,
+                name=new_layer_name,
+            )
+            metadata = selected_layer_metadata.copy()
+            if "filter_parameters" not in metadata:
+                metadata["filter_parameters"] = []
+            metadata["filter_parameters"].append({filter_name: kwargs.copy()})
+            metadata["used_filter"] = filter_name
+            new_layer.metadata = metadata
+
+
+def sphere_mask(
+    template: NDArray, center_x: float, center_y: float, center_z: float, radius: float
+):
+    return create_mask(
+        mask_type="ellipse",
+        shape=template.shape,
+        center=(center_x, center_y, center_z),
+        radius=radius,
+    )
+
+
+def ellipsod_mask(
+    template: NDArray,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    radius_x: float,
+    radius_y: float,
+    radius_z: float,
+):
+    return create_mask(
+        mask_type="ellipse",
+        shape=template.shape,
+        center=(center_x, center_y, center_z),
+        radius=(radius_x, radius_y, radius_z),
+    )
+
+
+def box_mask(
+    template: NDArray,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    height_x: int,
+    height_y: int,
+    height_z: int,
+):
+    return create_mask(
+        mask_type="box",
+        shape=template.shape,
+        center=(center_x, center_y, center_z),
+        height=(height_x, height_y, height_z),
+    )
+
+
+def tube_mask(
+    template: NDArray,
+    symmetry_axis: int,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    inner_radius: float,
+    outer_radius: float,
+    height: int,
+):
+    return create_mask(
+        mask_type="tube",
+        shape=template.shape,
+        symmetry_axis=symmetry_axis,
+        base_center=(center_x, center_y, center_z),
+        inner_radius=inner_radius,
+        outer_radius=outer_radius,
+        height=height,
+    )
+
+
+def wedge_mask(
+    template: NDArray,
+    tilt_start: float,
+    tilt_stop: float,
+    gaussian_sigma: float,
+    tilt_axis: int = 1,
+    omit_negative_frequencies: bool = True,
+    extrude_plane: bool = True,
+    infinite_plane: bool = True,
+):
+    wedge_mask = preprocessor.continuous_wedge_mask(
+        start_tilt=tilt_start,
+        stop_tilt=tilt_stop,
+        tilt_axis=tilt_axis,
+        shape=template.shape,
+        sigma=gaussian_sigma,
+        omit_negative_frequencies=omit_negative_frequencies,
+        extrude_plane=extrude_plane,
+        infinite_plane=infinite_plane,
+    )
+    wedge_mask = np.fft.fftshift(wedge_mask)
+    return wedge_mask
+
+def threshold_mask(
+    template: NDArray,
+    standard_deviation : float = 5.0,
+    invert : bool = False
+):
+    template_mean = template.mean()
+    template_deviation = standard_deviation * template.std()
+    upper = template_mean + template_deviation
+    lower = template_mean - template_deviation
+    mask = np.logical_and(template > lower, template < upper)
+    if invert:
+        np.invert(mask, out = mask)
+
+    return mask
+
+
+
+class MaskWidget(widgets.Container):
+    def __init__(self, viewer):
+        super().__init__(layout="vertical")
+
+        self.viewer = viewer
+        self.action_widgets = []
+
+        self.action_button = widgets.PushButton(text="Create mask", enabled=False)
+        self.action_button.changed.connect(self._action)
+
+        self.methods = {
+            "Sphere": sphere_mask,
+            "Ellipsoid": ellipsod_mask,
+            "Tube": tube_mask,
+            "Box": box_mask,
+            "Wedge": wedge_mask,
+            "Threshold" : threshold_mask
+        }
+
+        self.method_dropdown = widgets.ComboBox(
+            name="Choose Mask", choices=list(self.methods.keys())
+        )
+        self.method_dropdown.changed.connect(self._on_method_changed)
+
+        self.adapt_button = widgets.PushButton(text="Adapt to layer", enabled=False)
+        self.adapt_button.changed.connect(self._update_initial_values)
+
+        self.viewer.layers.selection.events.active.connect(
+            self._update_action_button_state
+        )
+
+        self.align_button = widgets.PushButton(text="Align to axis", enabled=False)
+        self.align_button.changed.connect(self._align_with_axis)
+        self.density_field = widgets.Label()
+        # self.density_field.value = f"Positive Density in Mask: {0:.2f}%"
+
+        self.append(self.method_dropdown)
+        self.append(self.adapt_button)
+        self.append(self.align_button)
+        self.append(self.action_button)
+        self.append(self.density_field)
+
+        # Create GUI for initially selected filtering method
+        self._on_method_changed(None)
+
+    def _update_action_button_state(self, event):
+        self.align_button.enabled = bool(self.viewer.layers.selection.active)
+        self.action_button.enabled = bool(self.viewer.layers.selection.active)
+        self.adapt_button.enabled = bool(self.viewer.layers.selection.active)
+
+    def _align_with_axis(self):
+        active_layer = self.viewer.layers.selection.active
+
+        if active_layer.metadata.get("is_aligned", False):
+            return
+
+        coords = np.array(np.where(active_layer.data > 0)).T
+        centered_coords = coords - np.mean(coords, axis=0)
+        cov_matrix = np.cov(centered_coords, rowvar=False)
+
+        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
+        principal_eigenvector = eigenvectors[:, np.argmax(eigenvalues)]
+
+        rotation_axis = np.cross(principal_eigenvector, [1, 0, 0])
+        rotation_angle = np.arccos(np.dot(principal_eigenvector, [1, 0, 0]))
+        k = rotation_axis / np.linalg.norm(rotation_axis)
+        K = np.array([[0, -k[2], k[1]], [k[2], 0, -k[0]], [-k[1], k[0], 0]])
+        rotation_matrix = np.eye(3)
+        rotation_matrix += np.sin(rotation_angle) * K
+        rotation_matrix += (1 - np.cos(rotation_angle)) * np.dot(K, K)
+
+        rotated_data = Density.rotate_array(
+            arr=active_layer.data,
+            rotation_matrix=rotation_matrix,
+            use_geometric_center=False,
+        )
+        eps = np.finfo(rotated_data.dtype).eps
+        rotated_data[rotated_data < eps] = 0
+
+        active_layer.metadata["is_aligned"] = True
+        active_layer.data = rotated_data
+
+    def _update_initial_values(self, event=None):
+        active_layer = self.viewer.layers.selection.active
+        center_of_mass = Density.center_of_mass(np.abs(active_layer.data), 0)
+        coordinates = np.array(np.where(active_layer.data > 0))
+        coordinates_min = coordinates.min(axis=1)
+        coordinates_max = coordinates.max(axis=1)
+        coordinates_heights = coordinates_max - coordinates_min
+        coordinate_radius = np.divide(coordinates_heights, 2)
+        center_of_mass = coordinate_radius + coordinates_min
+
+        defaults = dict(zip(["center_x", "center_y", "center_z"], center_of_mass))
+        defaults.update(
+            dict(zip(["radius_x", "radius_y", "radius_z"], coordinate_radius))
+        )
+        defaults.update(
+            dict(zip(["height_x", "height_y", "height_z"], coordinates_heights))
+        )
+
+        defaults["radius"] = np.min(coordinate_radius)
+        defaults["inner_radius"] = np.min(coordinate_radius)
+        defaults["outer_radius"] = np.max(coordinate_radius)
+        defaults["height"] = defaults["radius"]
+
+        for widget in self.action_widgets:
+            if widget.name in defaults:
+                widget.value = defaults[widget.name]
+
+    def _on_method_changed(self, event=None):
+        for widget in self.action_widgets:
+            self.remove(widget)
+        self.action_widgets.clear()
+
+        function = self.methods.get(self.method_dropdown.value)
+        widgets = widgets_from_function(function)
+        for widget in widgets:
+            self.action_widgets.append(widget)
+            self.insert(1, widget)
+
+    def _action(self):
+        function = self.methods.get(self.method_dropdown.value)
+
+        selected_layer = self.viewer.layers.selection.active
+        kwargs = {widget.name: widget.value for widget in self.action_widgets}
+        processed_data = function(template=selected_layer.data, **kwargs)
+        new_layer_name = f"{selected_layer.name} ({self.method_dropdown.value})"
+
+        if new_layer_name in self.viewer.layers:
+            selected_layer = self.viewer.layers[new_layer_name]
+
+        processed_data = processed_data.astype(np.float32)
+        metadata = selected_layer.metadata
+        mask = metadata.get("mask", False)
+        if mask == self.method_dropdown.value:
+            selected_layer.data = processed_data
+        else:
+            new_layer = self.viewer.add_image(
+                data=processed_data,
+                name=new_layer_name,
+            )
+            metadata = selected_layer.metadata.copy()
+            metadata["filter_parameters"] = {self.method_dropdown.value: kwargs.copy()}
+            metadata["mask"] = self.method_dropdown.value
+            metadata["origin_layer"] = selected_layer.name
+            new_layer.metadata = metadata
+
+        origin_layer = metadata["origin_layer"]
+        if origin_layer in self.viewer.layers:
+            origin_layer = self.viewer.layers[origin_layer]
+            if np.allclose(origin_layer.data.shape, processed_data.shape):
+                in_mask = np.sum(np.fmax(origin_layer.data * processed_data, 0))
+                in_mask /= np.sum(np.fmax(origin_layer.data, 0))
+                in_mask *= 100
+                self.density_field.value = f"Positive Density in Mask: {in_mask:.2f}%"
+
+
+class ExportWidget(widgets.Container):
+    def __init__(self, viewer):
+        super().__init__(layout="vertical")
+
+        self.viewer = viewer
+        self.selected_filename = ""
+
+        horizontal_container = widgets.Container(layout="horizontal")
+
+        self.gzip_output = widgets.CheckBox(name="gzip", value=False, label="gzip")
+        self.export_button = widgets.PushButton(name="Export", text="Export")
+        self.export_button.clicked.connect(self._get_save_path)
+
+        horizontal_container.append(self.export_button)
+        horizontal_container.append(self.gzip_output)
+
+        self.append(horizontal_container)
+
+        self.export_button.enabled = bool(self.viewer.layers.selection.active)
+        self.viewer.layers.selection.events.active.connect(
+            self._update_export_button_state
+        )
+
+    def _get_save_path(self, event):
+        options = QFileDialog.Options()
+        path, _ = QFileDialog.getSaveFileName(
+            self.native,
+            "Save As...",
+            "",
+            "MRC Files (*.mrc)",
+            options=options,
+        )
+        if path:
+            self.selected_filename = path
+            self._export_data()
+
+    def _update_export_button_state(self, event):
+        """Update the enabled state of the export button based on the active layer."""
+        self.export_button.enabled = bool(self.viewer.layers.selection.active)
+
+    def _export_data(self):
+        selected_layer = self.viewer.layers.selection.active
+        if selected_layer and isinstance(selected_layer, Image):
+            selected_layer.metadata["write_gzip"] = self.gzip_output.value
+            selected_layer.save(path=self.selected_filename)
+
+
+class PointCloudWidget(widgets.Container):
+    def __init__(self, viewer):
+        super().__init__(layout="vertical")
+
+        self.viewer = viewer
+        self.dataframes = {}
+
+        self.import_button = widgets.PushButton(
+            name="Import", text="Import Point Cloud"
+        )
+        self.import_button.clicked.connect(self._get_load_path)
+
+        self.export_button = widgets.PushButton(
+            name="Export", text="Export Point Cloud"
+        )
+        self.export_button.clicked.connect(self._export_point_cloud)
+        self.export_button.enabled = False
+
+        self.append(self.import_button)
+        self.append(self.export_button)
+        self.viewer.layers.selection.events.changed.connect(self._update_buttons)
+
+    def _update_buttons(self, event):
+        is_pointcloud = isinstance(
+            self.viewer.layers.selection.active, napari.layers.Points
+        )
+        if self.viewer.layers.selection.active and is_pointcloud:
+            self.export_button.enabled = True
+        else:
+            self.export_button.enabled = False
+
+    def _export_point_cloud(self, event):
+        options = QFileDialog.Options()
+        filename, _ = QFileDialog.getSaveFileName(
+            self.native,
+            "Save Point Cloud File...",
+            "",
+            "TSV Files (*.tsv);;All Files (*)",
+            options=options,
+        )
+
+        if filename:
+            layer = self.viewer.layers.selection.active
+            if layer and isinstance(layer, napari.layers.Points):
+                original_dataframe = self.dataframes.get(layer.name, pd.DataFrame())
+
+                export_data = pd.DataFrame(layer.data, columns=["z", "y", "x"])
+                merged_data = pd.merge(
+                    export_data, original_dataframe, on=["z", "y", "x"], how="left"
+                )
+                merged_data.to_csv(filename, sep="\t", index=False)
+
+    def _get_load_path(self, event):
+        options = QFileDialog.Options()
+        filename, _ = QFileDialog.getOpenFileName(
+            self.native,
+            "Open Point Cloud File...",
+            "",
+            "TSV Files (*.tsv);;All Files (*)",
+            options=options,
+        )
+        if filename:
+            self._load_point_cloud(filename)
+
+    def _load_point_cloud(self, filename):
+        dataframe = pd.read_csv(filename, sep="\t")
+        points = dataframe[["z", "y", "x"]].values
+        layer_name = filename.split("/")[-1]
+        self.viewer.add_points(points, size=10, name=layer_name)
+        self.dataframes[layer_name] = dataframe
+
+
+def main():
+    viewer = napari.Viewer()
+
+    filter_widget = FilterWidget(preprocessor, viewer)
+    mask_widget = MaskWidget(viewer)
+    export_widget = ExportWidget(viewer)
+    point_cloud = PointCloudWidget(viewer)
+
+    viewer.window.add_dock_widget(widget=filter_widget, name="Preprocess", area="right")
+    viewer.window.add_dock_widget(widget=mask_widget, name="Mask", area="right")
+    viewer.window.add_dock_widget(widget=point_cloud, name="PointCloud", area="left")
+
+    viewer.window.add_dock_widget(widget=export_widget, name="Export", area="right")
+
+    napari.run()
+
+
+if __name__ == "__main__":
+    main()