predictions class added
This commit is contained in:
rpotter6298
@@ -0,0 +1,194 @@
|
||||
"""PredictionStore — unified per-epoch prediction tensor across all folds.
|
||||
|
||||
Tensor shape: (n_folds, n_epochs, n_samples, n_heads, n_classes)
|
||||
|
||||
The meaning of "sample" depends on tower_mode:
|
||||
single — each eye is a sample; sample_ids like "5OD", "14OS"
|
||||
ensemble — each patient is a sample; sample_ids like "5", "14"
|
||||
fused — same as ensemble
|
||||
bilateral— same as ensemble
|
||||
|
||||
Head names by mode:
|
||||
single : ["fused", "img", "md"]
|
||||
ensemble : ["od_fused", "od_img", "od_md", "os_fused", "os_img", "os_md"]
|
||||
fused : ["od_fused", "od_img", "od_md", "os_fused", "os_img", "os_md", "bilat_fused"]
|
||||
bilateral : ["fused", "img_joint", "md_joint"]
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def head_names_for_mode(tower_mode: str, *, fused_head: bool = False) -> list[str]:
|
||||
"""Return canonical head name list for a given tower_mode."""
|
||||
if tower_mode in ("single", "classic"):
|
||||
return ["fused", "img", "md"]
|
||||
if tower_mode == "ensemble":
|
||||
names = ["od_fused", "od_img", "od_md", "os_fused", "os_img", "os_md"]
|
||||
return names + ["bilat_fused"] if fused_head else names
|
||||
if tower_mode == "bilateral":
|
||||
return ["fused", "img_joint", "md_joint"]
|
||||
raise ValueError(f"Unknown tower_mode: {tower_mode!r}")
|
||||
|
||||
|
||||
class PredictionStore:
|
||||
"""
|
||||
Stores per-epoch predictions for every sample, head, and fold in one tensor.
|
||||
|
||||
Usage
|
||||
-----
|
||||
# Build once before the fold loop:
|
||||
store = PredictionStore(
|
||||
sample_ids=all_eye_or_patient_ids,
|
||||
y_true=all_labels,
|
||||
head_names=head_names_for_mode(tower_mode, fused_head=args.fused_head),
|
||||
n_folds=n_folds,
|
||||
n_epochs=total_epochs,
|
||||
n_classes=num_classes,
|
||||
)
|
||||
|
||||
# Inside each epoch, after collecting probs:
|
||||
store.record(fold, epoch, patient_ids_batch, "od_fused", probs_od)
|
||||
store.set_split(fold, train_ids, "train")
|
||||
store.set_split(fold, val_ids, "val")
|
||||
|
||||
# After all folds:
|
||||
store.save(run_dir / "predictions.npz")
|
||||
|
||||
# Load and query:
|
||||
store = PredictionStore.load("predictions.npz")
|
||||
store.query("5", "od_fused", fold=0) # → (n_epochs, n_classes)
|
||||
store.query("5", "od_fused") # → (n_folds, n_epochs, n_classes)
|
||||
store.get_split("5", fold=0) # → "train"
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sample_ids: Sequence[str],
|
||||
y_true: Sequence[int],
|
||||
head_names: Sequence[str],
|
||||
n_folds: int,
|
||||
n_epochs: int,
|
||||
n_classes: int,
|
||||
):
|
||||
self.sample_ids = np.array(sample_ids, dtype=object)
|
||||
self.y_true = np.array(y_true, dtype=np.int64)
|
||||
self.head_names = np.array(head_names, dtype=object)
|
||||
self.n_folds = n_folds
|
||||
self.n_epochs = n_epochs
|
||||
self.n_classes = n_classes
|
||||
|
||||
n_samples = len(self.sample_ids)
|
||||
n_heads = len(self.head_names)
|
||||
|
||||
self.probs = np.full(
|
||||
(n_folds, n_epochs, n_samples, n_heads, n_classes),
|
||||
fill_value=np.nan,
|
||||
dtype=np.float32,
|
||||
)
|
||||
self.split = np.full((n_folds, n_samples), fill_value="", dtype=object)
|
||||
|
||||
self._sid_index: dict[str, int] = {str(s): i for i, s in enumerate(self.sample_ids)}
|
||||
self._head_index: dict[str, int] = {str(h): i for i, h in enumerate(self.head_names)}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Writing
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def record(
|
||||
self,
|
||||
fold: int,
|
||||
epoch: int,
|
||||
sample_ids: Sequence[str],
|
||||
head_name: str,
|
||||
probs: np.ndarray,
|
||||
) -> None:
|
||||
"""Record a batch of predictions for one head.
|
||||
|
||||
Args:
|
||||
fold: 0-indexed fold number
|
||||
epoch: 0-indexed epoch number
|
||||
sample_ids: sequence of sample ID strings (length B)
|
||||
head_name: which head — must be in self.head_names
|
||||
probs: (B, n_classes) probability array
|
||||
"""
|
||||
head_idx = self._head_index.get(head_name)
|
||||
if head_idx is None:
|
||||
return # head not active in this mode — skip silently
|
||||
for i, sid in enumerate(sample_ids):
|
||||
s_idx = self._sid_index.get(str(sid))
|
||||
if s_idx is not None:
|
||||
self.probs[fold, epoch, s_idx, head_idx, :] = probs[i]
|
||||
|
||||
def set_split(
|
||||
self,
|
||||
fold: int,
|
||||
sample_ids: Sequence[str],
|
||||
label: str,
|
||||
) -> None:
|
||||
"""Label a group of samples as 'train', 'val', or 'holdout' for a fold."""
|
||||
for sid in sample_ids:
|
||||
s_idx = self._sid_index.get(str(sid))
|
||||
if s_idx is not None:
|
||||
self.split[fold, s_idx] = label
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Querying
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def query(
|
||||
self,
|
||||
sample_id: str,
|
||||
head_name: str,
|
||||
fold: int | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Return epoch-level predictions for one sample + head.
|
||||
|
||||
Returns:
|
||||
fold=None → (n_folds, n_epochs, n_classes)
|
||||
fold=int → (n_epochs, n_classes)
|
||||
"""
|
||||
s_idx = self._sid_index[str(sample_id)]
|
||||
head_idx = self._head_index[str(head_name)]
|
||||
if fold is None:
|
||||
return self.probs[:, :, s_idx, head_idx, :]
|
||||
return self.probs[fold, :, s_idx, head_idx, :]
|
||||
|
||||
def get_split(self, sample_id: str, fold: int) -> str:
|
||||
"""Return the split label ('train'/'val'/'holdout') for a sample in a fold."""
|
||||
s_idx = self._sid_index[str(sample_id)]
|
||||
return str(self.split[fold, s_idx])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def save(self, path: str | Path) -> None:
|
||||
np.savez_compressed(
|
||||
path,
|
||||
probs=self.probs,
|
||||
split=self.split,
|
||||
sample_ids=self.sample_ids,
|
||||
y_true=self.y_true,
|
||||
head_names=self.head_names,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str | Path) -> "PredictionStore":
|
||||
data = np.load(path, allow_pickle=True)
|
||||
probs = data["probs"]
|
||||
n_folds, n_epochs, _, _, n_classes = probs.shape
|
||||
store = cls(
|
||||
sample_ids=data["sample_ids"].tolist(),
|
||||
y_true=data["y_true"],
|
||||
head_names=data["head_names"].tolist(),
|
||||
n_folds=n_folds,
|
||||
n_epochs=n_epochs,
|
||||
n_classes=n_classes,
|
||||
)
|
||||
store.probs = probs
|
||||
store.split = data["split"]
|
||||
return store
|
||||
Reference in New Issue
Block a user