src/syntax_pred/infer.py

# src/syntax_pred/infer.py — инференс SYNTAX (среднее по моделям), сортировка, Study

from __future__ import annotations
from dataclasses import dataclass
from typing import List, Dict, Any, Tuple

import os
import numpy as np
import torch
import pydicom

from .config import CFG, DEVICE
from .utils import discover_weights
from .model import SyntaxLightningModule
from .preprocess import read_dicom_uint8, ensure_length_center_crop, test_like_transform
from .hf_weights import fetch_weights
from .artery_cls import classify_artery


# -------- Data structure --------
@dataclass
class Study:
    name: str
    description: str
    files: List[str]


# -------- Weights & model --------
def build_model(weight_path: str) -> SyntaxLightningModule:
    m = SyntaxLightningModule(
        num_classes=CFG.num_classes,
        lr=1e-5,
        variant=CFG.variant,
        pl_weight_path=weight_path,
        rnn_hidden_div=CFG.a_rnn.hidden_div,
        rnn_dropout=CFG.a_rnn.dropout,
        bert_nhead=CFG.bert.nhead,
        bert_layers=CFG.bert.num_layers,
        bert_ff_div=CFG.bert.ff_div,
        bert_dropout=CFG.bert.dropout,
        precision=CFG.precision,
    )
    m = m.to(DEVICE).eval()
    return m


def list_weights() -> Tuple[List[str], List[str]]:
    """
    Приоритет:
      1) CFG.weights.left/right
      2) weights/left|right
      3) HF repo (ENV WEIGHTS_REPO или CFG.weights_repo)
    """
    left = list(getattr(CFG.weights, "left", [])) or []
    right = list(getattr(CFG.weights, "right", [])) or []

    if not left:
        left = discover_weights("weights/left")
    if not right:
        right = discover_weights("weights/right")

    if not left and not right:
        repo_id = os.environ.get("WEIGHTS_REPO", getattr(CFG, "weights_repo", "")) or ""
        if repo_id:
            pulled = fetch_weights(repo_id)
            left, right = pulled["left"], pulled["right"]

    return left, right


# -------- Глобальные модели, инициализируемые при импорте --------
_LEFT_MODEL_PATHS, _RIGHT_MODEL_PATHS = list_weights()
_LEFT_MODELS: List[SyntaxLightningModule] = []
_RIGHT_MODELS: List[SyntaxLightningModule] = []

for wp in _LEFT_MODEL_PATHS:
    try:
        _LEFT_MODELS.append(build_model(wp))
    except Exception as e:
        print(f"[WARN] failed to init LEFT model {wp}: {e}")

for wp in _RIGHT_MODEL_PATHS:
    try:
        _RIGHT_MODELS.append(build_model(wp))
    except Exception as e:
        print(f"[WARN] failed to init RIGHT model {wp}: {e}")

if not _LEFT_MODELS and not _RIGHT_MODELS:
    print("[ERROR] No models initialized: check weights/ and HF repo configuration.")


# -------- Stable sort --------
def _stable_sort_paths(paths: List[str]) -> List[str]:
    """
    Сортировка в духе старого пайплайна:
      (SeriesInstanceUID, basename, fullpath)
    """
    def key(p: str):
        try:
            ds = pydicom.dcmread(p, stop_before_pixels=True, specific_tags=["SeriesInstanceUID"])
            series = str(getattr(ds, "SeriesInstanceUID", ""))
        except Exception:
            series = ""
        return (series, os.path.basename(p), p)
    return sorted(paths, key=key)


def _filter_dicom_paths(paths: List[str]) -> List[str]:
    """Защитная фильтрация: оставляем только корректные DICOM."""
    out = []
    for p in paths:
        if not os.path.exists(p):
            continue
        try:
            pydicom.dcmread(p, stop_before_pixels=True)
            out.append(p)
        except Exception:
            continue
    return out


