Caution

This repository isn't supported, maintained or tested and should only be used experimentally.

Morphollm Transformer for Neuron Morphological Data- Usage Examples

Install

In root of repo

pip install -r requirements.txt
pip install -e .

Download SWC Files

Load morphological data swc from e.g. NeuroMorpho.org:

mkdir -p data/neuromorpho

1. Quickstart - Training

Basics Training

python scripts/train.py \
    --data_dir data/neuromorpho \
    --epochs 100 \
    --batch_size 32 \
    --lr 0.0001

Training /w Config

python scripts/train.py \
    --config config/model_config.yaml \
    --data_dir data/neuromorpho

2. Inference

Single File

python scripts/inference.py \
    --model checkpoints/best_model.pt \
    --input data/neuromorpho/neuron_001.swc \
    --visualize

Batch Inference on dir

python scripts/inference.py \
    --model checkpoints/best_model.pt \
    --input data/neuromorpho/ \
    --output predictions.txt

3. Python API

Dataset & DataLoader

from src.data.neuromorpho_loader import NeuromorphoDataset, create_dataloader

dataset = NeuromorphoDataset(
    data_dir='data/neuromorpho',
    max_sections=1000,
    normalize=True,
    augment=True
)

sample = dataset[0]
print(f"Features shape: {sample['features'].shape}")
print(f"Filename: {sample['filename']}")

# DataLoader
train_loader = create_dataloader(
    data_dir='data/neuromorpho',
    batch_size=32,
    max_sections=1000,
    shuffle=True
)

# Iterate Batches
for batch in train_loader:
    features = batch['features']  # (batch_size, max_sections, feature_dim)
    mask = batch['mask']          # (batch_size, max_sections)
    break

Model create/use

from src.model.transformer import NeuromorphoTransformer
import torch

# Erstelle Modell
model = NeuromorphoTransformer(
    input_dim=13,
    d_model=256,
    nhead=8,
    num_encoder_layers=6,
    dim_feedforward=1024,
    dropout=0.1,
    num_classes=10
)

# Forward Pass
features = torch.randn(4, 100, 13)  # (batch, seq_len, features)
mask = torch.ones(4, 100, dtype=torch.bool)

class_logits, regression_output = model(features, mask)
print(f"Class logits: {class_logits.shape}")      # (4, 10)
print(f"Regression: {regression_output.shape}")   # (4, 4)

# Inference
predictions = model.predict(features, mask)
print(f"Predicted classes: {predictions['class_predictions']}")

Preprocessing, Augmentation

from src.data.preprocessing import NeuronPreprocessor, DataAugmentor
from src.data.neuromorpho_loader import SWCParser

parser = SWCParser()
df = parser.parse_swc_file('data/neuromorpho/neuron_001.swc')

preprocessor = NeuronPreprocessor()
features = preprocessor.compute_morphological_features(df)
print(f"Num sections: {features['num_sections']}")
print(f"Total cable length: {features['total_cable_length']:.2f}")

path_metrics = preprocessor.compute_path_lengths(df)
print(f"Mean segment length: {path_metrics['mean_segment_length']:.2f}")

df_clean = preprocessor.detect_outliers(df, threshold=3.0)

df_aligned = preprocessor.align_to_principal_axes(df)

augmentor = DataAugmentor()

df_rotated = augmentor.random_rotation(df, axes='xyz')

df_jittered = augmentor.add_jitter(df, noise_level=0.05)

df_scaled = augmentor.random_scale(df, scale_range=(0.8, 1.2))

Custom Training Loop

from src.model.transformer import NeuromorphoTransformer
from src.data.neuromorpho_loader import create_dataloader
import torch
import torch.nn as nn
import torch.optim as optim

# Setup
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Model
model = NeuromorphoTransformer(
    input_dim=13,
    d_model=256,
    nhead=8,
    num_encoder_layers=6
).to(device)

# Data
train_loader = create_dataloader(
    data_dir='data/neuromorpho',
    batch_size=32,
    max_sections=1000
)

# Optimizer und Loss
optimizer = optim.AdamW(model.parameters(), lr=1e-4)
classification_loss = nn.CrossEntropyLoss()
regression_loss = nn.MSELoss()

# Training Loop
model.train()
for epoch in range(10):
    for batch in train_loader:
        features = batch['features'].to(device)
        mask = batch['mask'].to(device)
        
        # Dummy Labels
        labels = torch.randint(0, 10, (features.size(0),)).to(device)
        reg_targets = torch.randn(features.size(0), 4).to(device)
        
        # Forward
        optimizer.zero_grad()
        class_logits, reg_output = model(features, mask)
        
        # Loss
        loss_class = classification_loss(class_logits, labels)
        loss_reg = regression_loss(reg_output, reg_targets)
        loss = loss_class + 0.5 * loss_reg
        
        # Backward
        loss.backward()
        optimizer.step()
    
    print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

4. Model load/save

from src.model.transformer import NeuromorphoTransformer

model.save_checkpoint(
    'checkpoints/my_model.pt',
    epoch=50,
    val_loss=0.234
)

loaded_model = NeuromorphoTransformer.load_from_checkpoint(
    'checkpoints/my_model.pt'
)
loaded_model.eval()

5. Visualization

from scripts.inference import NeuronPredictor
import matplotlib.pyplot as plt

predictor = NeuronPredictor(
    model_path='checkpoints/best_model.pt',
    device='cuda'
)

results = predictor.predict_from_file('data/neuromorpho/neuron_001.swc')
fig = predictor.visualize_neuron('data/neuromorpho/neuron_001.swc', results)
plt.savefig('neuron_viz.png', dpi=150)
plt.show()

6. Hyperparameter Tuning

config/model_config.yaml:

model:
  d_model: 512         
  nhead: 16             
  num_encoder_layers: 8 

training:
  batch_size: 16        
  learning_rate: 0.00005 

7. Multi-GPU Training

import torch.nn as nn

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)
    print(f"Using {torch.cuda.device_count()} GPUs")

model = model.to('cuda')