Segmentation Module
===================

The segmentation module provides a unified interface for automated cellular structure segmentation across multiple microscopy modalities, leveraging deep learning architectures optimized for each imaging technique.

Overview
--------

**Unified API Design**

iPA implements a consistent segmentation API through the ``BaseSegmenter`` abstract class and modality-specific implementations. This design enables:

* **Consistent Interface**: Same methods (``load_model()``, ``predict()``) across all segmenters
* **Modality-Specific Optimization**: Custom preprocessing and postprocessing for each imaging type
* **Automatic Device Selection**: GPU/CPU auto-detection with manual override option
* **Model Management**: Unified model loading, saving, and training interfaces

**Supported Modalities**:

+-------------+------------------------------------------+---------------------------+
| Modality    | Description                              | Available Tasks           |
+=============+==========================================+===========================+
| SXT         | Soft X-ray Tomography                    | Cell, Nucleus, Mito, ISG  |
+-------------+------------------------------------------+---------------------------+
| SIM         | Structured Illumination Microscopy       | ER, Mito, ISG             |
+-------------+------------------------------------------+---------------------------+
| WFM         | Widefield Microscopy                     | Cell, Nucleus, General    |
+-------------+------------------------------------------+---------------------------+
| Cryo-ET     | Cryo-Electron Tomography                 | Filaments, Membranes      |
+-------------+------------------------------------------+---------------------------+

Architecture
------------

The segmentation module follows a layered architecture:

1. **Base Layer** (``base.py``): Abstract ``BaseSegmenter`` class defining the common interface
2. **Unified Layer** (``unified.py``): Concrete implementations for each modality/task combination
3. **Specialized Layers**: Modality-specific submodules with custom models and algorithms
   - ``segmentation_sxt/``: SXT-specific U-Net models
   - ``segmentation_sim_wfm/``: SIM/WFM ERNet and other models
   - ``segmentation_et/``: Cryo-ET filament and membrane detection

Key Components
--------------

BaseSegmenter Class
~~~~~~~~~~~~~~~~~~~

.. autoclass:: ipa.processing.segmentation.BaseSegmenter
   :members:
   :undoc-members:
   :show-inheritance:

Unified Segmenter Classes
~~~~~~~~~~~~~~~~~~~~~~~~~

The following segmenter classes are available in ``ipa.processing.segmentation.unified``:

**SXT Segmenters**:
- ``SXTCellSegmenter`` - Cell and nucleus segmentation (U-Net, 3-class)
- ``SXTMitoSegmenter`` - Mitochondria segmentation (U-Net)
- ``SXTISGSegmenter`` - Insulin secretory granule segmentation (U-Net)
- ``SXTISGMaskRCNNSegmenter`` - ISG instance segmentation (Mask R-CNN) *[Training in progress]*

**SIM/WFM Segmenters**:
- ``SIMERSegmenter`` - Endoplasmic reticulum segmentation (ERNet)
- ``SIMMitoSegmenter`` - Mitochondria segmentation
- ``WFMSegmenter`` - General widefield segmentation

**Cryo-ET Segmenters**:
- ``ETFilamentSegmenter`` - Actin/microtubule filament detection
- ``ETMembraneSegmenter`` - Membrane surface segmentation

Model Architecture and Methods
-------------------------------

This section provides detailed information about the deep learning models and algorithms used in each segmentation task.

SXT Segmentation Models
~~~~~~~~~~~~~~~~~~~~~~~~

**SXT Cell & Nucleus Segmentation** (`SXTCellSegmenter`)

* **Model Architecture**: U-Net with dropout
* **Input**: 2D slices from 3D SXT volumes (resized to 288×480)
* **Output**: 3-class segmentation (background, membrane, nucleus)
* **Preprocessing**: 
  - Normalization to [0, 255] uint8
  - Resize to (288, 480) using nearest-neighbor interpolation
  - Standard normalization: ``(x / 255.0 - 0.456) / 0.224``
* **Postprocessing**: Slice-by-slice prediction, assembled into 3D volume
* **Training Data**: PBC Consortium pancreatic β-cell SXT datasets
* **Model File**: ``models/sxt/cell_nucleus_unet_best.pth``

**SXT Mitochondria Segmentation** (`SXTMitoSegmenter`)

* **Model Architecture**: U-Net (similar to cell segmentation)
* **Input**: 2D slices from 3D SXT volumes
* **Output**: Binary mitochondria mask
* **Preprocessing**: Same as cell segmentation
* **Model File**: ``models/sxt/mito_unet_best.pth``

**SXT ISG Segmentation** (`SXTISGSegmenter`)

* **Model Architecture**: U-Net for semantic segmentation
* **Input**: 3D SXT volumes
* **Output**: Binary ISG mask (semantic segmentation)
* **Use Case**: Identifying insulin secretory granules as a population
* **Model File**: ``models/sxt/isg_unet_best.pth``

**SXT ISG Instance Segmentation** (`SXTISGMaskRCNNSegmenter`) *[In Development]*

* **Model Architecture**: Mask R-CNN with ResNet-50 FPN backbone
* **Input**: 3D SXT volumes (slice-by-slice or volumetric)
* **Output**: Instance-level masks with bounding boxes and confidence scores
* **Advantages over U-Net**:
  - Distinguishes individual ISG instances
  - Provides per-instance morphological features
  - Enables ISG-ISG spatial relationship analysis
