From be0c63795626297714dd63412e991b19a491c850 Mon Sep 17 00:00:00 2001
From: Talon Chandler <talonchandler@gmail.com>
Date: Fri, 17 May 2024 11:53:50 -0700
Subject: [PATCH] first-pass characterize

---
 mantis/analysis/AnalysisSettings.py | 16 ++++++
 mantis/analysis/analyze_psf.py      |  9 +++-
 mantis/cli/characterize.py          | 79 +++++++++++++++++++++++++++++
 mantis/cli/main.py                  |  2 +
 4 files changed, 105 insertions(+), 1 deletion(-)
 create mode 100644 mantis/cli/characterize.py

diff --git a/mantis/analysis/AnalysisSettings.py b/mantis/analysis/AnalysisSettings.py
index edc63ce3..e8833ad2 100644
--- a/mantis/analysis/AnalysisSettings.py
+++ b/mantis/analysis/AnalysisSettings.py
@@ -1,6 +1,7 @@
 from typing import Literal, Optional, Union
 
 import numpy as np
+import torch
 
 from pydantic import BaseModel, Extra, NonNegativeInt, PositiveFloat, PositiveInt, validator
 
@@ -77,3 +78,18 @@ class PsfFromBeadsSettings(MyBaseModel):
 
 class DeconvolveSettings(MyBaseModel):
     regularization_strength: PositiveFloat = 0.001
+
+
+class CharacterizeSettings(MyBaseModel):
+    block_size: list[NonNegativeInt] = (64, 64, 32)
+    blur_kernel_size: NonNegativeInt = 3
+    nms_distance: NonNegativeInt = 32
+    min_distance: NonNegativeInt = 50
+    threshold_abs: PositiveFloat = 200.0
+    max_num_peaks: NonNegativeInt = 2000
+    exclude_border: list[NonNegativeInt] = (5, 10, 5)
+    device: str = "cuda"
+
+    @validator("device")
+    def check_device(cls, v):
+        return "cuda" if torch.cuda.is_available() else "cpu"
diff --git a/mantis/analysis/analyze_psf.py b/mantis/analysis/analyze_psf.py
index 378733d3..c354248e 100644
--- a/mantis/analysis/analyze_psf.py
+++ b/mantis/analysis/analyze_psf.py
@@ -150,7 +150,14 @@ def generate_report(
 def extract_beads(
     zyx_data: ArrayLike, points: ArrayLike, scale: tuple, patch_size_voxels: tuple = None
 ):
-    patch_size = (scale[0] * patch_size_voxels[0], scale[1] * patch_size_voxels[1], scale[2] * patch_size_voxels[2])
+    if patch_size_voxels is None:
+        patch_size = (scale[0] * 15, scale[1] * 18, scale[2] * 18)
+    else:
+        patch_size = (
+            scale[0] * patch_size_voxels[0],
+            scale[1] * patch_size_voxels[1],
+            scale[2] * patch_size_voxels[2],
+        )
 
     # extract bead patches
     bead_extractor = BeadExtractor(
diff --git a/mantis/cli/characterize.py b/mantis/cli/characterize.py
new file mode 100644
index 00000000..4d80b4cf
--- /dev/null
+++ b/mantis/cli/characterize.py
@@ -0,0 +1,79 @@
+import click
+import time
+import gc
+import torch
+import warnings
+
+from iohub.ngff import open_ome_zarr
+from typing import List
+from mantis.analysis.AnalysisSettings import CharacterizeSettings
+from mantis.cli.parsing import input_position_dirpaths, output_dirpath, config_filepath
+from mantis.cli.utils import yaml_to_model
+from mantis.analysis.analyze_psf import detect_peaks, extract_beads, analyze_psf, generate_report
+
+
+@click.command()
+@input_position_dirpaths()
+@config_filepath()
+@output_dirpath()
+def characterize(
+    input_position_dirpaths: List[str],
+    config_filepath: str,
+    output_dirpath: str,
+):
+    """
+    Characterize the point spread function (PSF) from bead images in an html report
+
+    >> mantis characterize -i ./beads.zarr/*/*/* -c ./characterize_params.yml -o ./
+    """
+    click.echo(f"Loading data...")
+    with open_ome_zarr(str(input_position_dirpaths[0]), mode="r") as input_dataset:
+        T, C, Z, Y, X = input_dataset.data.shape
+        zyx_data = input_dataset["0"][0, 0]
+        zyx_scale = input_dataset.scale[-3:]
+
+    # Read settings
+    settings = yaml_to_model(config_filepath, CharacterizeSettings)
+    
+    click.echo(f"Detecting peaks...")
+    t1 = time.time()
+    peaks = detect_peaks(
+        zyx_data,
+        **settings.dict(),
+        verbose=True,
+    )
+    gc.collect()
+    torch.cuda.empty_cache()
+    t2 = time.time()
+    click.echo(f'Time to detect peaks: {t2-t1}')
+
+    t1 = time.time()
+    beads, offsets = extract_beads(
+        zyx_data=zyx_data,
+        points=peaks,
+        scale=zyx_scale,
+    )
+
+    click.echo(f"Analyzing PSFs...")
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        df_gaussian_fit, df_1d_peak_width = analyze_psf(
+            zyx_patches=beads,
+            bead_offsets=offsets,
+            scale=zyx_scale,
+        )
+    t2 = time.time()
+    click.echo(f'Time to analyze PSFs: {t2-t1}')
+
+    # Generate HTML report
+    generate_report(
+        output_dirpath,
+        input_position_dirpaths[0],
+        "",
+        beads,
+        peaks,
+        df_gaussian_fit,
+        df_1d_peak_width,
+        zyx_scale,
+        ["AXIS 0", "AXIS 1", "AXIS 2"],
+    )
\ No newline at end of file
diff --git a/mantis/cli/main.py b/mantis/cli/main.py
index 3d0a736e..290f5e9e 100644
--- a/mantis/cli/main.py
+++ b/mantis/cli/main.py
@@ -10,6 +10,7 @@
 from mantis.cli.update_scale_metadata import update_scale_metadata
 from mantis.cli.psf_from_beads import psf_from_beads
 from mantis.cli.deconvolve import deconvolve
+from mantis.cli.characterize import characterize
 
 CONTEXT = {"help_option_names": ["-h", "--help"]}
 
@@ -35,3 +36,4 @@ def cli():
 cli.add_command(update_scale_metadata)
 cli.add_command(psf_from_beads)
 cli.add_command(deconvolve)
+cli.add_command(characterize)