cdr_extraction.py

"""
cdr_extraction.py
Optic Disc & Cup segmentation + CDR computation.
Exact methodology from [IDSC]_D4.ipynb Cell 5.
"""

import cv2
import numpy as np
import base64
from typing import Optional, Tuple, Dict


# ─── 4.1 Optic Disc Segmentation ─────────────────────────────────────────────

def segment_optic_disc(img_rgb: np.ndarray):
    """
    Segment optic disc using brightness-based approach (LAB L-channel).
    Optic disc = brightest, large circular region in the retina.

    Returns: (disc_mask, bbox, centroid) or (None, None, None) on failure.
    """
    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)
    L_channel = img_lab[:, :, 0]

    # Otsu threshold on L channel (brightness)
    _, bright_mask = cv2.threshold(L_channel, 0, 255,
                                   cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # Morphological cleanup
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
    disc_mask = cv2.morphologyEx(bright_mask, cv2.MORPH_CLOSE, kernel)
    disc_mask = cv2.morphologyEx(disc_mask, cv2.MORPH_OPEN, kernel)

    # Largest connected component = optic disc
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
        disc_mask, connectivity=8
    )

    if num_labels < 2:
        return None, None, None

    largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
    disc_mask_final = (labels == largest_label).astype(np.uint8) * 255

    x = stats[largest_label, cv2.CC_STAT_LEFT]
    y = stats[largest_label, cv2.CC_STAT_TOP]
    w = stats[largest_label, cv2.CC_STAT_WIDTH]
    h = stats[largest_label, cv2.CC_STAT_HEIGHT]
    centroid = centroids[largest_label]

    return disc_mask_final, (x, y, w, h), centroid


# ─── 4.2 Optic Cup Segmentation ──────────────────────────────────────────────

def segment_optic_cup(img_rgb: np.ndarray, disc_bbox: Optional[Tuple]) -> Optional[np.ndarray]:
    """
    Segment optic cup within the optic disc region.
    Cup = brightest central area within the disc (75th percentile of L channel).

    Returns full-size cup mask or None.
    """
    if disc_bbox is None:
        return None

    x, y, w, h = disc_bbox
    margin = 10
    x1 = max(0, x - margin)
    y1 = max(0, y - margin)
    x2 = min(img_rgb.shape[1], x + w + margin)
    y2 = min(img_rgb.shape[0], y + h + margin)

    disc_region = img_rgb[y1:y2, x1:x2]
    if disc_region.size == 0:
        return None

    disc_lab = cv2.cvtColor(disc_region, cv2.COLOR_RGB2LAB)
    L_disc = disc_lab[:, :, 0]

    # 75th percentile threshold for cup
    threshold = np.percentile(L_disc, 75)
    cup_mask = (L_disc > threshold).astype(np.uint8) * 255

    # Morphological smoothing
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
    cup_mask = cv2.morphologyEx(cup_mask, cv2.MORPH_CLOSE, kernel)
    cup_mask = cv2.morphologyEx(cup_mask, cv2.MORPH_OPEN, kernel)

    # Largest connected component
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
        cup_mask, connectivity=8
    )
    if num_labels < 2:
        return None

    largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
    cup_mask_final = (labels == largest_label).astype(np.uint8) * 255

    # Return to full image size
    full_cup_mask = np.zeros(img_rgb.shape[:2], dtype=np.uint8)
    full_cup_mask[y1:y2, x1:x2] = cup_mask_final

    return full_cup_mask


# ─── 4.3 CDR Computation ─────────────────────────────────────────────────────

def compute_cdr(disc_mask: np.ndarray, cup_mask: np.ndarray) -> Optional[Dict]:
    """
    Compute Cup-to-Disc Ratio (CDR) metrics.
    CDR = diameter_cup / diameter_disc

    Returns dict with: vertical_cdr, horizontal_cdr, area_cdr, mean_cdr
    """
    if disc_mask is None or cup_mask is None:
        return None

    disc_contours, _ = cv2.findContours(disc_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cup_contours, _ = cv2.findContours(cup_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if not disc_contours or not cup_contours:
        return None

    # Bounding rects for diameter — use LARGEST contour (notebook Step 4.3)
    disc_rect = cv2.boundingRect(max(disc_contours, key=cv2.contourArea))
    cup_rect  = cv2.boundingRect(max(cup_contours,  key=cv2.contourArea))

    disc_w, disc_h = disc_rect[2], disc_rect[3]
    cup_w, cup_h = cup_rect[2], cup_rect[3]

    if disc_h == 0 or disc_w == 0:
        return None

    vertical_cdr = cup_h / disc_h
    horizontal_cdr = cup_w / disc_w

    # Area-based CDR
    disc_area = cv2.countNonZero(disc_mask)
    cup_area = cv2.countNonZero(cup_mask)
    area_cdr = cup_area / disc_area if disc_area > 0 else 0.0

    mean_cdr = (vertical_cdr + horizontal_cdr) / 2.0

    return {
        'vertical_cdr': round(float(vertical_cdr), 4),
        'horizontal_cdr': round(float(horizontal_cdr), 4),
        'area_cdr': round(float(area_cdr), 4),
        'mean_cdr': round(float(mean_cdr), 4),
    }


# ─── Contour Overlay for Display ─────────────────────────────────────────────

def generate_contour_overlay(img_rgb: np.ndarray,
                              disc_mask: Optional[np.ndarray],
                              cup_mask: Optional[np.ndarray]) -> str:
    """
    Overlay optic disc (green) and optic cup (yellow) contours on the image.
    Returns base64 JPEG string.
    """
    overlay = img_rgb.copy()

    if disc_mask is not None:
        disc_contours, _ = cv2.findContours(disc_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(overlay, disc_contours, -1, (0, 255, 80), 2)

    if cup_mask is not None:
        cup_contours, _ = cv2.findContours(cup_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(overlay, cup_contours, -1, (255, 230, 0), 2)

    overlay_bgr = cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR)
    _, buffer = cv2.imencode('.jpg', overlay_bgr, [cv2.IMWRITE_JPEG_QUALITY, 90])
    return base64.b64encode(buffer).decode('utf-8')


# ─── Full CDR Pipeline ────────────────────────────────────────────────────────

def run_cdr_pipeline(img_rgb: np.ndarray) -> Dict:
    """
    Full CDR extraction on an RGB image (already resized to 380x380).
    Returns CDR metrics + contour overlay base64.
    """
    disc_mask, disc_bbox, centroid = segment_optic_disc(img_rgb)
    cup_mask = segment_optic_cup(img_rgb, disc_bbox)
    cdr = compute_cdr(disc_mask, cup_mask)
    contour_b64 = generate_contour_overlay(img_rgb, disc_mask, cup_mask)

    if cdr is None:
        # Fallback values if segmentation fails
        cdr = {
            'vertical_cdr': 0.50,
            'horizontal_cdr': 0.50,
            'area_cdr': 0.25,
            'mean_cdr': 0.50,
        }

    return {
        'cdr': cdr,
        'contour_overlay_b64': contour_b64,
        'disc_detected': disc_mask is not None,
        'cup_detected': cup_mask is not None,
    }