KapoorLabs-Lightning 6.0.2

Lightning modules for KapoorLabs specific projects

KapoorLabs-Lightning

Developed by KapoorLabs

This product is a testament to our expertise at KapoorLabs, where we specialize in creating cutting-edge solutions. We offer bespoke pipeline development services, transforming your developmental biology questions into publishable figures with our advanced computer vision and AI tools. Leverage our expertise and resources to achieve end-to-end solutions that make your research stand out.

Note: The tools and pipelines showcased here represent only a fraction of what we can achieve. For tailored and comprehensive solutions beyond what was done in the referenced publication, engage with us directly. Our team is ready to provide the expertise and custom development you need to take your research to the next level. Visit us at KapoorLabs.

Lightning Modules for KapoorLabs Projects

PyTorch Lightning framework for training deep learning models on microscopy data, with specialized support for:

• ONEAT: Spatio-temporal event detection in 3D+T microscopy data
• Cell Fate Classification: Time series classification of cell fates (basal, goblet, radial) from tracking data

---

Key Features

• Modular Architecture: Base, ONEAT, and Cell Fate Lightning modules
• YOLO-style Detection: VolumeYoloLoss for multi-task learning (classification + localization)
• H5 Dataset Support: Memory-efficient streaming from HDF5 files with YOLO labels
• Segmentation-Guided Prediction: Uses instance segmentation to locate cells, carves patches from raw image, classifies each cell, and records global coordinates for positive events
• Transform Presets: Light, Medium, Heavy augmentation pipelines for microscopy data
• Multiple Optimizers: Adam, SGD, LARS, AdamW with learning rate schedulers
• SLURM Integration: Auto-requeue support for HPC clusters
• Hydra Configuration: YAML-based experiment configuration

Package Structure

kapoorlabs_lightning/
├── Lightning Modules
│   ├── base_module.py          # BaseModule - common functionality
│   ├── oneat_module.py         # OneatActionModule - event detection
│   └── cellfate_module.py      # CellFateModule - time series classification
├── Models
│   ├── pytorch_models.py       # DenseVollNet, DenseNet, InceptionNet
│   └── pytorch_losses.py       # VolumeYoloLoss, OneatClassificationLoss
├── Data
│   ├── pytorch_datasets.py        # H5VisionDataset, H5MitosisDataset
│   └── oneat_prediction_dataset.py # OneatPredictionDataset (seg-guided inference)
├── Transforms
│   ├── oneat_transforms.py     # Microscopy-specific augmentations
│   ├── oneat_presets.py        # Light/Medium/Heavy presets
│   └── time_series_presets.py  # Cell fate transforms + presets (order-preserving)
├── Training
│   ├── lightning_trainer.py    # MitosisInception trainer class
│   ├── optimizers.py           # Adam, SGD, LARS, AdamW
│   └── schedulers.py           # Cosine, WarmCosine, Step
├── Utilities
│   ├── utils.py                # H5 creation, normalization, plotting
│   ├── nms_utils.py            # Space-time NMS
│   └── pytorch_callbacks.py    # Checkpointing, progress bars
└── Logging
    └── pytorch_loggers.py      # CustomNPZLogger for metrics

Installation

You can install KapoorLabs-Lightning via pip:

pip install KapoorLabs-Lightning

To install latest development version :

pip install git+https://github.com/Kapoorlabs-CAPED/KapoorLabs-Lightning.git

Documentation

• ONEAT Training Guide - Complete workflow for event detection
• Cell Fate Classification Guide - Time series cell fate classification
• Lightning Modules - Module architecture details

Quick Start

ONEAT Event Detection

from kapoorlabs_lightning import MitosisInception

# Initialize trainer
trainer = MitosisInception(
    h5_file="training_data.h5",
    num_classes=2,
    epochs=100,
    batch_size=32,
    learning_rate=1e-3,
)

# Setup model and training
trainer.setup_densenet_vision_model(
    input_shape=(3, 8, 64, 64),  # (T, Z, Y, X)
    categories=2,
    box_vector=8,
)
trainer.setup_oneat_transforms_medium()
trainer.setup_vision_h5_datasets()
trainer.setup_adam()
trainer.setup_oneat_lightning_model()
trainer.train()

Cell Fate Classification

from kapoorlabs_lightning import MitosisInception

# Initialize trainer
trainer = MitosisInception(
    h5_file="cellfate_data.h5",  # H5 with train_arrays/train_labels/val_arrays/val_labels
    num_classes=3,               # e.g. basal, goblet, radial
    epochs=250,
    batch_size=64,
    learning_rate=1e-3,
    seq_len=25,                  # 25 timepoints per track
)

# Setup (no temporal order changes in transforms)
trainer.setup_cellfate_transforms_medium()
trainer.setup_gbr_h5_datasets()
trainer.setup_inception_qkv_model()
trainer.setup_adam()
trainer.setup_cellfate_lightning_model()
trainer.train()

Contributing

Contributions are very welcome. Tests can be run with tox, please ensure the coverage at least stays the same before you submit a pull request.

License

Distributed under the terms of the BSD-3 license, "KapoorLabs-Lightning" is free and open source software

Issues

If you encounter any problems, please file an issue along with a detailed description.