Experimental DeepD3: Incorporating Distance Constraints in Dendrite and Spine Segmentation

Project Overview

This project is an experimental study and extension of the original DeepD3 project. The primary aim is to integrate distance constraints between spine predictions and dendrite predictions in the original implementation DeepD3 project.

Objective

The main objective of this experimental study is to:

1. Implement a distance constraint mechanism between spine and dendrite predictions
2. Compare the performance of this experimental model with the original DeepD3 model
3. Evaluate the impact of the distance constraint on segmentation performance

Key Modifications

• Introduction of a Distance Transform Map (DTM) in the model architecture
• Implementation of a custom SpineDTMLayer to integrate the distance information
• Modification of the loss function to include a distance penalty

Features

• Modified DeepD3 neural network model with distance constraint
• Data generation and augmentation for training
• Inference pipeline for processing and comparing results
• Performance evaluation using Intersection over Union (IoU) metrics with respect to Benchmark dataset and Original DeepD3.
• Comparative analysis with the original DeepD3 model
• Visualization of comparative results

Installation

To set up the DeepD3 project, follow these steps:

1. Clone the repository

2. Install the required dependencies:

   pip install -r requirements.txt

   or

   conda create --name <env> --file requirements.txt

Project Structure

For the inference script to work correctly, organize your data in the following structure. And all the required datasets can be downloaded from DeepD3 website or we can use DeepD3 training method as mentioned here:

data_dir/
    models/
        Experimental_DeepD3.h5
        DeepD3_32F_94nm.h5
    input/
        DeepD3_Benchmark.tif
    ground_truth/
        Segmentation_U.mask
        Dendrite_U_dendrite.tif
        Segmentation_V.mask
        Dendrite_V_dendrite.tif
        Segmentation_W.mask
        Dendrite_W_dendrite.tif

Training

To train the experimental DeepD3 model:

1. Prepare your training data in the required format

2. Run the training script with desired parameters:

   python training.py --training_data path/to/training_data.d3set --validation_data path/to/validation_data.d3set --epochs 30 --learning_rate 0.0005 --distance_weight 0.5 --experiment_dir my_experiment

   Available arguments:

   • --training_data: Path to training data (default: "DeepD3_Training.d3set")
   • --validation_data: Path to validation data (default: "DeepD3_Validation.d3set")
   • --epochs: Number of epochs to train (default: 30)
   • --learning_rate: Initial learning rate (default: 0.0005)
   • --distance_weight: Weight for distance penalty in loss function (default: 0.5)
   • --batch_size: Batch size for training (default: 32)
   • --experiment_dir: Directory to save experiment results (default: "new_experiment")

Inference and Comparison

To run inference and compare results with the original DeepD3 model:

1. Organize your data according to the project structure described above

2. Run the inference script:

   python inference.py --data_dir path/to/your/data_directory --save_predictions True

   Available arguments:

   • --data_dir: Path to the base directory containing models, input, and ground truth data (required)
   • --save_predictions: Flag to save model predictions (optional)

The script will automatically:

• Load the experimental and original models
• Process the input image stack
• Perform inference using both models
• Calculate IoU scores for both models
• Generate and save a comparison plot
• Save predictions if the --save_predictions flag is used

Model Architecture

The experimental DeepD3 model is inspired by the original DeepD3's U-Net-like architecture with these key additions:

• SpineDTMLayer: A custom layer that integrates the Distance Transform Map
• Modified decoder output to include both spine predictions and DTM information

Here's a visual representation of our modified DeepD3 architecture:

Key components of the architecture:

1. Encoder: Convolutional layers with residual connections
2. Decoder: Upsampling and concatenation with skip connections
3. SpineDTMLayer: Custom layer for integrating distance information
4. Dual output: Dendrite segmentation and Spine segmentation with DTM

Evaluation

The project includes functions for calculating Intersection over Union (IoU) scores to compare the performance of the experimental model against the original DeepD3 model.

Results

In this experiment, we evaluated the experimental model with the original DeepD3 on the available Benchmark dataset. The results we got after the evaluation don't show any drastic improvement on existing results. These results show that integrating Distance Transform Map (DTM) doesn't make any significant improvement to the existing model.

Future Work

As we have evaluated different methods such as Distance Transform Map (DTM) and Signed Distance Function (SDF), these methods do not encapsulate the required distance constraint into the Segmentation task. Also, on the usage of U-Net-like architecture, we think if we can use modern architectures like Transformer or MLP Mixer, which work on patch embeddings and attention mechanisms, we will be able to provide more granular control on the integration of these distance constraints.

Contributing

Contributions to this experimental DeepD3 project are welcome. Please follow these steps:

1. Fork the repository
2. Create a new branch for your feature or experiment
3. Commit your changes
4. Push to your branch
5. Create a pull request with a detailed description of your modifications and their impact

Acknowledgements

This project builds upon the original DeepD3 project. We acknowledge their contributions to the field of neuronal image segmentation.

Contact Information

Munna Prithvinath Singh
Friedrich Alexander University, Erlangen-Nürnberg
mailto:munna.singh@fau.de

Prof. Dr. Andreas Kist
Friedrich Alexander University, Erlangen-Nürnberg