preprocessing.py

"""
preprocessing.py
Multi-gate fundus validation + preprocessing pipeline.

Gate 1 — Fundus Shape Gate:
  Retinal fundus photos have a distinct circular/oval bright region
  on a very dark (near-black) background. Non-fundus images fail this.

Gate 2 — Color Profile Gate:
  Fundus images are dominated by warm reddish-orange tones (blood vessels,
  retinal tissue). Screenshots, selfies, and random photos fail this.

Gate 3 — Quality Gate:
  Sharpness measured via Laplacian variance *inside* the fundus ROI.
  Low sharpness = blurry / out-of-focus scan → reject.

Inference transform (no augmentation, exact test_transform from notebook):
  Resize 380×380 → ImageNet Normalize → Tensor
"""

import cv2
import numpy as np
import base64
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2

IMG_SIZE = 380   # EfficientNet-B4 input size
CDR_SIZE = 512   # Notebook Step 4: CDR computed on 512×512
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

# ─── Exact test_transform from notebook ──────────────────────────────────────
test_transform = A.Compose([
    A.Resize(IMG_SIZE, IMG_SIZE),
    A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    ToTensorV2()
])

# ─── Display transform: CLAHE applied, no normalization ──────────────────────
display_transform = A.Compose([
    A.Resize(IMG_SIZE, IMG_SIZE),
    A.CLAHE(clip_limit=2.0, p=1.0),
])


# ═══════════════════════════════════════════════════════════════════════════
# GATE 1 — Fundus Shape Gate
# ═══════════════════════════════════════════════════════════════════════════

def check_fundus_shape(img_rgb: np.ndarray) -> tuple:
    """
    Gate 1: Verify the image contains a circular/oval bright region
    on a dark background — the hallmark of a retinal fundus photo.

    Checks:
      a) Dark border ratio: ≥ 20% of pixels must be near-black (L < 30)
         (fundus images have large dark surround from the camera aperture)
      b) Bright circular region: largest bright contour must be
         roughly circular (aspect ratio 0.5–2.0) and cover 10–85% of image.

    Returns: (passed: bool, reason: str, score: float 0–1)
    """
    h, w = img_rgb.shape[:2]
    total_pixels = h * w

    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)
    L = img_lab[:, :, 0]

    # ── a) Dark border check ─────────────────────────────────────────────
    dark_pixels = np.sum(L < 30)
    dark_ratio = dark_pixels / total_pixels

    if dark_ratio < 0.20:
        # Fundus images must have at least ~20% pitch black borders from camera aperture
        return False, f"Tidak terdeteksi sebagai foto fundus retina (dark border terlalu kecil: {dark_ratio:.0%})", 0.0

    # ── b) Bright circular region check ──────────────────────────────────
    _, bright_mask = cv2.threshold(L, 30, 255, cv2.THRESH_BINARY)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
    bright_mask = cv2.morphologyEx(bright_mask, cv2.MORPH_CLOSE, kernel)
    bright_mask = cv2.morphologyEx(bright_mask, cv2.MORPH_OPEN, kernel)

    contours, _ = cv2.findContours(bright_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return False, "Tidak ditemukan region cerah (bukan foto fundus retina)", 0.0

    largest = max(contours, key=cv2.contourArea)
    x, y, bw, bh = cv2.boundingRect(largest)
    bright_area = cv2.contourArea(largest)
    bright_ratio = bright_area / total_pixels

    if bright_ratio < 0.15:
        return False, f"Region retina terlalu kecil (area: {bright_ratio:.0%} dari gambar)", 0.0
    if bright_ratio > 0.80:
        return False, f"Seluruh gambar terang — bukan foto fundus dengan bingkai hitam ({bright_ratio:.0%})", 0.0

    # Aspect ratio check: fundus ROI must be roughly a perfect circle
    aspect = bw / bh if bh > 0 else 0
    if aspect < 0.7 or aspect > 1.3:
        return False, f"Bentuk region tidak bulat sempurna (aspect ratio: {aspect:.2f}) — bukan fundus retina", 0.0

    # ── Score: combination of dark border + circularity ──────────────────
    dark_score = min(dark_ratio / 0.40, 1.0)
    circ_score = 1.0 - abs(aspect - 1.0) / 0.3  # closer to 1.0 = more circular
    shape_score = round(max((dark_score + circ_score) / 2, 0.0), 3)

    return True, "OK", shape_score


# ═══════════════════════════════════════════════════════════════════════════
# GATE 2 — Color Profile Gate
# ═══════════════════════════════════════════════════════════════════════════

def check_fundus_color(img_rgb: np.ndarray) -> tuple:
    """
    Gate 2: Verify the image has a warm reddish-orange color profile
    typical of retinal fundus photography.

    Fundus images: dominant hue in HSV is orange-red (Hue 0–30 or 160–180),
    medium-high saturation. Grayscale/blue-dominant images fail.

    Returns: (passed: bool, reason: str, score: float 0–1)
    """
    img_hsv = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2HSV)
    H = img_hsv[:, :, 0]
    S = img_hsv[:, :, 1]
    V = img_hsv[:, :, 2]

    # Only look at pixels that aren't near-black (part of the fundus ROI)
    roi_mask = V > 40
    if roi_mask.sum() < 1000:
        return False, "Gambar terlalu gelap untuk dianalisis warnanya", 0.0

    H_roi = H[roi_mask]
    S_roi = S[roi_mask]

    # ── Reddish-orange hue check ─────────────────────────────────────────
    # In OpenCV HSV: Hue 0–17 = red-orange, 17–35 = orange-yellow
    # Red wraps around: also 160–180
    red_orange_mask = ((H_roi <= 20) | (H_roi >= 160))
    orange_yellow_mask = ((H_roi > 10) & (H_roi <= 35))
    warm_ratio = (red_orange_mask.sum() + orange_yellow_mask.sum()) / len(H_roi)

    # ── Saturation check ─────────────────────────────────────────────────
    mean_saturation = S_roi.mean()

    # Grayscale images or weakly colored diagrams: saturation is low
    if mean_saturation < 35:
        return False, f"Gambar tidak memiliki ketajaman warna fundus (saturasi: {mean_saturation:.0f} < 35)", 0.0

    if warm_ratio < 0.35:
        return False, (
            f"Warna tidak dominan merah/oranye fundus "
            f"(warm hue: {warm_ratio:.0%}, harusnya ≥35%). "
            f"Bukan foto medis retina yang valid."
        ), 0.0

    color_score = round(min(warm_ratio / 0.60, 1.0) * min(mean_saturation / 90, 1.0), 3)
    return True, "OK", color_score


