pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__cp311-cp311-macosx_15_0_arm64.whl

scripts/preprocessor_gui.py

@@ -19,14 +19,14 @@ from magicgui import widgets
 from numpy.typing import NDArray
 from napari.layers import Image
 from scipy.fft import next_fast_len
-from qtpy.QtWidgets import QFileDialog
 from napari.utils.events import EventedList
+from qtpy.QtWidgets import QFileDialog, QMessageBox
 
-from tme.backends import backend
+from tme.backends import backend as be
 from tme.rotations import align_vectors
-from tme import Preprocessor, Density, Orientations
 from tme.matching_utils import create_mask, load_pickle
-from tme.filters import BandPassFilter, CTFReconstructed
+from tme import Preprocessor, Density, Orientations
+from tme.filters import BandPassReconstructed, CTFReconstructed
 
 preprocessor = Preprocessor()
 SLIDER_MIN, SLIDER_MAX = 0, 25
@@ -43,17 +43,16 @@ def bandpass_filter(
     hard_edges: bool = False,
     sampling_rate=None,
 ) -> NDArray:
-    bpf = BandPassFilter(
+    bpf = BandPassReconstructed(
         lowpass=lowpass_angstrom,
         highpass=highpass_angstrom,
         sampling_rate=np.max(sampling_rate),
         use_gaussian=not hard_edges,
-        shape_is_real_fourier=True,
         return_real_fourier=True,
     )
     template_ft = np.fft.rfftn(template, s=template.shape)
 
-    mask = bpf(shape=template_ft.shape)["data"]
+    mask = bpf(shape=template.shape)["data"]
     np.multiply(template_ft, mask, out=template_ft)
     return np.fft.irfftn(template_ft, s=template.shape).real
 
@@ -70,7 +69,7 @@ def ctf_filter(
     flip_phase: bool = False,
 ) -> NDArray:
     fast_shape = [next_fast_len(x) for x in np.multiply(template.shape, 2)]
-    template_pad = backend.topleft_pad(template, fast_shape)
+    template_pad = be.topleft_pad(template, fast_shape)
     template_ft = np.fft.rfftn(template_pad, s=template_pad.shape)
     ctf = CTFReconstructed(
         shape=fast_shape,
@@ -86,7 +85,7 @@ def ctf_filter(
     )
     np.multiply(template_ft, ctf()["data"], out=template_ft)
     template_pad = np.fft.irfftn(template_ft, s=template_pad.shape).real
-    template = backend.topleft_pad(template_pad, template.shape)
+    template = be.topleft_pad(template_pad, template.shape)
     return template
 
 
@@ -393,7 +392,13 @@ class FilterWidget(widgets.Container):
             metadata = selected_layer_metadata.copy()
             if "filter_parameters" not in metadata:
                 metadata["filter_parameters"] = []
-            metadata["filter_parameters"].append({filter_name: kwargs.copy()})
+
+            payload = {filter_name: kwargs.copy()}
+            if isinstance(metadata["filter_parameters"], dict):
+                metadata["filter_parameters"].update(payload)
+            else:
+                metadata["filter_parameters"].append(payload)
+
             metadata["used_filter"] = filter_name
             new_layer.metadata = metadata
 
@@ -451,7 +456,29 @@ def box_mask(
         mask_type="box",
         shape=template.shape,
         center=(center_x, center_y, center_z),
-        height=(height_x, height_y, height_z),
+        size=(height_x, height_y, height_z),
+        sigma_decay=sigma_decay,
+    )
+
+
+def membrane_mask(
+    template: NDArray,
+    symmetry_axis: int,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    radius: float,
+    thickness: float = 1,
+    separation: float = 3,
+    sigma_decay: float = 0,
+    **kwargs,
+) -> NDArray:
+    return create_mask(
+        mask_type="membrane",
+        shape=template.shape,
+        radius=radius,
+        thickness=thickness,
+        separation=separation,
         sigma_decay=sigma_decay,
     )
 
@@ -472,7 +499,7 @@ def tube_mask(
         mask_type="tube",
         shape=template.shape,
         symmetry_axis=symmetry_axis,
-        base_center=(center_x, center_y, center_z),
+        center=(center_x, center_y, center_z),
         inner_radius=inner_radius,
         outer_radius=outer_radius,
         height=height,
@@ -585,6 +612,7 @@ class MaskWidget(widgets.Container):
             "Ellipsoid": ellipsod_mask,
             "Tube": tube_mask,
             "Box": box_mask,
+            "Membrane": membrane_mask,
             "Wedge": wedge_mask,
             "Threshold": threshold_mask,
             "Lowpass": lowpass_mask,
@@ -818,7 +846,7 @@ class AlignmentWidget(widgets.Container):
         principal_eigenvector = eigenvectors[:, np.argmax(eigenvalues)]
 
         rotation_matrix = align_vectors(principal_eigenvector, alignment_axis)
-        rotated_data, _ = backend.rigid_transform(
+        rotated_data, _ = be.rigid_transform(
             arr=active_layer.data,
             rotation_matrix=rotation_matrix,
             use_geometric_center=False,
@@ -954,7 +982,6 @@ class PointCloudWidget(widgets.Container):
             if not isinstance(layer, napari.layers.Points):
                 continue
 
-            layer.face_color = "white"
             if event == "Label":
                 if len(layer.properties.get("detail", ())) == 0:
                     continue
@@ -971,9 +998,7 @@ class PointCloudWidget(widgets.Container):
                 layer.face_color = "score_scaled"
                 layer.face_colormap = "turbo"
                 layer.face_color_mode = "colormap"
-
             layer.refresh_colors()
-
         return None
 
     def _set_positive(self, event):
@@ -1141,9 +1166,18 @@ class MatchingWidget(widgets.Container):
         self.viewer = viewer
         self.dataframes = {}
 
+        option_container = widgets.Container(layout="horizontal")
+        self.load_target_checkbox = widgets.CheckBox(text="Load Target", value=False)
+        self.load_rotations_checkbox = widgets.CheckBox(
+            text="Load Rotations", value=False
+        )
+        option_container.append(self.load_target_checkbox)
+        option_container.append(self.load_rotations_checkbox)
+
         self.import_button = widgets.PushButton(name="Import", text="Import Pickle")
         self.import_button.clicked.connect(self._get_load_path)
 
+        self.append(option_container)
         self.append(self.import_button)
 
     def _get_load_path(self, event):
@@ -1151,7 +1185,7 @@ class MatchingWidget(widgets.Container):
             self.native,
             "Open Pickle File...",
             "",
-            "Pickle Files (*.pickle);;All Files (*)",
+            "Pickle Files (*.pickle *pickle.gz);;All Files (*)",
         )
         if filename:
             self._load_data(filename)
@@ -1160,14 +1194,35 @@ class MatchingWidget(widgets.Container):
         data = load_pickle(filename)
 
         fname = basename(filename).replace(".pickle", "")
+
+        if self.load_target_checkbox.value:
+            try:
+                target = Density.from_file(data[-1][-1].target)
+                _ = self.viewer.add_image(
+                    data=target.data,
+                    name=f"{fname}_target",
+                    metadata={
+                        "origin": target.origin,
+                        "sampling_rate": target.sampling_rate,
+                    },
+                )
+            except Exception as e:
+                msg = QMessageBox(self.native)
+                msg.setIcon(QMessageBox.Warning)
+                msg.setWindowTitle("Loading Error")
+                msg.setText(str(e))
+                msg.setStandardButtons(QMessageBox.Ok)
+                msg.exec_()
+
         if data[0].ndim == data[2].ndim:
             metadata = {"origin": data[-1][1], "sampling_rate": data[-1][2]}
-            _ = self.viewer.add_image(
-                data=data[2],
-                name=f"{fname}_rotations",
-                colormap="orange",
-                metadata=metadata,
-            )
+            if self.load_rotations_checkbox.value:
+                _ = self.viewer.add_image(
+                    data=data[2],
+                    name=f"{fname}_rotations",
+                    colormap="orange",
+                    metadata=metadata,
+                )
             _ = self.viewer.add_image(
                 data=data[0],
                 name=f"{fname}_scores",
@@ -1175,11 +1230,8 @@ class MatchingWidget(widgets.Container):
                 metadata=metadata,
             )
             return None
-            detail = np.zeros_like(data[2])
-        else:
-            detail = data[3]
 
-        point_properties = {"score": data[2], "detail": detail}
+        point_properties = {"score": data[2], "detail": data[3]}
         point_properties["score_scaled"] = np.log1p(
             point_properties["score"] - point_properties["score"].min()
         )
@@ -1192,8 +1244,26 @@ class MatchingWidget(widgets.Container):
         )
 
 
+class CustomNapariViewer(napari.Viewer):
+    """
+    Custom viewer to ensure 3D image layers are by default shown as xy projection.
+    """
+
+    def add_image(self, data, **kwargs):
+        viewer_ndim = len(self.dims.order)
+        layer = super().add_image(data, **kwargs)
+
+        try:
+            # Set to xy view the first time data is opened
+            if viewer_ndim != 3 and data.ndim == 3:
+                self.dims.order = (2, 0, 1)
+        except Exception:
+            pass
+        return layer
+
+
 def main():
-    viewer = napari.Viewer()
+    viewer = CustomNapariViewer()
 
     filter_widget = FilterWidget(preprocessor, viewer)
     mask_widget = MaskWidget(viewer)