Skip to content

Future Works & Improvements: Explainable AI (XAI)

As part of the continuous evolution of the Cervical Cytology AI pipeline, our next major milestone is integrating Explainable AI (XAI).

The primary goal is to enhance the trust and clinical utility of the AI system by visualizing why the MobileViT model classified a specific cell as High Risk (Dyskeratotic or Koilocytotic) using methodologies like SHAP (SHapley Additive exPlanations) or Grad-CAM.

🎯 Clinical Motivation

In medical diagnostics, "black-box" models often face resistance from clinical practitioners. By providing visual evidence of the neural network's attention map, cytologists can:

  • Instantly verify if the AI is looking at the correct cellular structures (e.g., enlarged nuclei, irregular borders).
  • Build confidence in the automated triage system.
  • Catch potential model biases or artifacts.

Performance Considerations

Generating SHAP values or pixel-level gradients is computationally expensive. To maintain the pipeline's high throughput, we propose generating explanations only for High Risk cells (Dyskeratotic/Koilocytotic) that exceed a specific confidence threshold.

🚀 Proposed Architectural Enhancements

To implement this without disrupting the current production pipeline, we plan to introduce a modular XAI engine.

1. New Explainability Module (src/explainable/)

We will create a new package to handle all interpretability logic, completely decoupled from the core inference engine.

  • src/explainable/engine.py:

    • XAIEngine: A wrapper around the MobileViT inference model.
    • Uses shap.GradientExplainer or shap.PartitionExplainer.
    • Output: Raw SHAP values and a PIL image overlay (heatmap).

  • src/explainable/vis.py:

    • Generates clinician-friendly visualizations.
    • Example: A side-by-side view showing [Original Crop] | [Attention Heatmap] | [Prototypical Example].

2. Pipeline Integration (src/pipeline.py)

The orchestration layer will be updated to seamlessly route high-risk crops to the XAI engine.

  1. Filter: After classification, isolate cells predicted as Dyskeratotic or Koilocytotic with a confidence score > 0.80.
  2. Explain: Pass these specific crops to the XAIEngine.
  3. Store: Save the resulting heatmap visualizations to a new directory (data/reports/figures/).
  4. Link: Enrich the cell's JSON metadata with the explanation_path for downstream reporting.

📄 Reporting Layer Updates

The generated PDF report (handled by src/report_gen.py) will feature a new dedicated section.

🔍 Section: AI Logic & Interpretability

  • Content: Display the top 3 most critical high-risk cells alongside their generated SHAP heatmaps.
  • Visual Legend: Include a caption explaining the heatmap gradient (e.g., "Warmer colors [Red/Pink] indicate regions the model identified as highly suspicious.").

🧪 Verification & Release Strategy

Before merging the XAI module into the main branch, the following validation steps will be required:

  1. Automated Unit Tests: Ensure XAIEngine generates correctly dimensioned heatmaps without memory leaks.
  2. Integration Tests: Run src/pipeline.py on a standard batch of SIPaKMeD test images to ensure the data/reports/figures/ directory populates correctly.
  3. Clinical Review: Manually inspect the generated PDF reports to verify that the heatmaps accurately align with the biological structures of the cells and that the PDF layout remains professional and unbroken.