src/brain_tumor_flair_classifier/classifier.py

'''提供 FLAIR NIfTI 脑肿瘤病例级分类器'''

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from statistics import mean
from typing import Any

import nibabel as nib
import numpy as np
import torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification


DEFAULT_THRESHOLD = 0.548381
DEFAULT_MAX_SLICES = 25
REQUIRED_MODEL_FILES = ('config.json', 'model.safetensors', 'preprocessor_config.json')


@dataclass(frozen=True)
class FlairClassifier:
    '''封装权重加载和 FLAIR NIfTI 病例级推理'''

    model: Any
    processor: Any
    device: torch.device
    model_dir: Path
    source: str

    @property
    def checkpoint_name(self) -> str:
        '''返回权重文件名'''

        return 'model.safetensors'

    @classmethod
    def from_pretrained(
        cls,
        source: str | Path,
        *,
        device: str = 'auto',
        cache_dir: str | Path | None = None,
        revision: str | None = None,
        local_files_only: bool = False,
    ) -> 'FlairClassifier':
        '''从本地目录或 Hugging Face 模型仓库加载权重'''

        model_dir = _resolve_model_dir(
            source,
            cache_dir=cache_dir,
            revision=revision,
            local_files_only=local_files_only,
        )
        missing = [name for name in REQUIRED_MODEL_FILES if not (model_dir / name).is_file()]
        if missing:
            raise FileNotFoundError(f'模型目录缺少文件 {missing}')
        weights_path = model_dir / 'model.safetensors'
        if weights_path.stat().st_size < 1_024:
            raise ValueError('模型权重未完整下载')
        runtime_device = _resolve_device(device)
        processor = AutoImageProcessor.from_pretrained(
            model_dir,
            local_files_only=True,
            use_fast=False,
        )
        model = AutoModelForImageClassification.from_pretrained(
            model_dir,
            local_files_only=True,
            use_safetensors=True,
        ).to(runtime_device)
        model.eval()
        labels = {str(name).strip().lower(): int(index) for name, index in model.config.label2id.items()}
        if labels != {'no': 0, 'yes': 1}:
            raise ValueError(f'模型标签不符合预期 {labels}')
        return cls(
            model=model,
            processor=processor,
            device=runtime_device,
            model_dir=model_dir,
            source=str(source),
        )

    def predict_nifti(
        self,
        path: str | Path,
        *,
        threshold: float = DEFAULT_THRESHOLD,
        max_slices: int = DEFAULT_MAX_SLICES,
        batch_size: int = 25,
    ) -> dict[str, Any]:
        '''对单份 FLAIR NIfTI 返回病例级分类与切片证据'''

        if not 0 < threshold < 1:
            raise ValueError('threshold 必须位于 0 到 1 之间')
        if batch_size <= 0:
            raise ValueError('batch_size 必须大于 0')
        volume_path = Path(path).expanduser().resolve()
        volume, indices, canonical_shape, foreground_slices, intensity_window = _prepare_volume(
            volume_path,
            max_slices=max_slices,
        )
        scores = self._predict_slices(volume, batch_size=batch_size)
        yes_probability = mean(scores)
        predicted_class = 'yes' if yes_probability >= threshold else 'no'
        return {
            'source': self.source,
            'model_dir': str(self.model_dir),
            'checkpoint': 'model.safetensors',
            'device': str(self.device),
            'input_file': str(volume_path),
            'threshold': round(threshold, 6),
            'predicted_class': predicted_class,
            'yes_probability': round(yes_probability, 6),
            'no_probability': round(1 - yes_probability, 6),
            'evaluated_slices': len(scores),
            'slice_indices': indices,
            'slice_yes_probabilities': [round(score, 6) for score in scores],
            'canonical_shape': list(canonical_shape),
            'foreground_slices': foreground_slices,
            'intensity_window': list(intensity_window),
        }

    def predict_images(self, images: list[Image.Image], *, batch_size: int) -> list[dict[str, Any]]:
        '''返回与平台兼容的逐切片分类结果'''

        if not images:
            raise ValueError('至少需要一张切片')
        if batch_size <= 0:
            raise ValueError('batch_size 必须大于 0')
        scores = self._predict_slices(images, batch_size=batch_size)
        return [
            {
                'class': 'yes' if score >= 0.5 else 'no',
                'class_id': 1 if score >= 0.5 else 0,
                'confidence': max(score, 1 - score),
                'probabilities': {'no': 1 - score, 'yes': score},
            }
            for score in scores
        ]

    def _predict_slices(self, images: list[Image.Image], *, batch_size: int) -> list[float]:
        '''计算每张切片的阳性概率'''

        scores: list[float] = []
        with torch.inference_mode():
            for start in range(0, len(images), batch_size):
                inputs = self.processor(images=images[start : start + batch_size], return_tensors='pt')
                inputs = {key: value.to(self.device) for key, value in inputs.items()}
                probabilities = torch.softmax(self.model(**inputs).logits, dim=-1).cpu().numpy()
                scores.extend(float(row[1]) for row in probabilities)
        return scores


