Merge branch 'main' of https://github.com/mesoscope/cellpack into fea…

…ture/firebase-upload-gradients

cellpack/autopack/BaseGrid.py

@@ -208,8 +208,7 @@ def setup(self, boundingBox):
         self.free_points = list(range(self.gridVolume))
         self.nbFreePoints = len(self.free_points)
         self.log.info(
-            "Lookup: %d, bounding box: %r, gridSpacing %r, length compartment_ids %r",
-            boundingBox,
+            "gridSpacing %r, length compartment_ids %r",
             self.gridSpacing,
             len(self.compartment_ids),
         )

cellpack/autopack/Compartment.py

@@ -144,7 +144,7 @@ def __init__(self, name, object_info):
             self.filename = self.representations.get_mesh_path()
             self.path = autopack.fixOnePath(self.filename)
         self.stype = "mesh"
-        self.radius = object_info["radius"] if "radius" in object_info else 200.0
+        self.radius = object_info.get("radius", 200.0)
         self.height = 0.0
         self.axis = [0, 1, 0]
         self.area = 0.0
@@ -154,9 +154,8 @@ def __init__(self, name, object_info):
         self.bb = None
         self.diag = 9999.9
         self.ghost = None
-        self.encapsulating_radius = (
-            object_info["radius"] if "radius" in object_info else 200.0
-        )
+        self.encapsulating_radius = object_info.get("radius", 200.0)
+        self.color = object_info.get("color")
         self.checkinside = True
         self.innerRecipe = None
         self.surfaceRecipe = None
@@ -1078,7 +1077,7 @@ def prepare_buildgrid_box(self, env):
             len(env.grid.masterGridPositions),
         )
 
-    def get_surface_distances(self, env, master_grid_positions):
+    def set_surface_distances(self, env, master_grid_positions):
         surface_mask = numpy.equal(self.number, env.grid.compartment_ids)
         surface_ids = numpy.nonzero(surface_mask)
         surface_positions = master_grid_positions[surface_ids]
@@ -1095,7 +1094,7 @@ def get_surface_distances(self, env, master_grid_positions):
         all_surface_distances = numpy.full(master_grid_positions.shape[0], numpy.nan)
         all_surface_distances[grid_pt_indexes] = surface_distances
 
-        if self.parent is not None:
+        if self.parent is not None and parent_id != 0:
             grid_pts_between_surfaces = numpy.equal(
                 env.grid.compartment_ids, -parent_id
             )
@@ -1114,7 +1113,7 @@ def get_surface_distances(self, env, master_grid_positions):
             )
             self.scaled_distance_to_next_surface = scaled_distance_to_next_surface
 
-        self.max_distance = max(surface_distances)
+        self.max_distance = numpy.nanmax(surface_distances)
         self.surface_distances = all_surface_distances
 
     def BuildGrid_box(self, env, vSurfaceArea):
@@ -3120,6 +3119,84 @@ def getSurfaceInnerPoints(self, boundingBox, spacing, display=True, useFix=False
         self.grid_distances = distances
         return insidePoints, surfacePoints
 
+    def create_voxelized_mask(
+        self,
+        x_width,
+        y_width,
+        z_width,
+        center,
+        voxel_size,
+        mesh_store,
+        hollow=False,
+    ):
+        """
+        Creates a mask of the compartment voxelization
+
+        :param x_width: number of voxels in the x direction
+        :param y_width: number of voxels in the y direction
+        :param z_width: number of voxels in the z direction
+        :param center: center of the voxelization
+        :param voxel_size: size of the voxels
+        :param mesh_store: mesh store
+        :param hollow: if True, the mask will be hollow otherwise it will fill in the segmentation
+        """
+        if voxel_size is None:
+            voxel_size = numpy.array([1, 1, 1], dtype=int)
+
+        if center is None:  # use the middle of the grid
+            center = (self.bounding_box[0] + self.bounding_box[1]) / 2.0
+
+        Z, Y, X = numpy.meshgrid(
+            numpy.arange(z_width),
+            numpy.arange(y_width),
+            numpy.arange(x_width),
+            indexing="ij",
+        )
+        coords = numpy.column_stack(
+            [
+                (X * voxel_size[0] - center[0]).flatten(),
+                (Y * voxel_size[1] - center[1]).flatten(),
+                (Z * voxel_size[2] - center[2]).flatten(),
+            ]
+        )
+
+        mesh = mesh_store.get_mesh(self.gname)
+
+        if hollow:
+            trimesh_grid = creation.voxelize(mesh, pitch=numpy.min(voxel_size)).hollow()
+        else:
+            trimesh_grid = creation.voxelize(mesh, pitch=numpy.min(voxel_size)).fill()
+
+        mask_ravel = trimesh_grid.is_filled(coords)
+        mask = mask_ravel.reshape((x_width, y_width, z_width), order="F")
+
+        return mask
+
+    def create_voxelization(
+        self,
+        image_data,
+        bounding_box,
+        voxel_size,
+        image_size,
+        position,
+        mesh_store,
+        hollow=False,
+    ):
+        """
+        Creates a voxelization mask at the position of the compartment
+        """
+        relative_position = position - bounding_box[0]
+        mask = self.create_voxelized_mask(
+            *image_size,
+            center=relative_position,
+            voxel_size=voxel_size,
+            mesh_store=mesh_store,
+            hollow=hollow,
+        )
+        image_data[mask] = 255
+
+        return image_data
+
     def getSurfaceInnerPointsPandaRay(
         self, boundingBox, spacing, display=True, useFix=False
     ):