Lung Nodule Analysis using 3D CNNs

This project provides a complete pipeline for detecting and classifying lung nodules from 3D CT scans using a 3D Convolutional Neural Network (CNN) built with PyTorch. The pipeline covers data downloading, preprocessing, model training, and evaluation, using the public LUNA16 dataset.

Table of Contents

• Features
• Project Workflow
• Project Structure
• Getting Started
  • Prerequisites
  • Installation & Setup
• How to Use
  • Step 1: Configure the Project
  • Step 2: Download Data
  • Step 3: Preprocess Data
  • Step 4: Train the Model
  • detailed: Your PyTorch Project with Docker
• Configuration Details
• Model Architecture
• Future Improvements
• Acknowledgments
• License

Features

• End-to-End Pipeline: From raw LUNA16 data to a trained nodule classification model.
• 3D CNN Architecture: Utilizes a Simple3DCNN specifically designed to process 3D medical imaging patches.
• Efficient Preprocessing: A run-once script processes the entire dataset, resampling scans to a uniform spacing, normalizing Hounsfield Units (HU), and extracting 3D patches around nodule candidates. This significantly speeds up training.
• Class Imbalance Handling: Implements a simple but effective under-sampling strategy in the Dataset class to create a balanced training set.
• Configuration Driven: All key parameters (paths, patch size, learning rate, etc.) are centralized in config.py for easy modification and experimentation.
• Reproducibility: Uses Pipfile and Pipfile.lock to ensure a consistent and reproducible development environment.

Project Workflow

The project follows a clear, sequential workflow:

1. Download (download_data.py): The raw LUNA16 .mhd/.raw CT scan files are downloaded from the Zenodo archive. You can configure which data subsets to download.
2. Preprocess (preprocess.py): This crucial script iterates through the raw CT scans and the candidates.csv file. For each candidate, it:
   • Resamples the CT scan to a uniform 1x1x1 mm spacing.
   • Normalizes voxel values (Hounsfield Units) to a [0, 1] range.
   • Extracts a 3D patch (e.g., 48x48x48) centered on the candidate's coordinates.
   • Saves each patch as a separate .npy file, sorted into nodules and non_nodules folders.
3. Train (train.py):
   • The custom LunaDataset loads the preprocessed .npy patches, handling the train/validation split and class balancing.
   • The Simple3DCNN model is trained to classify the patches.
   • The script evaluates performance using metrics like Accuracy, Precision, Recall, and F1-Score.
   • The best-performing model (based on validation F1-score) is saved to outputs/models/.

Project Structure

Lung-Nodule-Analysis/
├── data/                       # Root for all raw LUNA16 data (created automatically)
├── outputs/                    # Root for all generated files (created automatically)
│   ├── preprocessed_data/      # Stores the extracted .npy patches
│   └── models/                 # Stores the trained .pth model files
├── config.py                   # Central configuration file for all paths and hyperparameters.
├── download_data.py            # Script to download LUNA16 dataset subsets.
├── preprocess.py               # Script to process raw scans and extract 3D patches.
├── helper.py                   # Utility functions (e.g., image resampling, HU normalization).
├── model.py                    # Defines the 3D CNN model architecture.
├── dataset.py                  # PyTorch Dataset class to load preprocessed patches.
├── train.py                    # Main script to run the training and validation loop.
├── evaluate.py                 # Script to evaluate the final model on the real test set.
├── test_synthetic.py           # Script to test the model's logic on synthetic data.
├── Pipfile                     # Lists project dependencies for pipenv.
├── Pipfile.lock                # Locks dependency versions for reproducibility.
└── README.md                   # You are here!

Getting Started

Prerequisites

• Python 3.8 or higher
• pipenv to manage dependencies.
• Sufficient disk space for the LUNA16 dataset and preprocessed patches (~20-50 GB depending on the number of subsets downloaded).

Installation & Setup

1. Clone the repository:

   git clone <repo-url>
   cd Lung-Nodule-Analysis

2. Build the Image

   docker compose build

3. Run the container

   docker compose up -d

4. Access the container

   docker compose exec app bash

5. Run the download script

   python3 download_data.py

6. Run the train script

   python3 train.py

