MONAI Breast Density Classification

A TorchServe-based medical imaging model for automated breast density classification from mammogram images using deep learning. This model classifies breast density into the standard BI-RADS categories (A, B, C, D) and generates comprehensive medical reports.

Overview

This project implements a production-ready breast density classification system that:

• Uses transfer learning with Inception v3 architecture
• Supports multiple mammography acquisition protocols (FFDM, C-View, Tomosynthesis)
• Generates AI-powered medical reports via LLM integration
• Deploys as a TorchServe model for scalable inference
• Integrates with cloud storage (GCS) for data management

Architecture

graph TB
    subgraph "Input Layer"
        A[Mammogram Image] --> B[Input Validation]
        B --> C[Image Preprocessing]
    end
    
    subgraph "Inference Pipeline"
        C --> D[MONAI Dataset Creation]
        D --> E[Data Augmentation]
        E --> F[Inception v3 Model]
        F --> G[Softmax Classification]
    end
    
    subgraph "Report Generation"
        G --> H[Output Processing]
        H --> I[LLM Client]
        I --> J[Med42 Report Generation]
        J --> K[PDF Conversion]
    end
    
    subgraph "Output Layer"
        K --> L[GCS Upload]
        L --> M[Response JSON]
    end
    
    subgraph "Configuration"
        N[Model Config] --> F
        O[TorchServe Config] --> P[Handler]
        P --> F
    end

Loading

Model Architecture

graph LR
    subgraph "Transfer Learning Pipeline"
        A[ImageNet Pre-trained<br/>Inception v3] --> B[Feature Extraction<br/>Layers]
        B --> C[Global Average<br/>Pooling]
        C --> D[Fully Connected<br/>Layer 1024]
        D --> E[Dropout<br/>Layer]
        E --> F[Output Layer<br/>4 classes]
        F --> G[Softmax<br/>Activation]
    end
    
    subgraph "Input Processing"
        H[Raw Mammogram<br/>299x299x3] --> I[Normalization<br/>0.0-1.0]
        I --> J[Channel First<br/>Format]
        J --> A
    end
    
    subgraph "Output Classes"
        G --> K[Class A: Almost Entirely Fat]
        G --> L[Class B: Scattered Fibroglandular]
        G --> M[Class C: Heterogeneously Dense]
        G --> N[Class D: Extremely Dense]
    end

Loading

Data Flow

sequenceDiagram
    participant Client
    participant TorchServe
    participant Handler
    participant Model
    participant LLM
    participant GCS
    
    Client->>TorchServe: POST /predictions/monai-breast-density-classification
    TorchServe->>Handler: preprocess(data)
    Handler->>Handler: validate_input()
    Handler->>Handler: create_dataset()
    Handler->>Model: inference()
    Model->>Handler: classification_output
    Handler->>LLM: generate_report()
    LLM->>Handler: medical_report.pdf
    Handler->>GCS: upload_files()
    GCS->>Handler: gs://bucket/path
    Handler->>TorchServe: response_json
    TorchServe->>Client: inference_result

Loading

Project Structure

fetch/
├── config/                          # Configuration files
│   ├── model-definition.json        # Model metadata and workflow
│   ├── model-config.yaml           # TorchServe model configuration
│   ├── config.properties           # TorchServe server settings
│   ├── archiver.properties         # MAR packaging configuration
│   └── long_description.md         # Model documentation
├── inference_monai/                 # Core inference package
│   ├── inference_handler.py        # Main TorchServe handler
│   ├── config.py                   # Environment configuration
│   ├── inference_types.py          # Input/Output type definitions
│   ├── inference_utils.py          # Utility functions
│   └── report_generator/           # Medical report generation
│       ├── report_generator.py     # Main report generation logic
│       ├── prompts.py              # LLM prompt templates
│       └── llm_interface/          # LLM client implementations
│           ├── llm_client.py       # Main LLM client
│           └── clients/            # Provider-specific clients
│               ├── openai.py       # OpenAI client
│               └── perplexity.py   # Perplexity client
├── weights/                        # Model weights
│   ├── get_weights.sh             # Weight download script
│   └── model.pt                   # Trained model weights
├── examples/                       # Sample data and configurations
├── resources/                      # Additional resources
│   └── paper.md                   # Research paper reference
├── utils/                         # Utility modules
└── setup.py                       # Package installation

Installation

Prerequisites

• Python 3.8+
• CUDA-compatible GPU (recommended)
• TorchServe
• Required API keys (OpenAI/Perplexity for report generation)

Setup

1. Clone and install dependencies:

git clone <repository-url>
cd fetch
pip install -e .

2. Download model weights:

cd weights
chmod +x get_weights.sh
./get_weights.sh

3. Set up environment variables:

export OPENAI_API_KEY="your-openai-key"
# or
export PERPLEXITY_API_KEY="your-perplexity-key"

Usage

