Training of Knoweledge Distillation Anomaly Detection Models

This repository contains the core pipeline for training and evaluating anomaly detection models (RD4AD and STFPM) using the Anomalib framework. It handles the transition from medical imaging formats (NIfTI) to deep learning inputs (PNG) and provides tools for detailed inference.

---

Repository Structure

Data Preparation

• data_conversion.py: Utilities to prepare datasets for training. It automatically converts NIfTI volumes into normalized 2D PNG slices while maintaining the required folder structure for Anomalib.

Model Training

• train_rd4ad.py / train_stfpm.py: Scripts to train the respective models with custom learning rate configurations.

Testing & Evaluation

• test_rd4ad.py / test_stfpm.py: Evaluation scripts for trained checkpoints.
• Note: These scripts are used for internal validation and checkpoint verification. They are not used for the final report metrics, as those are handled by the specialized post-processing pipeline (map extraction -> main_pipeline).

Inference & Map Extraction

• extract_rd4ad_maps.py / extract_stfpm_maps.py: Specialized inference scripts to extract raw data for downstream post-processing.
• Saves raw Anomaly Maps as .npy files and Prediction Masks as PNG or NIfTI.

---

🚀 Execution Guide

1. Training

To train the model (e.g., RD4AD), run the following command. The script will automatically prepare the PNG cache if NIfTI files are detected:

python train_rd4ad.py \
    --data_root /path/to/dataset/root \
    --name rd4ad_experiment_v1 \
    --batch_size 16 \
    --epochs 100 \
    --learning_rate 0.001

2. Map Extraction (Inference)

Once training is complete, extract the raw maps and masks required for the final evaluation pipeline:

python extract_rd4ad_maps.py \
  --data_root /path/to/dataset/png_cache \
  --checkpoint /path/to/results/checkpoints/last.ckpt \
  --output_dir ./extracted_output_maps \
  --mask_output_format png \
  --split test

3. Post-Processing

The final evaluation and metrics for the report are performed in a separate repository. The outputs generated by the Extraction scripts above serve as the input for that pipeline.

To run the final refinement and metrics pipeline, use the main_pipeline.py from the Post-Processing Repository:

python main_pipeline.py \
  --input-dir ./extracted_output_maps/prediction_masks/test \
  --anomaly-map-dir ./extracted_output_maps/anomaly_maps/test \
  --body-mask-dir /path/to/body_masks/ \
  --output-root ./post_process_results \
  --ground-truth-dir /path/to/dataset/test \
  --metrics-mean-fraction-thresholds 0.0 0.0001 0.0005 0.0008 0.001 0.0015 0.002 0.0025 0.003 0.0035 1.0