openmedaxis 0.0.6

Open framework for generative medical imaging research

OpenMedAxis

PyTorch Lightning for Generative Medical Imaging.

OpenMedAxis is an open-source framework designed to make generative medical imaging research reproducible, modular, and easy to experiment with. It provides clean implementations of datasets, generative models, training pipelines, and evaluation tools commonly used in medical imaging research.

The goal is to create a researcher-first framework that removes boilerplate and fragmented codebases while keeping the system low-level, flexible, and hackable.

---

Why OpenMedAxis?

Research in generative medical imaging often suffers from the same problems:

• Dataset preprocessing scripts are scattered across different repositories
• Training pipelines are inconsistent and difficult to reproduce
• Implementations of GANs or diffusion models vary widely
• Evaluation pipelines are rarely standardized

OpenMedAxis aims to solve this by providing a clean, unified research toolkit for the community.

Instead of every paper reinventing the same infrastructure, OpenMedAxis provides:

• standardized dataset pipelines
• reference implementations of generative models
• modular training components
• reproducible experiment structure

---

Key Features

Medical Imaging Dataset Tools

Utilities and standardized loaders for common datasets.

Planned initial support:

• IXI
• BraTS
• fastMRI

Each dataset module will include:

• download utilities
• preprocessing scripts
• standardized dataset interfaces

---

Generative Model Zoo

Reference implementations of commonly used generative models in medical imaging.

Planned models:

GAN-based models

• DCGAN
• pix2pix
• CycleGAN
• SAGAN

Diffusion-based models

• DDPM
• DDIM
• Latent Diffusion

Future additions may include:

• Flow Matching

---

Modular Training Framework

A lightweight and extensible training framework inspired by modern deep learning libraries.

Features include:

• experiment configuration
• modular trainer
• logging and checkpointing
• reproducible experiments

Researchers can easily swap components such as:

• models
• optimizers
• schedulers
• loss functions

---

Evaluation and Metrics

OpenMedAxis will provide standardized evaluation tools for generative medical imaging.

Planned metrics include:

• PSNR
• SSIM
• MS-SSIM
• FID
• LPIPS

Future work will include domain-specific metrics for medical image realism and structure preservation.

---

Research Utilities

Tools designed to reduce repetitive boilerplate code.

Examples include:

• preprocessing pipelines
• visualization tools
• slice viewers
• experiment templates
• training utilities

---

Installation

OpenMedAxis provides two ways to get started:

---

⚡ Recommended: Project Scaffolding (Best Experience)

First install uv:

curl -LsSf https://astral.sh/uv/install.sh | sh

Create a fully configured research project with one command:

uvx openmedaxis-init my_project

This will:

• create a new project
• set up a virtual environment
• install OpenMedAxis
• install required dependencies
• configure PyTorch backend

Then:

cd my_project
uv run python -m my_project.train

---

Backend Options

By default, OpenMedAxis uses CUDA 11.8 (cu118).

You can switch backend:

# CPU-only
uvx openmedaxis-init my_project --torch-backend cpu

# CUDA 11.8 (default)
uvx openmedaxis-init my_project --torch-backend cu118

---

🧠 Advanced: Manual Installation

If you prefer full control over your environment:

pip install openmedaxis

Then install full dependency stack:

pip install "openmedaxis[full]"

⚠️ PyTorch must be installed separately depending on your system:

# CPU
pip install torch torchvision torchaudio

# CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

---

Development Installation

Clone the repository:

git clone https://github.com/AllAxisAI/openmedaxis.git
cd openmedaxis

Install in editable mode:

pip install -e .

---

Currently the project is in pre-alpha development.

---

Repository Structure

openmedaxis/

├── README.md
├── LICENSE
├── pyproject.toml
│
├── docs/
├── examples/
├── configs/
├── scripts/
├── tests/
│
└── openmedaxis/

    ├── datasets/
    ├── models/
    │   ├── gan/
    │   └── diffusion/
    │
    ├── training/
    ├── evaluation/
    ├── transforms/
    ├── visualization/
    └── utils/

---

Example Vision

The framework aims to make research experiments simple and reproducible.

Example usage (future API):

from openmedaxis.datasets import IXI
from openmedaxis.models.diffusion import DDPM
from openmedaxis.training import Trainer


dataset = IXI(root="data/", modality="T2")
model = DDPM()

trainer = Trainer(model=model, dataset=dataset)
trainer.fit()

---

Roadmap

Phase 1 — Foundation

• project architecture
• dataset interfaces
• IXI dataset loader
• BraTS dataset loader
• fastMRI dataset loader
• DCGAN baseline
• pix2pix baseline
• DDPM baseline
• basic evaluation metrics

Phase 2 — Expansion

• CycleGAN and SAGAN
• improved evaluation tools
• visualization utilities
• configuration system

Phase 3 — Advanced Generative Methods

• latent diffusion
• 3D generative models
• efficient training pipelines
• community benchmark suite

---

Contributing

OpenMedAxis is at an early stage and contributions are welcome.

Areas where contributions are especially valuable:

• dataset integrations
• generative model implementations
• evaluation metrics
• documentation and examples

If you are interested in contributing, please open an issue to discuss ideas or improvements.

---

License

This project will be released under an open-source license (MIT or Apache-2.0).

---

Organization

OpenMedAxis is developed under AllAxisAI.

---

Status

🚧 Pre-alpha — initial framework design and core modules under development.