7. Run the evaluate script

   python3 evaluate.py

8. Run the evaluate script

   python3 test_synthetic.py

How to Use

Step 1: Configure the Project

Open config.py and review the parameters. The default paths are already set up, but you can adjust them if needed. Most importantly, check the LUNA16_RAW_DATA_PATH and PREPROCESSED_DATA_PATH.

Step 2: Download Data

The LUNA16 dataset is split into 10 subsets. You can choose which ones to download.

1. Open download_data.py.
2. Modify the SUBSETS_TO_DOWNLOAD list to include the desired subsets (e.g., ["subset0", "subset1"]).
3. Run the script:

   python3 download_data.py

   Warning: This will download several gigabytes of data and may take a significant amount of time depending on your internet connection. The script will display a progress bar.

Step 3: Preprocess Data

After downloading, you need to run the preprocessing script to extract the 3D patches for training.

python preprocess.py

Warning: This is a very computationally intensive step and can take several hours to complete. It will also consume a large amount of disk space (~20-50 GB) as it saves every patch as a .npy file.

Step 4: Train the Model

Once the data is preprocessed, you can start training the model.

python train.py

The script will:

• Print the device being used (CPU or CUDA).
• Load the training and validation datasets.
• Start the training loop, showing progress for each epoch.
• Print training loss/accuracy and validation loss/metrics after each epoch.
• Save the model with the best validation F1-score to outputs/models/best_model.pth.

Running Your PyTorch Project with Docker

This document outlines the steps to build and run this PyTorch application using Docker and Docker Compose. This ensures a consistent and reproducible environment for development and deployment.

Prerequisites

1. Docker: Ensure Docker Desktop (for Mac/Windows) or Docker Engine (for Linux) is installed and running. You can download it from docker.com.
2. Docker Compose: Docker Compose V2 is typically included with Docker Desktop. For Linux, you might need to install it separately.
3. (Optional) NVIDIA GPU Support:
   • If you intend to use NVIDIA GPUs, ensure you have the latest NVIDIA drivers installed on your host machine.
   • Install the NVIDIA Container Toolkit on your host machine. This allows Docker containers to access NVIDIA GPUs.
4. Project Files:
   • Dockerfile: Defines the Docker image for the application.
   • docker-compose.yml: Defines how to run the application services (including GPU support).
   • Pipfile: Specifies Python package dependencies.
   • Pipfile.lock: Locks package versions for reproducible builds.
   • Your application code (e.g., inference.py).

Building and Running the Application

We will use Docker Compose to manage the build and run process.

Step 1: Clone the Repository (if applicable)

If you haven't already, clone the project repository to your local machine:

git clone <your-repository-url>
cd <your-project-directory>

Step 2: Check/Generate Pipfile.lock

The Dockerfile uses pipenv install --deploy, which requires Pipfile.lock to be up-to-date with Pipfile.

Troubleshooting Pipfile.lock out-of-date error: If, during the Docker build process (Step 3), you encounter an error similar to:

Your Pipfile.lock (...) is out of date. Expected: (...).
ERROR:: Aborting deploy

This means your Pipfile.lock is not synchronized with your Pipfile. To fix this, run the following command in your project's root directory (where Pipfile is located) on your host machine:

pipenv lock

This will update Pipfile.lock. After running this command, proceed to Step 3.

Step 3: Build and Run with Docker Compose

Open your terminal in the root directory of the project (where docker-compose.yml and Dockerfile are located).

To build the image and run the application (e.g., execute inference.py):

docker-compose up --build

• --build: This flag tells Docker Compose to build the Docker image using the Dockerfile. You can omit this on subsequent runs if the Dockerfile or its dependencies haven't changed, and an image already exists.
• The application (defined by CMD in the Dockerfile, e.g., python3 inference.py) will start, and its output will be displayed in your terminal.

To run in detached mode (in the background):

docker-compose up --build -d

Step 4: Interacting with the Application