def predict_images(
    classifier: FlairClassifier,
    images: list[Image.Image],
    *,
    batch_size: int,
) -> list[dict[str, Any]]:
    '''以函数形式提供平台兼容的逐切片推理接口'''

    return classifier.predict_images(images, batch_size=batch_size)


def _resolve_model_dir(
    source: str | Path,
    *,
    cache_dir: str | Path | None,
    revision: str | None,
    local_files_only: bool,
) -> Path:
    '''解析本地目录或下载模型仓库快照'''

    local_path = Path(source).expanduser()
    if local_path.is_dir():
        return local_path.resolve()
    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise RuntimeError('远程模型加载需要安装 huggingface-hub') from exc
    return Path(
        snapshot_download(
            repo_id=str(source),
            revision=revision,
            cache_dir=str(cache_dir) if cache_dir else None,
            local_files_only=local_files_only,
            allow_patterns=list(REQUIRED_MODEL_FILES),
        )
    ).resolve()


def _resolve_device(name: str) -> torch.device:
    '''解析推理设备'''

    if name == 'auto':
        return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    runtime_device = torch.device(name)
    if runtime_device.type == 'cuda' and not torch.cuda.is_available():
        raise RuntimeError('当前环境没有可用 CUDA 设备')
    return runtime_device


def _prepare_volume(
    path: Path,
    *,
    max_slices: int,
) -> tuple[list[Image.Image], list[int], tuple[int, int, int], int, tuple[float, float]]:
    '''读取 FLAIR 并均匀采样有效轴位切片'''

    if max_slices <= 0:
        raise ValueError('max_slices 必须大于 0')
    try:
        data = np.asarray(nib.as_closest_canonical(nib.load(str(path))).dataobj, dtype=np.float32)
    except Exception as exc:
        raise ValueError(f'无法读取输入文件 {path}') from exc
    if data.ndim != 3:
        raise ValueError(f'输入必须是三维 NIfTI 文件 当前形状为 {data.shape}')
    if not np.isfinite(data).all():
        raise ValueError('输入包含 NaN 或无穷值')
    foreground = data != 0
    if not foreground.any():
        raise ValueError('输入体积为空')
    counts = foreground.sum(axis=(0, 1))
    candidates = np.flatnonzero(counts >= max(1, int(counts.max() * 0.01)))
    positions = np.linspace(0, candidates.size - 1, num=min(max_slices, candidates.size))
    indices = candidates[np.rint(positions).astype(np.int64)]
    lower, upper = np.percentile(data[foreground].astype(np.float64, copy=False), (1, 99))
    if not np.isfinite(lower) or not np.isfinite(upper) or upper <= lower:
        raise ValueError('输入没有有效的强度变化')
    images = []
    for index in indices:
        normalized = np.clip((data[:, :, int(index)] - lower) / (upper - lower), 0, 1)
        pixels = np.rint(normalized * 255).astype(np.uint8)
        images.append(Image.fromarray(pixels, mode='L').convert('RGB'))
    return (
        images,
        [int(index) for index in indices],
        tuple(int(size) for size in data.shape),
        int(candidates.size),
        (round(float(lower), 6), round(float(upper), 6)),
    )