prediction mask threshold API change and doc update (#7)

teamtomo · Dec 3, 2022 · 741d1d0 · 741d1d0
1 parent ba113f6
commit 741d1d0
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/docs/index.md b/docs/index.md
@@ -7,7 +7,7 @@
 [![codecov](https://codecov.io/gh/teamtomo/fidder/branch/main/graph/badge.svg)](https://codecov.io/gh/teamtomo/fidder)
 
 
-*fidder* is a package for detecting and erasing gold fiducials in cryo-EM
+*fidder* is a Python package for detecting and erasing gold fiducials in cryo-EM
 images.
 
 <script
@@ -31,6 +31,7 @@ and the
 # Quickstart
 
 ## Python
+
 ```python
 import mrcfile
 import torch
@@ -41,7 +42,9 @@ from fidder.erase import erase_masked_region
 image = torch.tensor(mrcfile.read('my_image_file.mrc'))
 
 # use a pretrained model to predict a mask
-mask, probabilities = predict_fiducial_mask(image, pixel_spacing=1.35)
+mask, probabilities = predict_fiducial_mask(
+    image, pixel_spacing=1.35, probability_threshold=0.5
+)
 
 # erase fiducials
 erased_image = erase_masked_region(image=image, mask=mask)
@@ -52,6 +55,7 @@ erased_image = erase_masked_region(image=image, mask=mask)
 # predict fiducial mask
 fidder predict \
 --input-image example.mrc \
+--probability-threshold 0.5 \
 --output-mask mask.mrc
 
 # erase masked region

diff --git a/src/fidder/predict/cli.py b/src/fidder/predict/cli.py
@@ -42,7 +42,7 @@ def predict_fiducial_mask(
     mask, probabilities = _predict_fiducial_mask(
         image=image,
         pixel_spacing=pixel_spacing,
-        mask_threshold=probability_threshold,
+        probability_threshold=probability_threshold,
         model_checkpoint_file=model_checkpoint_file,
     )
     mask = mask.cpu().numpy().astype(np.int8)

diff --git a/src/fidder/predict/predict.py b/src/fidder/predict/predict.py
@@ -15,7 +15,7 @@
 def predict_fiducial_mask(
     image: torch.Tensor,
     pixel_spacing: float,
-    mask_threshold: float,
+    probability_threshold: float,
     model_checkpoint_file: Optional[Path] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Predict fiducial masks for a batch of arbitrarily sized images.
@@ -55,7 +55,7 @@ def predict_fiducial_mask(
     [probabilities] = Trainer(accelerator="auto").predict(model, image)
     mask = probabilities_to_mask(
         probabilities=probabilities,
-        threshold=mask_threshold,
+        threshold=probability_threshold,
         connected_pixel_count_threshold=(PIXELS_PER_FIDUCIAL // 4),
     )