| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import nibabel as nib |
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset |
|
|
|
|
| def normalize_pet(volume: np.ndarray, eps: float = 1e-6) -> np.ndarray: |
| volume = np.asarray(volume, dtype=np.float32) |
| finite = np.isfinite(volume) |
| if not finite.any(): |
| return np.zeros_like(volume, dtype=np.float32) |
| lo, hi = np.percentile(volume[finite], [0.5, 99.5]) |
| volume = np.clip(volume, lo, hi) |
| volume = (volume - lo) / max(float(hi - lo), eps) |
| return volume.astype(np.float32, copy=False) |
|
|
|
|
| def resize_volume(volume: np.ndarray, output_size: tuple[int, int, int]) -> torch.Tensor: |
| tensor = torch.from_numpy(volume)[None, None] |
| tensor = F.interpolate(tensor, size=output_size, mode="trilinear", align_corners=False) |
| return tensor[0] |
|
|
|
|
| def load_suvr_vector(csv_path: str | Path, include_background: bool = False) -> tuple[list[str], torch.Tensor]: |
| df = pd.read_csv(csv_path) |
| if not include_background: |
| df = df[df["label_name"] != "Background"].copy() |
| labels = df["label_name"].astype(str).tolist() |
| values = torch.tensor(df["mean_scalar"].astype(float).to_numpy(), dtype=torch.float32) |
| return labels, values |
|
|
|
|
| def suvr_to_text(labels: list[str], values: torch.Tensor, top_k: int = 8) -> str: |
| pairs = sorted(zip(labels, values.tolist()), key=lambda x: x[1], reverse=True) |
| high = ", ".join(f"{name} {value:.3f}" for name, value in pairs[:top_k]) |
| low = ", ".join(f"{name} {value:.3f}" for name, value in pairs[-top_k:]) |
| mean_value = float(values.mean()) |
| return f"FDG-PET regional SUVR summary. Mean SUVR {mean_value:.3f}. Highest regions: {high}. Lowest regions: {low}." |
|
|
|
|
| class PETSUVRDataset(Dataset): |
| def __init__( |
| self, |
| manifest_path: str | Path, |
| output_size: tuple[int, int, int] = (96, 96, 96), |
| include_background: bool = False, |
| ) -> None: |
| self.manifest = pd.read_csv(manifest_path) |
| self.output_size = output_size |
| self.include_background = include_background |
|
|
| def __len__(self) -> int: |
| return len(self.manifest) |
|
|
| def __getitem__(self, index: int) -> dict[str, object]: |
| row = self.manifest.iloc[index] |
| volume = nib.load(str(row["pet_path"])).get_fdata(dtype=np.float32) |
| volume = normalize_pet(volume) |
| image = resize_volume(volume, self.output_size) |
| labels, suvr = load_suvr_vector(row["suvr_csv_path"], self.include_background) |
| text = suvr_to_text(labels, suvr) |
| return { |
| "sample_id": row["sample_id"], |
| "image": image, |
| "suvr": suvr, |
| "region_labels": labels, |
| "text": text, |
| } |
|
|
|
|
| def collate_pet_suvr(batch: list[dict[str, object]]) -> dict[str, object]: |
| return { |
| "sample_id": [item["sample_id"] for item in batch], |
| "image": torch.stack([item["image"] for item in batch]), |
| "suvr": torch.stack([item["suvr"] for item in batch]), |
| "text": [item["text"] for item in batch], |
| "region_labels": batch[0]["region_labels"], |
| } |
|
|
