""" 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]), }, }