File size: 5,989 Bytes
f621d73 75c3625 f621d73 75c3625 f621d73 75c3625 f621d73 75c3625 f621d73 75c3625 f621d73 75c3625 f621d73 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | import os
import json
from typing import Callable, Optional, Tuple, Any
from pathlib import Path
import numpy as np
import pydicom
import torch
from torch import Tensor
from torch.utils.data import Dataset
class SyntaxDataset(Dataset):
"""
PyTorch dataset for training the 3D backbone on DICOM videos.
Expected JSON entries contain:
- "path": relative path to the DICOM file from the JSON directory
- "artery": 0 (left) or 1 (right)
- "<label>": numeric SYNTAX score value, for example "syntax_left"
- videos are loaded from multi-frame DICOM files via pydicom as 3D arrays
"""
def __init__(
self,
root: str,
meta: str,
train: bool,
length: int,
label: str,
artery_bin: int,
validation: bool = False,
transform: Optional[Callable] = None,
) -> None:
self.root = Path(root).resolve()
self.train = train
self.length = int(length)
self.label = label
self.transform = transform
self.validation = validation
if artery_bin not in (0, 1):
raise ValueError("artery_bin must be 0 (left) or 1 (right)")
self.artery_bin = artery_bin
meta_path = meta if os.path.isabs(meta) else self.root / meta
meta_path = Path(meta_path).resolve()
json_dir = meta_path.parent
print(f"Backbone dataset: root={self.root}, meta={meta_path}, json_dir={json_dir}")
with open(meta_path, "r", encoding="utf-8") as f:
dataset = json.load(f)
dataset = [rec for rec in dataset if rec.get("artery") == artery_bin]
if validation:
dataset = [rec for rec in dataset if float(rec.get(self.label, 0.0)) > 0]
self.json_dir = json_dir
self.dataset = dataset
for rec in self.dataset:
rec.setdefault("weight", 1.0)
print(f"Backbone dataset loaded: {len(self.dataset)} samples after filtering")
def get_sample_weights(self) -> Tensor:
"""
Return sample weights for WeightedRandomSampler.
Logic:
- split records into score bins separately for the left and right artery;
- rarer bins receive larger weights.
"""
bin_thresholds = {
0: [0, 5, 10, 15],
1: [0, 2, 5, 8],
}
thr0, thr1, thr2, thr3 = bin_thresholds[self.artery_bin]
def in_bin(score: float) -> int:
if score == thr0:
return 0
if thr0 < score <= thr1:
return 1
if thr1 < score <= thr2:
return 2
if thr2 < score <= thr3:
return 3
return 4
scores = [float(rec.get(self.label, 0.0)) for rec in self.dataset]
bins = [in_bin(s) for s in scores]
counts = np.bincount(np.array(bins, dtype=np.int64), minlength=5)
total = int(counts.sum())
weights_by_bin = np.array(
[(total / counts[b]) if counts[b] > 0 else 0.0 for b in range(5)],
dtype=np.float64,
)
weights = np.array([weights_by_bin[b] for b in bins], dtype=np.float64)
return torch.as_tensor(weights, dtype=torch.double)
def __len__(self) -> int:
return len(self.dataset)
def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor, Tensor, float, str, Tensor]:
"""
Return:
video: Tensor video clip (T, H, W, C) before transform
label: Tensor(1,) binary classification label
target: Tensor(1,) regression target (log1p(score))
sample_weight: float original record weight from JSON
path: str relative DICOM path as stored in JSON
original_label: Tensor(1,) original score value
"""
rec = self.dataset[idx]
rel_path = rec["path"]
sample_weight = float(rec.get("weight", 1.0))
full_path = (self.json_dir / rel_path).resolve()
if not full_path.exists():
raise FileNotFoundError(
f"DICOM not found: {full_path}\n"
f" json_dir={self.json_dir}\n"
f" rel_path='{rel_path}'"
)
video = pydicom.dcmread(str(full_path)).pixel_array
if video.ndim != 3:
raise ValueError(f"Expected 3D video array, got shape={video.shape} for {rel_path}")
if video.shape[0] > 128 and video.shape[-1] <= 128:
video = np.moveaxis(video, -1, 0)
if video.dtype == np.uint16:
vmax = int(np.max(video))
if vmax <= 0:
raise ValueError(f"Invalid vmax={vmax} for {rel_path}")
video = (video.astype(np.float32) * (255.0 / vmax)).clip(0, 255).astype(np.uint8)
else:
video = video.astype(np.uint8)
score = float(rec.get(self.label, 0.0))
bin_thresholds = {
0: 15,
1: 5,
}
label = torch.tensor(
[1.0 if score > bin_thresholds[self.artery_bin] else 0.0],
dtype=torch.float32,
)
target = torch.tensor([np.log1p(score)], dtype=torch.float32)
original_label = torch.tensor([score], dtype=torch.float32)
while video.shape[0] < self.length:
video = np.concatenate([video, video], axis=0)
t = int(video.shape[0])
if self.train:
begin = torch.randint(low=0, high=t - self.length + 1, size=(1,)).item()
video = video[begin: begin + self.length]
else:
video = video[:self.length]
video = torch.from_numpy(np.stack([video, video, video], axis=-1))
if self.transform is not None:
video = self.transform(video)
return video, label, target, sample_weight, str(rel_path), original_label
|