preprocessing.py

"""
preprocessing.py
================
Pipeline tiền xử lý ảnh đáy mắt (fundus) cho khâu suy luận (Inference Pipeline).
Hỗ trợ loại bỏ viền đen, resize giữ tỷ lệ (letterbox), lọc nhiễu Ben Graham và chuẩn hóa Tensor PyTorch.
"""

from __future__ import annotations
import os
import io
from typing import Tuple, Union
import cv2
import numpy as np
from PIL import Image
import torch
from torchvision import transforms

TARGET_SIZE = (224, 224)
CROP_TOLERANCE = 12
BEN_SIGMA = 10
BLACK_BORDER_RATIO_THRESH = 0.05
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def crop_fundus_circle(img: np.ndarray, tolerance: int = CROP_TOLERANCE) -> np.ndarray:
    """Crop bounding box của vùng sáng trên ảnh BGR (loại viền đen)."""
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    _, mask = cv2.threshold(gray, tolerance, 255, cv2.THRESH_BINARY)
    coords = cv2.findNonZero(mask)
    if coords is None:
        return img
    x, y, w, h = cv2.boundingRect(coords)
    return img[y: y + h, x: x + w]


def auto_detect_border(img: np.ndarray, thresh: float = BLACK_BORDER_RATIO_THRESH) -> bool:
    """Tự động phát hiện xem ảnh có viền đen xung quanh hay không dựa trên tỷ lệ pixel tối."""
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    return float((gray < CROP_TOLERANCE).mean()) > thresh


def letterbox_resize(
    img: np.ndarray,
    target_size: Tuple[int, int] = TARGET_SIZE,
    interpolation: int = cv2.INTER_CUBIC,
) -> np.ndarray:
    """Resize ảnh về kích thước target_size giữ nguyên aspect ratio (đệm viền đen)."""
    h, w = img.shape[:2]
    th, tw = target_size
    scale = min(tw / w, th / h)
    nw, nh = int(w * scale), int(h * scale)
    resized = cv2.resize(img, (nw, nh), interpolation=interpolation)
    canvas = np.zeros((th, tw, 3), dtype=np.uint8)
    pad_y = (th - nh) // 2
    pad_x = (tw - nw) // 2
    canvas[pad_y: pad_y + nh, pad_x: pad_x + nw] = resized
    return canvas


def ben_graham_transform(img: np.ndarray, sigma_x: int = BEN_SIGMA) -> np.ndarray:
    """Xử lý tăng cường tương phản Ben Graham: output = 4*img - 4*Blur + 128."""
    blur = cv2.GaussianBlur(img, (0, 0), sigmaX=sigma_x)
    enhanced = cv2.addWeighted(img, 4, blur, -4, 128)
    return np.clip(enhanced, 0, 255).astype(np.uint8)


def full_preprocess_pipeline(
    img: np.ndarray,
    target_size: Tuple[int, int] = TARGET_SIZE,
    use_ben_graham: bool = True,
    force_crop: bool | None = None,
) -> np.ndarray:
    """
    Pipeline xử lý ảnh OpenCV đầy đủ:
    1. Tự động kiểm tra & Crop viền đen
    2. Resize letterbox về target_size
    3. Ben Graham contrast enhancement (nếu use_ben_graham=True)
    Returns: numpy array BGR
    """
    do_crop = auto_detect_border(img) if force_crop is None else force_crop
    if do_crop:
        img = crop_fundus_circle(img)
    img = letterbox_resize(img, target_size)
    if use_ben_graham:
        img = ben_graham_transform(img)
    return img


def load_image(image_input: Union[str, bytes, Image.Image, np.ndarray]) -> np.ndarray:
    """
    Chuyển đổi các định dạng đầu vào (Path, Bytes, PIL Image, BGR Numpy Array) -> OpenCV BGR array.
    """
    if isinstance(image_input, (str, bytes, bytearray)):
        if isinstance(image_input, str):
            img = cv2.imread(image_input)
            if img is None:
                raise ValueError(f"Không thể đọc file ảnh từ đường dẫn: {image_input}")
            return img
        else:
            buf = np.frombuffer(image_input, dtype=np.uint8)
            img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
            if img is None:
                raise ValueError("Không thể giải mã dữ liệu bytes thành ảnh.")
            return img
    elif isinstance(image_input, Image.Image):
        # PIL (RGB) -> OpenCV (BGR)
        img_rgb = np.array(image_input.convert("RGB"))
        return cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
    elif isinstance(image_input, np.ndarray):
        if image_input.ndim == 2:  # Grayscale
            return cv2.cvtColor(image_input, cv2.COLOR_GRAY2BGR)
        elif image_input.shape[2] == 4:  # RGBA
            return cv2.cvtColor(image_input, cv2.COLOR_RGBA2BGR)
        return image_input.copy()
    else:
        raise TypeError(f"Kiểu dữ liệu đầu vào không được hỗ trợ: {type(image_input)}")