# ═══════════════════════════════════════════════════════════════════════════
# GATE 3 — Quality Gate (Sharpness within fundus ROI)
# ═══════════════════════════════════════════════════════════════════════════

def compute_quality_score(img_rgb: np.ndarray) -> float:
    """
    Gate 3: Sharpness-based Quality Score measured INSIDE the fundus ROI.
    Uses Laplacian variance on the bright region only.

    Maps to 1–5 scale matching dataset QS convention.
    QS < 3 → reject.

    Thresholds calibrated for retinal fundus images (naturally smooth).
    """
    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)
    L = img_lab[:, :, 0]
    _, roi_mask = cv2.threshold(L, 30, 255, cv2.THRESH_BINARY)

    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    gray_resized = cv2.resize(gray, (300, 300))
    roi_resized  = cv2.resize(roi_mask, (300, 300))

    # Measure sharpness only inside the fundus ROI
    lap = cv2.Laplacian(gray_resized, cv2.CV_64F)
    if roi_resized.sum() > 5000:
        lap_var = lap[roi_resized > 0].var()
    else:
        lap_var = lap.var()  # fallback to full image

    print(f"[Preprocessing] Laplacian Variance (ROI): {lap_var:.2f}")

    if lap_var < 3.0:
        return 1.0   # Extremely blurry
    elif lap_var < 8.0:
        return 2.0   # Blurry → reject
    elif lap_var < 15.0:
        return 3.0   # Acceptable
    elif lap_var < 30.0:
        return 4.0   # Good
    else:
        return 5.0   # Excellent


# ═══════════════════════════════════════════════════════════════════════════
# Auto-crop Fundus ROI
# ═══════════════════════════════════════════════════════════════════════════

def auto_crop_fundus(img_rgb: np.ndarray) -> np.ndarray:
    """
    Auto-detect and crop the circular fundus region.
    Falls back to full image if detection fails.
    """
    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)
    L = img_lab[:, :, 0]
    _, mask = cv2.threshold(L, 30, 255, cv2.THRESH_BINARY)

    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (20, 20))
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return img_rgb

    largest = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(largest)
    pad = 10
    x1 = max(0, x - pad);  y1 = max(0, y - pad)
    x2 = min(img_rgb.shape[1], x + w + pad)
    y2 = min(img_rgb.shape[0], y + h + pad)
    cropped = img_rgb[y1:y2, x1:x2]
    return cropped if cropped.size > 0 else img_rgb


