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 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