# Quick Mitigation Options for Class B Bias

This document outlines immediate workarounds for the Class B bias while you retrain the model.

## Option 1: Logit-Bias Correction (Inference-Time Hack)

Add this function to `main.py` in the `/analyze` endpoint before softmax:

```python
def apply_logit_bias_correction(logits, bias_correction=None):
    """
    Apply bias correction to logits before softmax.
    This temporarily mitigates the Class B overprediction.
    
    Args:
        logits: Raw model output (before softmax)
        bias_correction: Dict with per-class bias adjustments
    
    Returns:
        Adjusted logits
    """
    if bias_correction is None:
        # DEFAULT: Reduce Class B, boost others
        bias_correction = {
            0: +0.3,   # Adenocarcinoma: slight boost
            1: -1.2,   # Small Cell: heavy penalty
            2: +0.3,   # Large Cell: slight boost
            3: +0.3,   # Squamous Cell: slight boost
        }
    
    for class_idx, bias in bias_correction.items():
        logits[class_idx] += bias
    
    return logits
```

**Usage in classification code:**
```python
# in the /analyze endpoint, after model inference:
with torch.no_grad():
    logits = model(input_tensor)  # Raw output
    logits_corrected = apply_logit_bias_correction(logits[0])  # Apply correction
    probs = torch.softmax(logits_corrected, dim=0)  # Then softmax
```

## Option 2: Confidence-Based Filtering

Add confidence thresholds per-class:

```python
CONFIDENCE_THRESHOLDS = {
    0: 0.30,  # Adenocarcinoma: accept if >= 30%
    1: 0.60,  # Small Cell: require higher confidence (70%+)
    2: 0.25,  # Large Cell: accept lower confidence
    3: 0.25,  # Squamous Cell: accept lower confidence
}

if confidence < CONFIDENCE_THRESHOLDS[class_idx]:
    logger.warning(f"LOW CONFIDENCE: {class_idx} at {confidence*100:.1f}% (threshold: {CONFIDENCE_THRESHOLDS[class_idx]*100}%)")
    # Optionally flag as uncertain
```

## Option 3: Report Confidence Scores

Return all confidence scores to the frontend and let the UI handle the ambiguity:

**API Response Structure:**
```json
{
  "success": true,
  "detections": [
    {
      "tumor_id": 1,
      "prediction": "Small Cell (Class B)",
      "confidence": 75.97,
      "all_confidences": {
        "Adenocarcinoma": 3.5,
        "Small Cell": 75.97,
        "Large Cell": 6.89,
        "Squamous Cell": 13.64
      },
      "confidence_status": "UNCERTAIN",  // NEW: Add this
      "note": "Low confidence in this prediction. Multiple classes viable."
    }
  ]
}
```

Then in Flutter UI:
```dart
// Show warning if confidence is borderline
if (allConfidences.values.any((c) => c > 20 && c != maxConfidence)) {
  showConfidenceWarning("Multiple classes possible. Review with specialist.");
}
```

## Option 4: Ensemble Multiple Models

If you have other trained models:

```python
def ensemble_predict(input_tensor, models_list):
    """Average predictions from multiple models"""
    all_probs = []
    
    for model in models_list:
        with torch.no_grad():
            probs = torch.softmax(model(input_tensor), dim=1)[0]
        all_probs.append(probs)
    
    ensemble_probs = torch.stack(all_probs).mean(dim=0)
    return ensemble_probs
```

## Comparison of Options

| Option | Ease | Effectiveness | Drawbacks |
|--------|------|---------------|-----------|
| **Logit Bias Correction** | ⭐⭐⭐ Easy | ⭐⭐ Moderate | Doesn't fix root issue; magic numbers |
| **Confidence Filtering** | ⭐⭐⭐ Easy | ⭐⭐ Moderate | Still wrong predictions, just flagged |
| **Report All Scores** | ⭐⭐⭐ Easy | ⭐⭐ Moderate | Defers to human; adds complexity |
| **Ensemble Models** | ⭐⭐ Medium | ⭐⭐⭐ Good | Need multiple models |
| **Retrain with Weights** | ⭐ Hard | ⭐⭐⭐⭐⭐ Excellent | Takes time; need training data |

## ⚠️ Recommendation

**Use Option 3 (Report All Scores) + Option 1 (Logit Correction) as SHORT-TERM fix**
- Apply light logit correction to reduce Class B dominance
- Return all confidence scores to frontend
- Show warnings when confidence is low (< 60%)
- **Plan to retrain model properly** (highest priority)

**Then retrain the model** (Option 5 - see `retrain_with_balanced_weights.py`)

## Testing Your Fix

After implementing mitigation:

```bash
# Test with images from each class
LABELS=("Class_A" "Class_B" "Class_E" "Class_E")

for label in "${LABELS[@]}"; do
    img=$(ls test-images/${label}_*.png | head -1)
    echo "Testing: $img"
    curl -X POST -F "file=@$img" http://localhost:5001/analyze | grep -A 10 '"prediction"'
done
```

Expected results after fix:
```
Class A image → Mostly Class A (not Class B)
Class B image → Mostly Class B (correct)
Class E image → Mix of E and others (not all B)
Class G image → Mix of G and others (not all B)
```