Merge branch 'main' of github.com:KosinskiLab/colabseg

colabseg/new_gui_functions.py

@@ -16,12 +16,11 @@
 from tqdm.notebook import tqdm
 
 from .image_io import ImageIO
+from .parametrization import PARAMETRIZATION_TYPE
 from .utilities import (
     plane_fit,
     make_plot_array,
     R_2vect,
-    create_sphere_points,
-    lstsq_sphere_fitting,
 )
 
 class ColabSegData(object):
@@ -536,19 +535,11 @@ def flip_normals(self):
             "surface_normals"
         ] * (-1)
 
-    def interpolate_membrane_sphere(self, cluster_index=0):
+    def interpolate_membrane_closed_surface(self, shape_type, cluster_index=0):
         """Least square fit for a perfect sphere and adding of points.
         For vesicles and spherical viruses.
         """
-        radius, x0, y0, z0 = lstsq_sphere_fitting(
-            np.asarray(self.cluster_list_tv)[cluster_index]
-        )
-        # pixel_based_point_count is some approximate heursitic to have every pixel filled
-        # might need to adapt this later to exact solution
-        pixel_based_point_count = int(
-            np.round((radius * np.pi) / (self.pixel_size[0] * 0.25))
-        )
-        interpxyz = create_sphere_points(radius, x0, y0, z0, n=pixel_based_point_count)
+        interpxyz = PARAMETRIZATION_TYPE[shape_type].fit(np.asarray(self.cluster_list_tv)[cluster_index]).sample(1000)
         self.cluster_list_fits.append(interpxyz)
         return
 

colabseg/py3dmol_controls.py

@@ -183,8 +183,9 @@ def load_all_models_fit(self, cluster_list, start_index):
         Use to initialize gui. Use to reload models after editing point cloud
         """
         for i, cluster_positions in enumerate(cluster_list):
+            downsample_fit = int(np.round(len(cluster_positions) / 50000))
             i = i + start_index
-            xyz = self.make_xyz_string(cluster_positions[:: self.downsample])
+            xyz = self.make_xyz_string(cluster_positions[::downsample_fit])
             self.view.addModel(xyz, "xyz")
             self.view.setStyle(
                 {"model": i},

colabseg/segmentation_gui.py

@@ -157,6 +157,12 @@ def gui_elements_cluster_analysis(self):
         self.all_widgets["fit_sphere"] = widgets.Button(description="Fit Sphere")
         self.all_widgets["fit_sphere"].on_click(self.fit_sphere)
 
+        self.all_widgets["fit_ellipsoid"] = widgets.Button(description="Fit Ellipsoid")
+        self.all_widgets["fit_ellipsoid"].on_click(self.fit_ellipsoid)
+
+        self.all_widgets["fit_cylinder"] = widgets.Button(description="Fit Cylinder")
+        self.all_widgets["fit_cylinder"].on_click(self.fit_cylinder)
+
         self.all_widgets["crop_fit"] = widgets.Button(description="Crop fit around")
         self.all_widgets["crop_fit"].on_click(self.crop_fit)
 
@@ -493,10 +499,17 @@ def gui_elements_cluster_analysis(self):
         self.all_widgets["cropping"] = widgets.HBox(
             [self.all_widgets["crop_fit"], self.all_widgets["distance_tolerance"]]
         )
+        self.all_widgets["fit_closed_surface"] = widgets.HBox(
+            [
+                self.all_widgets["fit_sphere"],
+                self.all_widgets["fit_ellipsoid"],
+                self.all_widgets["fit_cylinder"]
+            ]
+        )
         self.all_widgets["fitting"] = widgets.VBox(
             [
                 self.all_widgets["fitting_rbf"],
-                self.all_widgets["fit_sphere"],
+                self.all_widgets["fit_closed_surface"],
                 self.all_widgets["cropping"],
                 self.all_widgets["delete_fit"],
             ]
@@ -808,7 +821,36 @@ def fit_sphere(self, obj):
             print("Nothing or too many clusters selected!")
             print("Please select a single cluster for the fitting procedure!")
             return
-        self.data_structure.interpolate_membrane_sphere(
+        self.data_structure.interpolate_membrane_closed_surface(
+            "sphere",
+            self.all_widgets["cluster_sel"].value
+        )
+        self.reload_gui()
+        return
+
+    def fit_ellipsoid(self, obj):
+            """Fit lstsq ellipsoid to selected cluster"""
+            self.data_structure.backup_step_to_previous()
+            if len(self.all_widgets["cluster_sel"].value) != 1:
+                print("Nothing or too many clusters selected!")
+                print("Please select a single cluster for the fitting procedure!")
+                return
+            self.data_structure.interpolate_membrane_closed_surface(
+                "ellipsoid",
+                self.all_widgets["cluster_sel"].value
+            )
+            self.reload_gui()
+            return
+
+    def fit_cylinder(self, obj):
+        """Fit lstsq ellipsoid to selected cluster"""
+        self.data_structure.backup_step_to_previous()
+        if len(self.all_widgets["cluster_sel"].value) != 1:
+            print("Nothing or too many clusters selected!")
+            print("Please select a single cluster for the fitting procedure!")
+            return
+        self.data_structure.interpolate_membrane_closed_surface(
+            "cylinder",
             self.all_widgets["cluster_sel"].value
         )
         self.reload_gui()