# -------- Packing one side --------
def _pack_study_side_to_tensor(
    file_paths: List[str],
    frames_per_clip: int,
    video_size: Tuple[int, int],
) -> torch.Tensor | None:
    if not file_paths:
        return None
    tx = test_like_transform(video_size)
    clips = []
    for p in file_paths:
        arr = read_dicom_uint8(p)
        arr = ensure_length_center_crop(arr, frames_per_clip)
        thwc = np.stack([arr, arr, arr], axis=-1)  # (T,H,W,3)
        cthw = tx(torch.tensor(thwc, dtype=torch.uint8))
        clips.append(cthw)
    if not clips:
        return None
    return torch.stack(clips, dim=0).unsqueeze(0)  # (1,S,C,T,H,W)


@torch.no_grad()
def _score_side_by_models(
    side_paths: List[str],
    models: List[SyntaxLightningModule],
    model_paths: List[str],
    frames_per_clip: int,
    video_size: Tuple[int, int],
) -> Dict[str, Any]:
    if not side_paths:
        return {"mean": 0.0, "per_model": [], "n_files": 0}

    x = _pack_study_side_to_tensor(side_paths, frames_per_clip, video_size)
    if x is None:
        return {"mean": 0.0, "per_model": [], "n_files": 0}
    x = x.to(DEVICE)

    per_model_scores: List[float] = []
    used_models: List[str] = []

    for m, wp in zip(models, model_paths):
        try:
            y = m(x)  # (1,2)
            reg_log = float(y[0, 1].detach().cpu().numpy())
            score = float(max(0.0, np.exp(reg_log) - 1.0))  # inverse log(1+score)
            per_model_scores.append(score)
            used_models.append(os.path.basename(wp))
        except Exception as e:
            print(f"[WARN] model {wp} failed: {e}")

    mean_score = float(np.mean(per_model_scores)) if per_model_scores else 0.0
    return {
        "mean": mean_score,
        "per_model": [{"model": n, "score": round(s, 3)} for n, s in zip(used_models, per_model_scores)],
        "n_files": len(side_paths),
    }


# -------- Top-level inference --------
def run_inference(studies: List[Study]) -> Dict[str, Any]:
    if not _LEFT_MODEL_PATHS and not _RIGHT_MODEL_PATHS:
        return {"error": "No weights found. Upload to weights/left and weights/right."}
    if not _LEFT_MODELS and not _RIGHT_MODELS:
        return {"error": "Models failed to initialize. See logs."}

    results = {"studies": []}
    thr = CFG.thresholds.both
    video_size = tuple(CFG.video_size)
    frames = CFG.frames_per_clip

    for st in studies:
        st_files = _filter_dicom_paths(st.files)

        cls = classify_artery(st_files)  # {"left":[{path,prob}], "right":[...], "other":[...]}

        left_paths = _stable_sort_paths([r["path"] for r in cls["left"]])
        right_paths = _stable_sort_paths([r["path"] for r in cls["right"]])

        left_res = _score_side_by_models(
            left_paths,
            _LEFT_MODELS,
            _LEFT_MODEL_PATHS,
            frames,
            video_size,
        )
        right_res = _score_side_by_models(
            right_paths,
            _RIGHT_MODELS,
            _RIGHT_MODEL_PATHS,
            frames,
            video_size,
        )

        total = left_res["mean"] + right_res["mean"]

        results["studies"].append({
            "study": st.name,
            "description": st.description or "",
            "left": {
                "mean": round(left_res["mean"], 3),
                "per_model": left_res["per_model"],
                "n_files": left_res["n_files"],
                "files": cls["left"],
            },
            "right": {
                "mean": round(right_res["mean"], 3),
                "per_model": right_res["per_model"],
                "n_files": right_res["n_files"],
                "files": cls["right"],
            },
            "other": cls["other"],
            "total": {
                "mean": round(total, 3),
                f"High-risk (≥{thr:.1f})": bool(total >= thr),
            },
        })
    return results