MrBlankness/LightM-UNet

Pytorch implementation of "LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation"

What it solves

LightM-UNet addresses the need for high-performance medical image segmentation models that are computationally efficient. It aims to provide a lightweight alternative to state-of-the-art models like nnU-Net and U-Mamba, which often have significantly larger parameter counts.

How it works

The project implements a lightweight fusion of the UNet architecture and Mamba. It is built on the nnU-Net framework, allowing it to handle both 2D and 3D real-world medical datasets. By integrating Mamba, the model achieves high segmentation accuracy while maintaining a very small footprint.

Who it’s for

This tool is designed for researchers and developers in the medical imaging field who need efficient models for segmenting 2D and 3D medical images without requiring massive computational resources.

Highlights

  • Extremely Lightweight: Features a parameter count of only 1M.
  • High Efficiency: Reduces parameter counts by 116X compared to nnU-Net and 224X compared to U-Mamba.
  • Versatile: Supports both 2D and 3D medical image segmentation.
  • Framework Integration: Built on the popular nnU-Net framework for data preprocessing and training.

Related

  • Project
    kyegomez/MultiModalMamba

    A multi-modal AI model that integrates Vision Transformer (ViT) and Mamba to efficiently process and interpret text, images, audio, and video data simultaneously.

  • Project
    JoHof/lungmask

    A package of trained U-net models for automated lung segmentation in CT scans, capable of handling severe pathologies and lobe-specific masking.

  • Project
    state-spaces/mamba

    Mamba is a linear-time sequence modeling architecture based on selective state space models (SSMs) that provides a high-performance, subquadratic alternative to Transformers.

  • Project
    UbiquitousLearning/mllm

    A lightweight multimodal LLM inference engine optimized for mobile and edge devices, supporting various NPUs and GPUs with advanced quantization.

  • Project
    lmcinnes/umap

    A fast non-linear dimension reduction library for visualizing high-dimensional data and preprocessing for machine learning tasks.