* **Status**: Currently training, expected completion TBD
* **Alternative**: Use ``SXTISGSegmenter`` + ``scipy.ndimage.label()`` for instance extraction

SIM/WFM Segmentation Models
~~~~~~~~~~~~~~~~~~~~~~~~~~~~

**SIM ER Segmentation** (`SIMERSegmenter`)

* **Model Architecture**: ERNet (specialized U-Net variant for ER structures)
* **Input**: SIM super-resolution images (~100 nm resolution)
* **Output**: ER network mask
* **Key Features**: Optimized for tubular network structures
* **Model File**: ``models/sim/ernet_final.pth``

**SIM Mitochondria Segmentation** (`SIMMitoSegmenter`)

* **Model Architecture**: U-Net based
* **Input**: SIM images
* **Output**: Mitochondria mask
* **Application**: High-resolution mitochondrial morphology analysis

**WFM General Segmentation** (`WFMSegmenter`)

* **Model Architecture**: Flexible U-Net framework
* **Input**: Widefield microscopy images
* **Output**: Configurable segmentation (cell, nucleus, organelles)
* **Flexibility**: Can be trained for various targets

Cryo-ET Segmentation Methods
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

**Cryo-ET Filament Detection** (`ETFilamentSegmenter`)

* **Method**: Skeletonization-based filament extraction
* **Input**: Cryo-ET tomograms with membrane/vesicle masks
* **Output**: Filament coordinates and lengths
* **Applications**: 
  - F-actin filament length measurement
  - Microtubule tracking
  - Cytoskeleton network analysis
* **Algorithm**: 
  1. Extract skeleton from binary mask
  2. Build KDTree for efficient neighbor search
  3. Calculate filament lengths by summing adjacent voxel distances

**Cryo-ET Membrane Segmentation** (`ETMembraneSegmenter`)

* **Method**: Membrane surface detection
* **Input**: Cryo-ET tomograms
* **Output**: Membrane surface coordinates
* **Applications**: Vesicle-membrane interaction analysis

Model Selection Guide
----------------------

Choosing the Right Model
~~~~~~~~~~~~~~~~~~~~~~~~~

**For SXT Data**:

+-------------------+------------------+----------------------------------+
| Task              | Recommended      | When to Use                      |
+===================+==================+==================================+
| Cell/Nucleus      | SXTCellSegmenter | Whole-cell and nuclear boundary  |
|                   |                  | detection                        |
+-------------------+------------------+----------------------------------+
| Mitochondria      | SXTMitoSegmenter | Mitochondrial morphology and     |
|                   |                  | distribution analysis            |
+-------------------+------------------+----------------------------------+
| ISG (population)  | SXTISGSegmenter  | ISG density, RDF analysis,       |
|                   |                  | docking ratio                    |
+-------------------+------------------+----------------------------------+
| ISG (instances)   | SXTISGMaskRCNN   | Individual ISG features, ISG-ISG |
|                   | *[Training]*     | interactions, instance-level     |
|                   |                  | statistics                       |
+-------------------+------------------+----------------------------------+

**For SIM/WFM Data**:

+-------------------+------------------+----------------------------------+
| Task              | Recommended      | When to Use                      |
+===================+==================+==================================+
| ER Network        | SIMERSegmenter   | ER morphology, ER-organelle      |
|                   |                  | contacts                         |
+-------------------+------------------+----------------------------------+
| Mitochondria      | SIMMitoSegmenter | High-res mitochondrial analysis  |
+-------------------+------------------+----------------------------------+
| Custom Targets    | WFMSegmenter     | Train on your specific organelle |
+-------------------+------------------+----------------------------------+

**For Cryo-ET Data**:

+-------------------+------------------+----------------------------------+
| Task              | Recommended      | When to Use                      |
+===================+==================+==================================+
| Filaments         | ETFilamentSeg.   | F-actin/microtubule length,      |
|                   |                  | cytoskeleton analysis            |
+-------------------+------------------+----------------------------------+
| Membranes         | ETMembraneSeg.   | Vesicle-membrane interactions,   |
|                   |                  | contact site analysis            |
+-------------------+------------------+----------------------------------+

Quick Start
-----------

Basic Usage Example:

.. code-block:: python

    from ipa.processing.segmentation import SXTCellSegmenter
    from ipa.data_loader import UniversalDataLoader
    
    # Load data
    volume = UniversalDataLoader.load_data('sxt_volume.mrc')
    
    # Initialize segmenter
    segmenter = SXTCellSegmenter(device='cuda')  # or 'cpu'
    
    # Load pre-trained model (auto-loads default model)
    segmenter.load_model()
    
    # Predict
    results = segmenter.predict(volume)
    
    # Access results
    cell_mask = results['cell_mask']
    nucleus_mask = results['nucleus_mask']
    
    print(f"Cell mask shape: {cell_mask.shape}")
    print(f"Nucleus mask shape: {nucleus_mask.shape}")

Submodules
----------

Detailed documentation for each modality:

.. toctree::
   :maxdepth: 2

   sxt_segmentation
   sim_wfm_segmentation
   et_segmentation