# ═══════════════════════════════════════════════════════════════════════════
# Main Preprocessing Pipeline
# ═══════════════════════════════════════════════════════════════════════════

def preprocess_image(img_bytes: bytes) -> dict:
    """
    Full multi-gate preprocessing pipeline:

    Gate 1 — Fundus Shape: circular bright region on dark background
    Gate 2 — Color Profile: reddish-orange dominant (retinal tissue)
    Gate 3 — Quality:       sharpness ≥ 3.0 inside fundus ROI

    Returns dict with:
        passed_gate      : bool
        gate_failed      : str | None  ('shape' | 'color' | 'quality' | None)
        rejection_reason : str | None
        quality_score    : float (1–5)
        shape_score      : float (0–1)
        color_score      : float (0–1)
        original_b64, preprocessed_b64, tensor, cdr_img_rgb
    """
    # ── Decode ───────────────────────────────────────────────────────────
    nparr = np.frombuffer(img_bytes, np.uint8)
    img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img_bgr is None:
        raise ValueError("Tidak dapat membaca gambar. Upload file JPEG/PNG yang valid.")
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    # ── Gate 1: Fundus Shape ─────────────────────────────────────────────
    shape_ok, shape_reason, shape_score = check_fundus_shape(img_rgb)
    print(f"[Gate 1 - Shape]  passed={shape_ok}, score={shape_score}, reason={shape_reason}")
    if not shape_ok:
        return _rejection_response(img_rgb, 'shape', shape_reason, shape_score, 0.0, 0.0)

    # ── Gate 2: Color Profile ────────────────────────────────────────────
    color_ok, color_reason, color_score = check_fundus_color(img_rgb)
    print(f"[Gate 2 - Color]  passed={color_ok}, score={color_score}, reason={color_reason}")
    if not color_ok:
        return _rejection_response(img_rgb, 'color', color_reason, shape_score, color_score, 0.0)

    # ── Gate 3: Quality ──────────────────────────────────────────────────
    quality_score = compute_quality_score(img_rgb)
    quality_ok = quality_score >= 3.0
    print(f"[Gate 3 - Quality] passed={quality_ok}, score={quality_score}/5.0")
    if not quality_ok:
        reason = (
            f"Kualitas gambar terlalu rendah (skor {quality_score}/5.0). "
            f"Gambar terlalu blur atau tidak fokus. Gunakan hasil scan fundus yang jelas."
        )
        return _rejection_response(img_rgb, 'quality', reason, shape_score, color_score, quality_score)

    # ── All gates passed → process ───────────────────────────────────────
    cropped = auto_crop_fundus(img_rgb)

    original_display = cv2.resize(cropped, (IMG_SIZE, IMG_SIZE))
    original_b64 = ndarray_to_b64(original_display)

    display_result = display_transform(image=cropped)
    preprocessed_b64 = ndarray_to_b64(display_result['image'])

    test_result = test_transform(image=cropped)
    tensor = test_result['image']

    cdr_img_rgb = cv2.resize(cropped, (CDR_SIZE, CDR_SIZE))

    return {
        'passed_gate':      True,
        'gate_failed':      None,
        'rejection_reason': None,
        'quality_score':    quality_score,
        'shape_score':      shape_score,
        'color_score':      color_score,
        'original_b64':     original_b64,
        'preprocessed_b64': preprocessed_b64,
        'tensor':           tensor,
        'cdr_img_rgb':      cdr_img_rgb,
    }


def _rejection_response(img_rgb: np.ndarray, gate: str, reason: str,
                         shape_score: float, color_score: float,
                         quality_score: float) -> dict:
    """Return a standardised rejection dict with a display thumbnail."""
    # Always give user a preview of what was uploaded
    thumb = cv2.resize(img_rgb, (IMG_SIZE, IMG_SIZE))
    original_b64 = ndarray_to_b64(thumb)
    return {
        'passed_gate':      False,
        'gate_failed':      gate,
        'rejection_reason': reason,
        'quality_score':    quality_score,
        'shape_score':      shape_score,
        'color_score':      color_score,
        'original_b64':     original_b64,
        'preprocessed_b64': original_b64,  # show same thumb
        'tensor':           None,
        'cdr_img_rgb':      None,
    }


def ndarray_to_b64(img_rgb: np.ndarray) -> str:
    """Convert RGB numpy array (uint8) to base64 JPEG string."""
    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
    _, buffer = cv2.imencode('.jpg', img_bgr, [cv2.IMWRITE_JPEG_QUALITY, 90])
    return base64.b64encode(buffer).decode('utf-8')