Skip to content

Commit

Permalink
Merge branch 'finetuning' into finetune
Browse files Browse the repository at this point in the history
  • Loading branch information
LorenzLamm authored Sep 17, 2024
2 parents d8db3a1 + 1317d16 commit 1a6e349
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 1 deletion.
5 changes: 4 additions & 1 deletion src/membrain_seg/segmentation/cli/ske_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
store_tomogram,
)


from ..skeletonize import skeletonization as _skeletonization
from .cli import cli

Expand Down Expand Up @@ -53,19 +54,21 @@ def skeletonize(
--batch-size <batch-size>
"""
# Assuming _skeletonization function is already defined and can handle batch_size

segmentation = load_tomogram(label_path)
ske = _skeletonization(segmentation=segmentation.data, batch_size=batch_size)

# Update the segmentation data with the skeletonized output while preserving the original header and voxel_size
segmentation.data = ske


if not os.path.exists(out_folder):
os.makedirs(out_folder)

out_file = os.path.join(
out_folder,
os.path.splitext(os.path.basename(label_path))[0] + "_skel.mrc",
)

store_tomogram(filename=out_file, tomogram=segmentation)
print("Skeleton saved to ", out_file)
4 changes: 4 additions & 0 deletions src/membrain_seg/segmentation/skeletonize.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import scipy.ndimage as ndimage
import torch


from membrain_seg.segmentation.skeletonization.diff3d import (
compute_gradients,
compute_hessian,
Expand All @@ -23,6 +24,7 @@
from membrain_seg.segmentation.training.surface_dice import apply_gaussian_filter



def skeletonization(segmentation: np.ndarray, batch_size: int) -> np.ndarray:
"""
Perform skeletonization on a tomogram segmentation.
Expand All @@ -36,6 +38,7 @@ def skeletonization(segmentation: np.ndarray, batch_size: int) -> np.ndarray:
segmentation : ndarray
Tomogram segmentation as a numpy array, where non-zero values represent
the structures of interest.
batch_size : int
The number of elements to process in one batch during eigen decomposition.
Useful for managing memory usage.
Expand All @@ -58,6 +61,7 @@ def skeletonization(segmentation: np.ndarray, batch_size: int) -> np.ndarray:
--batch-size 1000000
This command runs the skeletonization process from the command line.
"""

# Convert non-zero segmentation values to 1.0
labels = (segmentation > 0) * 1.0

Expand Down

0 comments on commit 1a6e349

Please sign in to comment.