def prepare_image_tensor(
    image_input: Union[str, bytes, Image.Image, np.ndarray],
    target_size: Tuple[int, int] = TARGET_SIZE,
    mean: Tuple[float, float, float] = IMAGENET_MEAN,
    std: Tuple[float, float, float] = IMAGENET_STD,
    use_ben_graham: bool = True,
) -> torch.Tensor:
    """
    Nhận đầu vào linh hoạt -> Tiền xử lý OpenCV -> Chuyển thành PyTorch Tensor (1, C, H, W).
    """
    img_bgr = load_image(image_input)
    processed_bgr = full_preprocess_pipeline(
        img_bgr, target_size=target_size, use_ben_graham=use_ben_graham
    )
    # OpenCV BGR -> PIL RGB -> PyTorch Tensor
    processed_rgb = cv2.cvtColor(processed_bgr, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(processed_rgb)

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std),
    ])
    tensor = transform(pil_img)  # shape: (3, H, W)
    return tensor


class DRPredictor:
    def __init__(self, convnext_path: str | None = None, vit_path: str | None = None):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Candidate paths for ConvNeXt
        convnext_candidates = [
            convnext_path,
            "convnext_inference.pt",
            "convnext_results/convnext_inference.pt",
            os.path.join(os.path.dirname(__file__), "convnext_inference.pt"),
            os.path.join(os.path.dirname(__file__), "convnext_results", "convnext_inference.pt"),
        ]
        self.convnext_path = None
        for cand in convnext_candidates:
            if cand and os.path.exists(cand):
                self.convnext_path = cand
                break
        if self.convnext_path is None:
            self.convnext_path = "convnext_results/convnext_inference.pt"

        # Candidate paths for ViT
        vit_candidates = [
            vit_path,
            "vit_inference.pt",
            "vit_results/vit_inference.pt",
            os.path.join(os.path.dirname(__file__), "vit_inference.pt"),
            os.path.join(os.path.dirname(__file__), "vit_results", "vit_inference.pt"),
        ]
        self.vit_path = None
        for cand in vit_candidates:
            if cand and os.path.exists(cand):
                self.vit_path = cand
                break
        if self.vit_path is None:
            self.vit_path = "vit_results/vit_inference.pt"

        print(f"[DRPredictor] Loading ConvNeXt from {self.convnext_path} on {self.device}...")
        self.convnext = torch.jit.load(self.convnext_path, map_location=self.device)
        self.convnext.eval()
        
        print(f"[DRPredictor] Loading ViT from {self.vit_path} on {self.device}...")
        self.vit = torch.jit.load(self.vit_path, map_location=self.device)
        self.vit.eval()
        
        self.class_names = {
            0: "No DR",
            1: "Mild",
            2: "Moderate",
            3: "Severe",
            4: "Proliferative DR",
        }
        
        self.threshold_multipliers = np.array([1.7077, 0.6497, 1.1177, 0.9007, 0.6243], dtype=np.float32)

    def predict(self, image_input: Union[str, bytes, Image.Image, np.ndarray], use_ben_graham: bool = True) -> dict:
        # Prepare input tensor using existing prepare_image_tensor function
        tensor = prepare_image_tensor(image_input, use_ben_graham=use_ben_graham).unsqueeze(0).to(self.device)
        
        # Model forward passes
        with torch.no_grad():
            prob_cn = self.convnext(tensor)[0].cpu().numpy()
            prob_vit = self.vit(tensor)[0].cpu().numpy()
            
        # Soft Voting Ensemble
        prob_ensemble = 0.5 * prob_cn + 0.5 * prob_vit
        
        # Apply Optimized Threshold Multipliers
        scaled_probs = prob_ensemble * self.threshold_multipliers
        final_probs = scaled_probs / np.sum(scaled_probs)
        class_id = int(np.argmax(final_probs))
        
        return {
            "class_id": class_id,
            "class_name": self.class_names[class_id],
            "confidence": float(final_probs[class_id]),
            "probabilities": {
                "No DR": float(final_probs[0]),
                "Mild": float(final_probs[1]),
                "Moderate": float(final_probs[2]),
                "Severe": float(final_probs[3]),
                "Proliferative DR": float(final_probs[4]),
            },
        }