TorchServe Deployment

1. Create MAR file:

torch-model-archiver \
    --model-name monai-breast-density-classification \
    --version 1.0 \
    --handler inference_monai/inference_handler.py \
    --config-file config/model-config.yaml \
    --serialized-file weights/model.pt \
    --export-path model-store \
    --extra-files config/,inference_monai/,resources/

2. Start TorchServe:

torchserve --start \
    --model-store model-store \
    --models monai-breast-density-classification.mar \
    --ts-config config/config.properties

3. Make inference requests:

curl -X POST http://localhost:8085/predictions/monai-breast-density-classification \
    -H "Content-Type: application/json" \
    -d '{
        "id": "test-001",
        "bucket": "your-bucket",
        "root_path": "test-session",
        "input_folder": "input",
        "input_files": ["mammogram.jpg"],
        "output_folder": "output"
    }'

Input Format

The model expects JSON input with the following structure:

{
    "id": "unique-session-id",
    "bucket": "gcs-bucket-name",
    "root_path": "session-directory",
    "input_folder": "input-subfolder",
    "input_files": ["mammogram.jpg"],
    "output_folder": "output-subfolder"
}

Output Format

The model returns a JSON response containing:

{
    "id": "unique-session-id",
    "model_name": "BreastDensityClassification",
    "outputs": [
        {
            "name": "output_path",
            "shape": [1],
            "datatype": "BYTES",
            "data": ["gs://bucket/path/to/report.pdf"]
        }
    ]
}

Model Details

Training Data

• Source: Mayo Clinic PACS (June-September 2021)
• Total Images: 151,164 mammograms
• Patients: 26,411 patients
• Acquisition Types: FFDM, C-View, Tomosynthesis Projection, Intelligent 2D

Model Performance

• Architecture: Inception v3 with transfer learning
• Input Size: 299×299×3 (RGB)
• Classes: 4 (BI-RADS categories A, B, C, D)
• Training: Multi-protocol dataset for improved generalization

BI-RADS Classification

• A: Almost entirely fat
• B: Scattered fibroglandular density
• C: Heterogeneously dense
• D: Extremely dense

Report Generation

The system generates comprehensive medical reports using:

1. Model Output Processing: Converts raw classification scores to interpretable results
2. LLM Integration: Uses OpenAI or Perplexity APIs for report generation
3. Medical Context: Incorporates research literature and clinical guidelines
4. PDF Generation: Converts markdown reports to professional PDF format

Report Structure

• Model Details and Technical Information
• Key Findings and Observations
• Clinical Significance and Implications
• Recommendations for Follow-up
• Limitations and Considerations
• References to Medical Literature

Configuration

TorchServe Configuration

• Workers: 1 (configurable in model-config.yaml)
• Batch Size: 16
• Timeout: 60 seconds
• Memory: 32GB recommended
• GPU: NVIDIA L4 or equivalent

Environment Variables

LOCAL_EXECUTION=false          # Set to true for local testing
OPENAI_API_KEY=your-key       # For OpenAI-based report generation
PERPLEXITY_API_KEY=your-key   # For Perplexity-based report generation

Development

Running Tests

# Install development dependencies
pip install -e ".[dev]"

# Run tests (if available)
python -m pytest tests/

Local Development

# Set local execution mode
export LOCAL_EXECUTION=true

# Test inference handler directly
python -c "
from inference_monai.inference_handler import BreastDensityClassificationHandler
handler = BreastDensityClassificationHandler()
# Add your test code here
"

Monitoring and Metrics

TorchServe provides built-in metrics at:

• Inference API: http://localhost:8085
• Management API: http://localhost:8086
• Metrics API: http://localhost:8082 (Prometheus format)

Limitations

• Single Image Processing: Currently processes one mammogram at a time
• Image Format: Requires JPEG format mammograms
• Size Requirements: Images are resized to 299×299 pixels
• Internet Connection: Report generation requires API access to LLM services
• Medical Disclaimer: Outputs are for research purposes and require clinical validation

Contributing

1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests if applicable
5. Submit a pull request

License

This project is licensed under the MIT License - see the LICENSE file for details.

References

1. Gupta, V., et al. "A Multi Reconstruction Study of Breast Density Estimation Using Deep Learning." arXiv:2202.08238 (2022)
2. MONAI Model Zoo: https://monai.io/model-zoo.html
3. BI-RADS Atlas: https://www.acr.org/Clinical-Resources/Reporting-and-Data-Systems/Bi-Rads

Citation

If you use this model in your research, please cite:

@article{gupta2022multi,
  title={A Multi Reconstruction Study of Breast Density Estimation Using Deep Learning},
  author={Gupta, Vikash and Demirer, Mutlu and Maxwell, Robert W and White, Richard D and Erdal, Barbaros S},
  journal={arXiv preprint arXiv:2202.08238},
  year={2022}
}