MLP Pipeline for Z-Stack Position Prediction

This work presents a machine learning pipeline for Z-stack position prediction in microscopy imaging, addressing the challenge of automated focus position prediction with micrometer accuracy. The system combines computer vision feature extraction techniques with multi-layered perceptrons to predict Z-position offsets from single microscopy images.

This pipeline is part of the EAP4EMSIG (Experiment Automation Pipeline for Event-Driven Microscopy to Smart Microfluidic Single-Cell Analysis) project, a collaborative research effort between Karlsruhe Institute of Technology (KIT) and Forschungszentrum Jülich (FZJ) to enhance event-driven microscopy for microfluidic single-cell analysis. For academic use, please cite: Friederich et al., "EAP4EMSIG - Enhancing Event-Driven Microscopy for Microfluidic Single-Cell Analysis", at – Automatisierungstechnik, De Gruyter, 2025 (In Press).

The approach leverages multi-resolution wavelet decomposition using biorthogonal wavelets (bior1.1, bior1.3) and geometric descriptors including image orientation and shape-based features. The pipeline processes high-resolution microscopy images (2560×2160 pixels) and extracts 135 features that capture texture, frequency domain characteristics, and spatial patterns indicative of focus quality. These features are fed into a multi-layer perceptron with dropout regularization, achieving typical accuracy of MAE < 0.105μm on real microscopy datasets. The system supports end-to-end processing from raw .nd2 files through feature extraction, data scaling, k-fold cross-validation, and model training, with built-in uncertainty estimation via Monte Carlo Dropout and comprehensive error analysis capabilities.

Key practical advantages include inference times under 100ms suitable for real-time pipelines, compatibility with standard laboratory computers without requiring high-end GPUs for training and deployment (complete workflow achievable in under a few hours), and accessibility to non-deep learning experts due to the simplicity of MLPs compared to sophisticated architectures like transformers. The approach is designed as an experiment-oriented solution rather than a universal model, requiring retraining when experimental setups change dramatically but offering robust performance within consistent imaging conditions.

Citation

If you use this pipeline in your research, please cite:

@article{friederich2025eap4emsig,
  title={EAP4EMSIG - Enhancing Event-Driven Microscopy for Microfluidic Single-Cell Analysis},
  author={Nils Friederich and Angelo Jovin Yamachui Sitcheu and Annika Nassal and Erenus Yildiz and Matthias Pesch and Maximilian Beichter and Lukas Scholtes and Bahar Akbaba and Thomas Lautenschlager and Oliver Neumann and Dietrich Kohlheyer and Hanno Scharr and Johannes Seiffarth and Katharina Nöh and Ralf Mikut},
  journal={at – Automatisierungstechnik},
  year={2025},
  publisher={De Gruyter},
  note={In Press},
  url={https://jugit.fz-juelich.de/emsig/mlp_pipeline}
}

Key Features

• End-to-End Pipeline: From raw .nd2 files to trained models
• High Precision: MAE < 0.105 μm prediction accuracy
• Advanced Features: Wavelet transforms with biorthogonal wavelets (bior1.1, bior1.3)
• Uncertainty Estimation: Monte Carlo Dropout for prediction confidence
• Batch Processing: Efficient handling of large microscopy datasets
• Comprehensive Analysis: Built-in error analysis and visualization tools
• Data Augmentation: Flip and rotation augmentation support
• Performance Optimized: CPU inference with timing analysis

Quick Start

1. Environment Setup

# Clone and setup virtual environment
cd mlp_venv
bash setup.sh
source activate.sh

2. Configure Your Data

Edit info.json to point to your .nd2 files:

{
    "data_directory": "/path/to/your/nd2/files/",
    "nd2_files": [
        {
            "name": "your_experiment_name",
            "step_size": 0.1,
            "use": true
        }
    ]
}

3. Run Complete Pipeline

# Run all stages (extract → features → scale → kfolds → train)
python main.py --suffix testmae0105_dense512dr512dr256dr128_test005_AugOn

# Or start from specific stage
python main.py --stage generate_features --suffix testmae0105_dense512dr512dr256dr128_test005_AugOn

4. Test Your Model

# Single image prediction
python predict_image_regressor.py --image-path /path/to/image.tiff --suffix testmae0105_dense512dr512dr256dr128_test005_AugOn

# Comprehensive model testing with uncertainty analysis
python test_model.py --model saved_models/regressor_testmae0105_dense512dr512dr256dr128_test005_AugOn.keras --uncertainty

Project Structure

mlp_pipeline/
├── main.py                        # Main pipeline orchestrator
├── predict_image_regressor.py     # Single image prediction + benchmarking
├── test_model.py                  # Comprehensive model testing + uncertainty
├── info.json                      # Configuration file
├── batch_runners/                 # Pipeline stage processors
│   ├── batch_extract_images.py
│   ├── batch_extract_features.py
│   ├── batch_scale_data.py
│   └── batch_prepare_kfolds.py
├── classes/                       # Core feature extraction & scaling
│   ├── features.py
│   └── scalers.py
├── tasks/                         # Model training implementations
├── utilities/                     # Helper functions and utilities
├── saved_models/                  # Trained models and scalers
├── error_analysis/                # Model performance analysis
├── generated_features/            # Extracted features (no augmentation)
├── generated_features_aug_on/     # Extracted features (with augmentation)
├── mlp_venv/                      # Virtual environment setup
├── explore_features.py            # Feature visualization tools
└── explore_stack.py               # Stack analysis tools

Pipeline Stages

The pipeline consists of 5 sequential stages that can be run individually or together:

Stage 1: Image Extraction

python main.py --stage extract_images

• Input: .nd2 files from configured data_directory
• Process: Extracts individual frame images from microscopy files
• Output: Organized image directories for each nd2 file

Stage 2: Feature Generation

python main.py --stage generate_features

• Input: Extracted images
• Process: Generates mathematical features using wavelet transforms and geometric descriptors
• Output: features.csv files with extracted feature vectors
• Note: This stage can take significant time depending on feature complexity

Stage 3: Data Scaling

python main.py --stage scale

• Input: Raw feature CSV files
• Process: Normalizes features using configured scaler (MinMax, Standard, etc.)
• Output: scaled_features.csv files + saved scaler model

Stage 4: K-Fold Preparation

python main.py --stage prepare_kfolds

• Input: Scaled features
• Process: Splits data into train/validation/test sets with k-fold cross-validation
• Output: Training splits ready for model training

Stage 5: Model Training

python main.py --stage train --suffix testmae0105_dense512dr512dr256dr128_test005_AugOn

• Input: Prepared k-fold data
• Process: Trains neural network regressor with specified architecture
• Output: Trained model saved as regressor_testmae0105_dense512dr512dr256dr128_test005_AugOn.keras

Feature Engineering

The pipeline extracts 135 features from microscopy images using wavelet transforms and geometric descriptors:

Feature Breakdown

• Wavelet Transforms (132 features):
  • Biorthogonal wavelets (bior1.1, bior1.3) with 22 decomposition levels each
  • Statistics: mean, standard deviation, entropy (3 stats × 22 levels × 2 wavelets = 132 features)
• Geometric Features (3 features):
  • Image orientation angle (1 feature)
  • Shape descriptors: row and column dimensions (2 features)

Total Features: 135

Active Feature Details

Wavelet Transform

Multi-resolution analysis using biorthogonal wavelets:

{
    "name": "wavelet",
    "params": {
        "wavelet": "bior1.1",
        "mode": "periodization", 
        "level": 7
    },
    "stat": ["mean", "std", "ent"],
    "active": true,
    "tag": "wvlt_bior_1_1"
}

• Active wavelets: bior1.1 and bior1.3 (22 levels each)
• Statistics: Mean, standard deviation, entropy (3 stats × 22 levels × 2 wavelets = 132 features)
• Edge handling: 2.5% pixel radius trimmed to avoid artifacts

Geometric Features

• Orientation: Image rotation angle analysis (1 feature)
• Resolution: Shape-based features (2 features: row and column dimensions)

Available But Inactive Features

The following features are available but currently disabled ("active": false) to optimize training time and maintain the current model performance. To enable them, set "active": true in info.json:

Gaussian Pyramid

Multi-scale image representation:

• Successive low-pass filtering and downsampling
• Captures information at different scales
• Configurable number of pyramid levels

Laplacian Pyramid

Edge and detail enhancement:

• Difference of Gaussian operations
• Highlights high-frequency content
• Excellent for capturing texture information

Log-Gabor Filters

Frequency domain analysis:

{
    "name": "log_gabor",
    "params": {
        "n_scale": 5,
        "n_orient": 6,
        "min_wavelength": 3.0,
        "mult": 1.6,
        "sigma_on_f": 0.75
    },
    "active": false,
    "tag": "log_gabor"
}

• Warning: Computationally expensive!
• Artifact handling: 20-pixel radius trimmed

DPFS Offset

• Custom positional features (currently inactive)

Adding Custom Features

1. Implement: Add new method to GenerateFeatures class in classes/features.py

def my_custom_feature(self, image, **params):
    # Your feature extraction logic
    return feature_array

2. Configure: Add entry to info.json:

{
    "name": "my_custom_feature",
    "params": {"param1": "value1"},
    "stat": ["mean", "std"],
    "active": true,
    "tag": "custom"
}

3. Activate: Set "active": true for training

Feature Column Naming Convention

Example: feature_wvlt_bior_1_1_mean_4

• feature_: Standard prefix
• wvlt_bior_1_1: Feature tag from config
• mean: Calculated statistic
• 4: Index in feature series

Note: The current info.json configuration may generate a different number of features than the published model depending on active feature settings. To reproduce exact published results, use the pre-generated feature files in generated_features/ directories.

Model Prediction & Analysis

Model-Dataset Compatibility

IMPORTANT: The pipeline automatically detects and warns about model-dataset compatibility issues. Models trained with different augmentation settings (flip/rotate) may produce inaccurate predictions if used with incompatible configurations.

How it works:

• Each model is automatically tested against both augmented and non-augmented datasets during inference
• The system displays which augmentation setting the model was trained with
• Compatibility warnings are shown when model training and current config don't match

Example compatibility output:

Detecting model training dataset...
Model was trained with: AUGMENTATION OFF
Current config uses: AUGMENTATION ON
COMPATIBILITY WARNING
Feature extraction settings don't match model training!
Predictions may be inaccurate.

To fix compatibility issues:

1. Update info.json data_augmentation settings to match the model, OR
2. Use a model trained with compatible augmentation settings

Single Image Prediction

Predict Z-position for individual images with timing analysis:

# Basic prediction (includes automatic compatibility checking)
python predict_image_regressor.py --image-path /path/to/image.tiff --suffix testmae0105_dense512dr512dr256dr128_test005_AugOn

# With detailed timing analysis
python predict_image_regressor.py --image-path /path/to/image.tiff --timing --suffix testmae0105_dense512dr512dr256dr128_test005_AugOn

Output Example:

CPU inference (optimized for single image predictions)
START: PREDICTING IMAGE
1/1 [==============================] - 0s 53ms/step
=== TIMING ANALYSIS ===
Preprocessing time: 9.24ms
Model inference time: 78.22ms

DONE: PREDICTING IMAGE

Predicted offset: 3.561255931854248

Comprehensive Model Testing

Advanced testing with uncertainty estimation using Monte Carlo Dropout:

# Basic model testing on all test sets (includes automatic compatibility detection)
python test_model.py --model saved_models/regressor_testmae0105_dense512dr512dr256dr128_test005_AugOn.keras

# Custom uncertainty parameters
python test_model.py --model saved_models/regressor_testmae0105_dense512dr512dr256dr128_test005_AugOn.keras --uncertainty 
    --mc-samples 50 --max-uncertainty-samples 100 --save-output 

Model-Dataset Compatibility: The testing script automatically uses the correct dataset directory (either generated_features or generated_features_aug_on) based on the model's training history, ensuring accurate performance evaluation.

Comprehensive Analysis Includes:

• Performance Metrics: MAE, RMSE, R² for each dataset
• Uncertainty Estimation: Monte Carlo Dropout confidence intervals
• Visualization: Error distribution plots, prediction vs uncertainty
• Timing Analysis: Preprocessing, inference, and total prediction times
• Per-Dataset Breakdown: Individual performance analysis

Uncertainty Estimation Features

• Monte Carlo Dropout: Sample predictions with dropout enabled
• Confidence Intervals: 95% prediction confidence bounds
• Uncertainty Quantification: Standard deviation of MC samples
• Risk Assessment: Identify high-uncertainty predictions

Configuration Reference

Core Settings (info.json)

{
    "data_directory": "/path/to/nd2/files/",
    "data_augmentation": {
        "flip": true,      // Enable horizontal/vertical flipping
        "rotate": true     // Enable rotation augmentation
    },
    "current": {
        "random_state": 1,     // Reproducibility seed
        "test_size": 0.05,     // Test set proportion (5%)
        "batch_size": 512,     // Training batch size
        "scaler": "min_max",   // Scaling method
        "epochs": 500          // Training epochs
    }
}

Data Augmentation

When enabled, creates separate feature directories:

• generated_features/: Original features
• generated_features_aug_on/: Augmented features

Scaling Options

• "min_max": MinMax scaler (0-1 range)
• "standard": Standard scaler (mean=0, std=1)
• "robust": Robust scaler (median-based)

Feature Exploration Tools

Individual Feature Visualization

Explore what specific features extract from images:

python explore_features.py --image-path /path/to/image.tiff \
    --directory /output/dir --feature wavelet

Stack-wise Feature Analysis

Analyze how features change across Z-stack:

python explore_stack.py --base-directory experiment_name \
    --stack stack_folder --output-directory /plots/output

Model Performance & Analysis

Automatic Error Analysis

Models are automatically evaluated with:

• Cross-validation performance during training
• Test set evaluation with comprehensive metrics
• Error distribution analysis saved in error_analysis/model_name/
• Visualization plots for prediction accuracy

Performance Benchmarks

Typical performance on microscopy data:

• Accuracy: Sub-micrometer precision (MAE = 0.105μm)
• Speed: ~80 ms inference time (CPU)
• Uncertainty: Reliable confidence estimation via MC Dropout

Advanced Features

Multi-Model Comparison

The pipeline supports multiple model configurations:

# Train different architectures with consistent naming: {performance}_{architecture}_{test_split}_{augmentation}
python main.py --suffix testmae019_dense64dr32dr16_test005_AugOff
python main.py --suffix testmae0105_dense512dr512dr256dr128_test005_AugOn
python main.py --suffix testmae0119_dense256dr256dr128_test005_AugOff

# Compare performance by running test on different models and collecting the analysis plots for each under their respective directories
python test_model.py --model saved_models/regressor_testmae019_dense64dr32dr16_test005_AugOff.keras
python test_model.py --model saved_models/regressor_testmae0105_dense512dr512dr256dr128_test005_AugOn.keras
python test_model.py --model saved_models/regressor_testmae0119_dense256dr256dr128_test005_AugOff.keras

Batch Processing Optimization

• Memory Efficient: Processes large datasets in configurable batches
• Progress Tracking: Built-in progress bars for long operations
• GPU Support: Automatic GPU detection and utilization for training
• CPU Fallback: Optimized CPU inference for inference/deployment

Troubleshooting

Common Issues

Import Errors

# Ensure virtual environment is activated
source mlp_venv/activate.sh
pip install -r mlp_venv/requirements.txt

GPU Performance Issues

For inference, CPU is way faster than GPU for single images. The code automatically forces CPU usage:

# This is already implemented in predict_image_regressor.py
tf.config.set_visible_devices([], 'GPU')  # Forces CPU usage

Feature Extraction Slow

• Disable expensive features in info.json (set "active": false)
• Log-Gabor filters are particularly compute-intensive!
• Consider reducing wavelet levels for faster processing

Low Model Accuracy

• Verify data quality and Z-position labels
• Experiment with different feature combinations
• Increase training epochs in info.json
• Check for data leakage between train/test sets

High Variance Datasets (different cells, chamber designs, illumination, etc.)

• Increase feature complexity: Raise wavelet levels or add new features to capture additional details
• Expand network architecture: Increase MLP neurons and widen the network
• Experiment extensively: Try different architectures but be careful of diminishing returns!
• Consider feature engineering: Design features specific to the experimental variations

Dependencies

Pip Requirements

numpy==1.26.1          # Numerical computing
scikit-learn==1.3.2    # Machine learning utilities  
tensorflow==2.14.0     # Deep learning framework
opencv-python==4.8.1.78 # Computer vision
pywavelets==1.4.1      # Wavelet transforms
seaborn==0.13.0        # Statistical visualization
tifffile==2023.8.30    # TIFF image handling
pims-nd2==1.1          # ND2 file reading
nd2==0.7.2             # Alternative ND2 support
tqdm==4.66.1           # Progress bars
requests==2.26.0       # HTTP utilities

Installation

All dependencies are automatically installed via the virtual environment setup:

cd mlp_venv
bash setup.sh          # Creates venv and installs dependencies
source activate.sh     # Activates environment

Architecture Overview

Neural Network Architecture

• Input Layer: Feature vector (size varies by configuration)
• Hidden Layers: Configurable dense layers with dropout
• Output Layer: Single neuron (regression output)
• Activation: ReLU for hidden layers, linear for output
• Regularization: Dropout layers for overfitting prevention

Feature Processing Pipeline

Raw Image → Feature Extraction → Scaling → Neural Network → Z-Position
     ↓              ↓               ↓            ↓
  2560x2160       135 dims       Normalized   Single Value

Supported Model Variants

The codebase supports multiple architectures with clear naming convention:

• testmae019_dense64dr32dr16_test005_AugOff: 64→32→16 neurons with dropout, MAE 0.19, 0.5% test split, no augmentation
• testmae0105_dense512dr512dr256dr128_test005_AugOn: 512→512→256→128 with dropout, MAE 0.105, 0.5% test split, with augmentation
• testmae0119_dense256dr256dr128_test005_AugOff: 256→256→128 with dropout, MAE 0.119, 0.5% test split, no augmentation
• testmae012_dense128dr64dr32_test005_AugOff: 128→64→32 with dropout, MAE 0.12, 0.5% test split, no augmentation
• testmae018_dense64dr32dr16_test015_AugOff: 64→32→16 with dropout, MAE 0.18, 1.5% test split, no augmentation
• Custom architectures via configuration

Naming Convention: {performance}_{architecture}_{test_split}_{augmentation}

• Performance: testmae019 = MAE 0.19, testmae0105 = MAE 0.105, etc.
• Architecture: dense{neurons} for single layer, dense{n1}dr{n2}dr{n3} for multi-layer with dropout
• Test Split: test005 = 0.5%, test015 = 1.5%, etc.
• Augmentation: AugOn = with data augmentation, AugOff = without augmentation

Development Workflow

Code Organization

• batch_runners/: Modular pipeline stages
• classes/: Core feature extraction and scaling
• tasks/: Training and model management
• utilities/: Shared helper functions

Adding New Features

1. Implement feature extraction in classes/features.py
2. Configure feature parameters in info.json
3. Test with exploration tools
4. Validate with model training

Model Experimentation

1. Train with different suffixes for comparison
2. Analyze using comprehensive testing tools
3. Save results in organized directory structure

Contributing

Guidelines

• Follow existing code style and documentation patterns
• Add comprehensive docstrings for new functions
• Test new features with exploration tools before integration
• Update info.json examples for new configuration options

Feature Requests

When adding new image features:

• Consider computational complexity
• Implement proper edge handling
• Add corresponding exploration capabilities
• Document mathematical background

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

• Built on TensorFlow and scikit-learn ecosystems
• Inspired by advances in computational microscopy
• Special thanks to all my colleagues from KIT and FZJ alike.

---

For questions or issues, please open an issue.