• Viewing Logs (if running in detached mode):

  docker-compose logs -f app

  (Replace app with your service name if it's different in docker-compose.yml). Press Ctrl+C to stop following logs.

• Accessing a Shell Inside the Container (for debugging): If you need to explore the container's environment or run commands manually:

  1. Ensure the container is running (e.g., using docker-compose up -d).
  2. Open a shell:

     docker-compose exec app bash

     (Replace app with your service name if it's different).
  3. Inside the container, you can navigate to /app (the working directory) and run Python scripts or other commands.

• Port Mapping (if applicable): If your application (inference.py) runs a web server (e.g., on port 8000) and you have configured port mapping in docker-compose.yml (e.g., ports: - "8000:8000"), you can access it via http://localhost:8000 in your web browser.

Step 5: Stopping the Application

To stop and remove the containers, networks, and (optionally, depending on docker-compose down flags) volumes defined by Docker Compose:

docker-compose down

If you want to remove the volumes as well:

docker-compose down -v

Important Notes

• PyTorch Versions & CUDA: The Pipfile specifies PyTorch versions and a CUDA source (pytorch-cu111). Ensure these versions are valid and available from the specified PyTorch wheel index. If pipenv install fails during the Docker build due to version conflicts or "Could not find a version" errors, you will need to:
  1. Consult PyTorch Previous Versions to find compatible torch, torchvision, and torchaudio versions for your desired CUDA version (e.g., CUDA 11.1).
  2. Update the versions in your Pipfile.
  3. Run pipenv lock locally to regenerate Pipfile.lock.
  4. Re-run docker-compose up --build.
• GPU Usage: The docker-compose.yml is configured to attempt GPU access using NVIDIA. This requires the prerequisites mentioned above (NVIDIA drivers and NVIDIA Container Toolkit on the host). If GPUs are not available or not configured correctly, PyTorch will typically fall back to CPU mode.
• Development Mode Volume Mount: The docker-compose.yml includes volumes: - .:/app. This mounts your local project directory into the container. Code changes made locally will be reflected inside the container, which is useful for development. For production, you might remove this volume mount to rely solely on the code baked into the image.

Further Actions

• Cleaning up Docker Resources:
  • To remove unused Docker images: docker image prune
  • To remove unused Docker volumes: docker volume prune
  • To remove unused Docker networks: docker network prune
  • To remove all unused Docker resources (images, containers, volumes, networks): docker system prune -a (Use with caution!)

Configuration Details

All major parameters can be tweaked in config.py without changing the core logic:

• LUNA16_RAW_DATA_PATH: Path to the raw LUNA16 dataset.
• PREPROCESSED_DATA_PATH: Where to save the .npy patches.
• MODEL_SAVE_PATH: Where to save trained models.
• DEVICE: Automatically selects CUDA if available, otherwise CPU.
• TARGET_SPACING: The isotropic spacing [x, y, z] to resample all scans to.
• PATCH_SIZE: The [D, H, W] dimensions of the 3D patches to extract.
• HU_MIN, HU_MAX: Hounsfield Unit range for clipping.
• BATCH_SIZE, LEARNING_RATE, NUM_EPOCHS: Standard training hyperparameters.

Model Architecture

The project uses a Simple3DCNN defined in model.py. The architecture consists of three sequential 3D convolutional blocks followed by a fully connected classification head.

Each convolutional block contains:

1. Conv3d: Learns 3D spatial features.
2. ReLU: Activation function.
3. BatchNorm3d: Stabilizes and accelerates training.
4. MaxPool3d: Downsamples the feature maps.

The final block includes a Flatten layer, Linear layers, and Dropout for regularization. The output is a single logit, which is passed to a BCEWithLogitsLoss function for binary classification.

Future Improvements

• Data Augmentation: Implement 3D augmentations (e.g., random flips, rotations, scaling) in dataset.py to improve model generalization.
• Advanced Architectures: Experiment with more complex models like 3D ResNets or DenseNets.
• Cross-Validation: Implement K-fold cross-validation for more robust evaluation.
• Advanced Imbalance Handling: Explore other techniques like over-sampling (e.g., SMOTE) or using weighted loss functions (e.g., Focal Loss).
• Inference Script: Create a separate script (predict.py) to run a trained model on a new, unseen CT scan.

Acknowledgments

This project uses the LUNA16 (LUng Nodule Analysis 2016) dataset. We extend our gratitude to the challenge organizers and data providers.

License

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