Deconvolution as a preprocessing option

README.md

@@ -38,8 +38,7 @@ Explore MemBrain-seg, use it for your needs, and let us know how it works for yo
 Preliminary [documentation](https://github.com/teamtomo/membrain-seg/blob/main/docs/index.md) is available, but far from perfect. Please let us know if you encounter any issues, and we are more than happy to help (and get feedback what does not work yet).
 
 ```
-[1] Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2020). nnU-Net: a self-configuring method 
-for deep learning-based biomedical image segmentation. Nature Methods, 1-9.
+[1] Isensee, F., Jaeger, P.F., Kohl, S.A.A., Petersen, J., Maier-Hein, K.H., 2021. nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods 18, 203-211. https://doi.org/10.1038/s41592-020-01008-z
 ```
 
 # Installation
@@ -52,8 +51,9 @@ Please find more detailed instructions [here](./docs/Usage/Segmentation.md).
 
 ## Preprocessing
 Currently, we provide the following two [preprocessing](https://github.com/teamtomo/membrain-seg/tree/main/src/tomo_preprocessing) options:
-- pixel size matching: Rescale your tomogram to match the training pixel sizes
+- Pixel size matching: Rescale your tomogram to match the training pixel sizes
 - Fourier amplitude matching: Scale Fourier components to match the "style" of different tomograms
+- Deconvolution: denoises the tomogram by applying the deconvolution filter from Warp
 
 For more information, see the [Preprocessing](./docs/Usage/Preprocessing.md) subsection.
 

src/membrain_seg/tomo_preprocessing/readme.md → src/membrain_seg/tomo_preprocessing/README.md

@@ -29,6 +29,7 @@ already some rules of thumb:
 2. The Fourier amplitude matching only works in some cases, depending on the CTFs of input 
 and target tomograms. Our current recommendation is: If you're not satisfied with MemBrain's 
 segmentation performance, why not give the amplitude matching a shot?
+3. Deconvolution can help with the segmentation performance if your tomogram has not already been denoised somehow (e.g. using cryo-CARE, IsoNet or Warp). Deconvolving an already denoised tomogram is not recommended, it will most likely make things worse.
 
 More detailed guidelines are in progress!
 
@@ -53,7 +54,8 @@ For help on a specific command, use:
 `tomo_preprocessing extract_spectrum --input_path <path-to-tomo> --output_path <path-to-output>`
 - **match_spectrum**: Match amplitude of Fourier spectrum from input tomogram to target spectrum. Example:  
 `tomo_preprocessing match_spectrum --input <path-to-tomo> --target <path-to-spectrum> --output <path-to-output>`
-
+- **deconvolve**: Denoises the tomogram by deconvolving the contrast transfer function. Example:  
+`tomo_preprocessing deconvolve --input <path-to-tomo> --output <path-to-output-tomo> --df <defocus-value>`
 
 ### **Pixel Size Matching**
 Pixel size matching is recommended when your tomogram pixel sizes differs strongly from the training pixel size range (roughly 10-14&Aring;). You can perform it using the command
@@ -78,4 +80,15 @@ Fourier amplitude matching is performed in two steps:
 This extracts the radially averaged Fourier spectrum and stores it into a .tsv file.
 2. Matching of the input tomogram to the extracted spectrum:  
 `tomo_preprocessing match_spectrum --input <path-to-tomo> --target <path-to-spectrum> --output <path-to-output>`  
-Now, the input tomograms Fourier components are re-scaled based on the equalization kernel computed from the input tomogram's radially averaged Fourier intensities, and the previously extracted .tsv file.
+Now, the input tomograms Fourier components are re-scaled based on the equalization kernel computed from the input tomogram's radially averaged Fourier intensities, and the previously extracted .tsv file.
+
+### **Deconvolution**
+
+Deconvolution is a denoising method that works by "removing" the effects of the contrast transfer function (CTF) from the tomogram. This is based on an ad-hoc model of the spectral signal-to-noise-ratio (SSNR) in the data, following the implementation in the Warp package [1]. Effectively what the filter does is to boost the very low frequencies, thus enhancing the tomogram contrast, while low-pass filtering beyond the first zero-crossing of the CTF.
+For the filter to work, you need to provide the CTF parameters, namely a defocus value for the tomogram, as well as the acceleration voltage, spherical aberration and amplitude contrast, if those differ from the defaults. This is typically the defocus value of the zero tilt. It does not need to be super accurate, a roughly correct value already produces decent results. While the defaults usually work well, you can play with the filter parameters, namely the deconvolution strength and the falloff, to fine-tune the results.
+Example detailed command:
+`tomo_preprocessing deconvolve --input <path-to-tomo> --output <path-to-output-tomo> --df 45000 --ampcon 0.07 --cs 2.7 --kv 300 --strength 1.0 --falloff 1.0`
+
+```
+[1] Tegunov, D., Cramer, P., 2019. Real-time cryo-electron microscopy data preprocessing with Warp. Nature Methods 16, 1146–1152. https://doi.org/10.1038/s41592-019-0580-y
+```

src/membrain_seg/tomo_preprocessing/__init__.py

@@ -7,4 +7,5 @@
 # These imports are necessary to register CLI commands. Do not remove!
 from .amplitude_spectrum_matching._cli import extract, match_spectrum  # noqa: F401
 from .cli import cli  # noqa: F401
+from .deconvolution._cli import deconvolve  # noqa: F401
 from .pixel_size_matching._cli import match_pixel_size, match_seg_to_tomo  # noqa: F401

src/membrain_seg/tomo_preprocessing/deconvolution/__init__.py

@@ -0,0 +1 @@
+"""Empty init."""

src/membrain_seg/tomo_preprocessing/deconvolution/_cli.py

@@ -0,0 +1,87 @@
+from typer import Option
+
+from ..cli import OPTION_PROMPT_KWARGS as PKWARGS
+from ..cli import cli
+from .deconvolve import deconvolve as run_deconvolve
+
+
+@cli.command(name="deconvolve", no_args_is_help=True)
+def deconvolve(
+    input: str = Option(  # noqa: B008
+        None, help="Tomogram to deconvolve (.mrc/.rec format)", **PKWARGS
+    ),
+    output: str = Option(  # noqa: B008
+        None,
+        help="Output location for deconvolved tomogram (.mrc/.rec format)",
+        **PKWARGS,
+    ),
+    pixel_size: float = Option(  # noqa: B008
+        None,
+        help="Input pixel size (optional). If not specified, it will be read from the \
+tomogram's header. ATTENTION: This can lead to severe errors if the header pixel size \
+is not correct.",
+    ),
+    df: float = Option(  # noqa: B008
+        50000,
+        help="The defocus value to be used for deconvolution, in Angstroms. This is \
+typically the defocus of the zero tilt. Underfocus is positive.",
+        **PKWARGS,
+    ),
+    # df2: float = Option(
+    #     None,
+    #     help="Defocus 2 (or Defocus V in some notations) in Angstroms. Defocus axis \
+    # orthogonal to the U axis. Only mandatory for astigmatic data.",
+    # ),
+    # ast: float = Option(
+    #     0.0,
+    #     help="Angle for astigmatic data (in degrees). Astigmatism is currently not \
+    # used in deconvolution (only the axis of largest defocus is considered), but maybe\
+    # some better model in the future will use it?",
+    # ),
+    ampcon: float = Option(  # noqa: B008
+        0.07,
+        help="Amplitude contrast fraction (between 0.0 and 1.0).",
+    ),
+    cs: float = Option(  # noqa: B008
+        2.7,
+        help="Spherical aberration (in mm).",
+    ),
+    kv: float = Option(  # noqa: B008
+        300.0,
+        help="Acceleration voltage of the TEM (in kV).",
+    ),
+    strength: float = Option(  # noqa: B008
+        1.0,
+        help="Strength parameter for the denoising filter.",
+    ),
+    falloff: float = Option(  # noqa: B008
+        1.0,
+        help="Falloff parameter for the denoising filter.",
+    ),
+    hp_fraction: float = Option(  # noqa: B008
+        0.02,
+        help="Fraction of Nyquist frequency to be cut off on the lower end (since it \
+will be boosted the most)",
+    ),
+    skip_lowpass: bool = Option(  # noqa: B008
+        False,
+        help="The denoising filter by default will have a smooth low-pass effect that \
+enforces filtering out any information beyond the first zero of the CTF. Use this \
+option to skip this filter i.e. potentially include information beyond the first CTF \
+zero (not recommended).",
+    ),
+):
+    """Deconvolve the input tomogram using the Warp deconvolution filter."""
+    run_deconvolve(
+        input,
+        output,
+        df,
+        ampcon,
+        cs,
+        kv,
+        pixel_size,
+        strength,
+        falloff,
+        hp_fraction,
+        skip_lowpass,
+    )