LA Train

1
import os
2
import torch
3
import numpy as np
4
from glob import glob
5
from torch.utils.data import Dataset
6
import h5py
7
import itertools
8
from torch.utils.data.sampler import Sampler
9

10

11
class BraTS2019(Dataset):
12
    """ BraTS2019 Dataset """
13

14
    def __init__(self, base_dir=None, split='train', num=None, transform=None):
15
        self._base_dir = base_dir
16
        self.transform = transform
17
        self.sample_list = []
18

19
        train_path = self._base_dir+'/train.txt'
20
        test_path = self._base_dir+'/val.txt'
21

22
        if split == 'train':
23
            with open(train_path, 'r') as f:
24
                self.image_list = f.readlines()
25
        elif split == 'test':
26
            with open(test_path, 'r') as f:
27
                self.image_list = f.readlines()
28

29
        self.image_list = [item.replace('\n', '').split(",")[0] for item in self.image_list]
30
        if num is not None:
31
            self.image_list = self.image_list[:num]
32
        print("total {} samples".format(len(self.image_list)))
33

34
    def __len__(self):
35
        return len(self.image_list)
36

37
    def __getitem__(self, idx):
38
        image_name = self.image_list[idx]
39
        h5f = h5py.File(self._base_dir + "/data/{}.h5".format(image_name), 'r')
40
        image = h5f['image'][:]
41
        label = h5f['label'][:]
42
        sample = {'image': image, 'label': label.astype(np.uint8)}
43
        if self.transform:
44
            sample = self.transform(sample)
45
        return sample
46

47

48
class CenterCrop(object):
49
    def __init__(self, output_size):
50
        self.output_size = output_size
51

52
    def __call__(self, sample):
53
        image, label = sample['image'], sample['label']
54

55
        # pad the sample if necessary
56
        if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \
57
                self.output_size[2]:
58
            pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0)
59
            ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0)
60
            pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0)
61
            image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)],
62
                           mode='constant', constant_values=0)
63
            label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)],
64
                           mode='constant', constant_values=0)
65

66
        (w, h, d) = image.shape
67

68
        w1 = int(round((w - self.output_size[0]) / 2.))
69
        h1 = int(round((h - self.output_size[1]) / 2.))
70
        d1 = int(round((d - self.output_size[2]) / 2.))
71

72
        label = label[w1:w1 + self.output_size[0], h1:h1 +
73
                      self.output_size[1], d1:d1 + self.output_size[2]]
74
        image = image[w1:w1 + self.output_size[0], h1:h1 +
75
                      self.output_size[1], d1:d1 + self.output_size[2]]
76

77
        return {'image': image, 'label': label}
78

79

80
class RandomCrop(object):
81
    """
82
    Crop randomly the image in a sample
83
    Args:
84
    output_size (int): Desired output size
85
    """
86

87
    def __init__(self, output_size, with_sdf=False):
88
        self.output_size = output_size
89
        self.with_sdf = with_sdf
90

91
    def __call__(self, sample):
92
        image, label = sample['image'], sample['label']
93
        if self.with_sdf:
94
            sdf = sample['sdf']
95

96
        # pad the sample if necessary
97
        if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \
98
                self.output_size[2]:
99
            pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0)
100
            ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0)
101
            pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0)
102
            image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)],
103
                           mode='constant', constant_values=0)
104
            label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)],
105
                           mode='constant', constant_values=0)
106
            if self.with_sdf:
107
                sdf = np.pad(sdf, [(pw, pw), (ph, ph), (pd, pd)],
108
                             mode='constant', constant_values=0)
109

110
        (w, h, d) = image.shape
111
        # if np.random.uniform() > 0.33:
112
        #     w1 = np.random.randint((w - self.output_size[0])//4, 3*(w - self.output_size[0])//4)
113
        #     h1 = np.random.randint((h - self.output_size[1])//4, 3*(h - self.output_size[1])//4)
114
        # else:
115
        w1 = np.random.randint(0, w - self.output_size[0])
116
        h1 = np.random.randint(0, h - self.output_size[1])
117
        d1 = np.random.randint(0, d - self.output_size[2])
118

119
        label = label[w1:w1 + self.output_size[0], h1:h1 +
120
                      self.output_size[1], d1:d1 + self.output_size[2]]
121
        image = image[w1:w1 + self.output_size[0], h1:h1 +
122
                      self.output_size[1], d1:d1 + self.output_size[2]]
123
        if self.with_sdf:
124
            sdf = sdf[w1:w1 + self.output_size[0], h1:h1 +
125
                      self.output_size[1], d1:d1 + self.output_size[2]]
126
            return {'image': image, 'label': label, 'sdf': sdf}
127
        else:
128
            return {'image': image, 'label': label}
129

130

131
class RandomRotFlip(object):
132
    """
133
    Crop randomly flip the dataset in a sample
134
    Args:
135
    output_size (int): Desired output size
136
    """
137

138
    def __call__(self, sample):
139
        image, label = sample['image'], sample['label']
140
        k = np.random.randint(0, 4)
141
        image = np.rot90(image, k)
142
        label = np.rot90(label, k)
143
        axis = np.random.randint(0, 2)
144
        image = np.flip(image, axis=axis).copy()
145
        label = np.flip(label, axis=axis).copy()
146

147
        return {'image': image, 'label': label}
148

149

150
class RandomNoise(object):
151
    def __init__(self, mu=0, sigma=0.1):
152
        self.mu = mu
153
        self.sigma = sigma
154

155
    def __call__(self, sample):
156
        image, label = sample['image'], sample['label']
157
        noise = np.clip(self.sigma * np.random.randn(
158
            image.shape[0], image.shape[1], image.shape[2]), -2*self.sigma, 2*self.sigma)
159
        noise = noise + self.mu
160
        image = image + noise
161
        return {'image': image, 'label': label}
162

163

164
class CreateOnehotLabel(object):
165
    def __init__(self, num_classes):
166
        self.num_classes = num_classes
167

168
    def __call__(self, sample):
169
        image, label = sample['image'], sample['label']
170
        onehot_label = np.zeros(
171
            (self.num_classes, label.shape[0], label.shape[1], label.shape[2]), dtype=np.float32)
172
        for i in range(self.num_classes):
173
            onehot_label[i, :, :, :] = (label == i).astype(np.float32)
174
        return {'image': image, 'label': label, 'onehot_label': onehot_label}
175

176

177
class ToTensor(object):
178
    """Convert ndarrays in sample to Tensors."""
179

180
    def __call__(self, sample):
181
        image = sample['image']
182
        image = image.reshape(
183
            1, image.shape[0], image.shape[1], image.shape[2]).astype(np.float32)
184
        if 'onehot_label' in sample:
185
            return {'image': torch.from_numpy(image), 'label': torch.from_numpy(sample['label']).long(),
186
                    'onehot_label': torch.from_numpy(sample['onehot_label']).long()}
187
        else:
188
            return {'image': torch.from_numpy(image), 'label': torch.from_numpy(sample['label']).long()}
189

190

191
class TwoStreamBatchSampler(Sampler):
192
    """Iterate two sets of indices
193

194
    An 'epoch' is one iteration through the primary indices.
195
    During the epoch, the secondary indices are iterated through
196
    as many times as needed.
197
    """
198

199
    def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size):
200
        self.primary_indices = primary_indices
201
        self.secondary_indices = secondary_indices
202
        self.secondary_batch_size = secondary_batch_size
203
        self.primary_batch_size = batch_size - secondary_batch_size
204

205
        assert len(self.primary_indices) >= self.primary_batch_size > 0
206
        assert len(self.secondary_indices) >= self.secondary_batch_size > 0
207

208
    def __iter__(self):
209
        primary_iter = iterate_once(self.primary_indices)
210
        secondary_iter = iterate_eternally(self.secondary_indices)
211
        return (
212
            primary_batch + secondary_batch
213
            for (primary_batch, secondary_batch)
214
            in zip(grouper(primary_iter, self.primary_batch_size),
215
                   grouper(secondary_iter, self.secondary_batch_size))
216
        )
217

218
    def __len__(self):
219
        return len(self.primary_indices) // self.primary_batch_size
220

221

222
def iterate_once(iterable):
223
    return np.random.permutation(iterable)
224

225

226
def iterate_eternally(indices):
227
    def infinite_shuffles():
228
        while True:
229
            yield np.random.permutation(indices)
230
    return itertools.chain.from_iterable(infinite_shuffles())
231

232

233
def grouper(iterable, n):
234
    "Collect data into fixed-length chunks or blocks"
235
    # grouper('ABCDEFG', 3) --> ABC DEF"
236
    args = [iter(iterable)] * n
237
    return zip(*args)

1
import os
2
import argparse
3
import torch
4
import numpy as np
5
import h5py
6
import nibabel as nib
7
from tqdm import tqdm
8
from medpy import metric
9

10
from networks.net_factory import net_factory
11
from utils import val_2d
12

13
from thop import profile, clever_format
14

15

16
parser = argparse.ArgumentParser()
17
parser.add_argument('--root_path',   type=str,   default='../data/ACDC',  help='Root path of dataset')
18
parser.add_argument('--exp',         type=str,   default='CPAM',          help='Experiment name')
19
parser.add_argument('--model',       type=str,   default='unet',          help='Model name')
20
parser.add_argument('--gpu',         type=str,   default='0',             help='GPU to use')
21
parser.add_argument('--detail',      type=int,   default=1,               help='Print metrics for every sample?')
22
parser.add_argument('--nms',         type=int,   default=1,               help='Apply NMS (LargestCC) post-processing?')
23
parser.add_argument('--labelnum',    type=int,   default=3,               help='Number of labeled patients used in training')
24
parser.add_argument('--stage_name',  type=str,   default='self_train',    help='Stage: self_train or pre_train')
25
FLAGS = parser.parse_args()
26

27
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
28

29
# ── Paths (mirror the training script conventions) ──────────────────────────
30
snapshot_path  = "./model/CPAM/ACDC_{}_{}_labeled/{}".format(
31
    FLAGS.exp, FLAGS.labelnum, FLAGS.stage_name)
32
test_save_path = "./model/CPAM/ACDC_{}_{}_labeled/{}_predictions/".format(
33
    FLAGS.exp, FLAGS.labelnum, FLAGS.model)
34

35
num_classes = 4   # background(0) + RV(1) + Myo(2) + LV(3)
36

37
if not os.path.exists(test_save_path):
38
    os.makedirs(test_save_path)
39
print("Predictions will be saved to:", test_save_path)
40

41
# ── Load test list ────────────────────────────────────────────────────────────
42
# test.list entries look like: "patient001_frame01"  (one per line)
43
with open(os.path.join(FLAGS.root_path, 'test.list'), 'r') as f:
44
    image_list = [line.strip() for line in f.readlines() if line.strip()]
45

46
# Full path to each patient .h5 file, e.g.  ../data/ACDC/data/patient001_frame01.h5
47
image_list = [os.path.join(FLAGS.root_path, "data", item + ".h5")
48
              for item in image_list]
49

50
print("Total test cases:", len(image_list))
51

52

53
# ── NIfTI saving ─────────────────────────────────────────────────────────────
54
def save_nii(array, save_path, dtype=np.float32):
55
    """Save a numpy array (D, H, W) as a .nii.gz file."""
56
    nii_img = nib.Nifti1Image(array.astype(dtype), affine=np.eye(4))
57
    nib.save(nii_img, save_path)
58

59

60
# ── Metric helpers ────────────────────────────────────────────────────────────
61
def calculate_metric_percase(pred, gt):
62
    """Return (dice, jc, hd, asd) for a single binary foreground mask."""
63
    if pred.sum() > 0 and gt.sum() > 0:
64
        dice = metric.binary.dc(pred, gt)
65
        jc   = metric.binary.jc(pred, gt)
66
        hd   = metric.binary.hd95(pred, gt)
67
        asd  = metric.binary.asd(pred, gt)
68
    elif pred.sum() == 0 and gt.sum() == 0:
69
        dice, jc, hd, asd = 1.0, 1.0, 0.0, 0.0
70
    else:
71
        dice, jc, hd, asd = 0.0, 0.0, 100.0, 100.0
72
    return dice, jc, hd, asd
73

74

75
# ── Model complexity ──────────────────────────────────────────────────────────
76
def compute_model_complexity(model, input_size=(1, 1, 256, 256)):
77
    """
78
    Print parameter count and FLOPs for a single 2-D slice (B=1, C=1, H, W).
79
    Uses thop if available, otherwise falls back to manual param counting.
80
    """
81
    print("\n" + "=" * 50)
82
    print("Model Complexity")
83
    print("-" * 50)
84

85
    # ── Parameters ──────────────────────────────────────────────────────────
86
    total_params     = sum(p.numel() for p in model.parameters())
87
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
88
    print(f"  Total params    : {total_params:,}  "
89
          f"({total_params / 1e6:.2f} M)")
90
    print(f"  Trainable params: {trainable_params:,}  "
91
          f"({trainable_params / 1e6:.2f} M)")
92

93
    # ── FLOPs via thop ───────────────────────────────────────────────────────
94
    dummy = torch.randn(input_size).cuda()
95
    model.eval()
96
    with torch.no_grad():
97
        try:
98
            flops, _ = profile(model, inputs=(dummy,), verbose=False)
99
            flops_str, _ = clever_format([flops, total_params], "%.3f")
100
            print(f"  FLOPs (per slice): {flops_str}  "
101
                f"({flops / 1e9:.3f} GFLOPs)")
102
        except Exception as e:
103
            print(f"  FLOPs calculation failed: {e}")
104
    print("=" * 50 + "\n")
105

106

107
# ── Main evaluation ───────────────────────────────────────────────────────────
108
def test_calculate_metric():
109
    # Load model
110
    model = net_factory(net_type=FLAGS.model, in_chns=1,
111
                        class_num=num_classes, mode="train")
112
    save_model_path = os.path.join(snapshot_path,
113
                                   '{}_best_model.pth'.format(FLAGS.model))
114
    model.load_state_dict(torch.load(save_model_path))
115
    print("Loaded weights from:", save_model_path)
116
    model.eval()
117

118
    # ── Parameters & FLOPs ───────────────────────────────────────────────────
119
    compute_model_complexity(model, input_size=(1, 1, 256, 256))
120

121
    # Accumulate metrics across all test cases  (3 foreground classes)
122
    # Shape: [num_cases, num_fg_classes, 4_metrics]
123
    all_metrics = []
124

125
    for case_path in tqdm(image_list, desc="Testing"):
126
        # ── Read h5 ──────────────────────────────────────────────────────────
127
        with h5py.File(case_path, 'r') as f:
128
            image = f['image'][:]   # (H, W, D) or (D, H, W) — adjust if needed
129
            label = f['label'][:]   # same shape, integer class labels
130

131
        # ── Save dir ─────────────────────────────────────────────────────────
132
        case_name = os.path.basename(case_path).replace('.h5', '')
133
        case_save_dir = os.path.join(test_save_path, case_name)
134
        os.makedirs(case_save_dir, exist_ok=True)
135

136
        # ── Inference: build full prediction volume slice-by-slice ────────────
137
        from scipy.ndimage import zoom as scipy_zoom
138
        patch_size = [256, 256]
139
        prediction = np.zeros_like(label)
140
        for ind in range(image.shape[0]):
141
            slc = image[ind, :, :]
142
            x, y = slc.shape
143
            slc_resized = scipy_zoom(slc, (patch_size[0] / x, patch_size[1] / y), order=0)
144
            inp = torch.from_numpy(slc_resized).unsqueeze(0).unsqueeze(0).float().cuda()
145
            with torch.no_grad():
146
                output = model(inp)
147
                if len(output) > 1:
148
                    output = output[0]
149
                out = torch.argmax(torch.softmax(output, dim=1), dim=1).squeeze(0)
150
                out = out.cpu().numpy()
151
                pred_slc = scipy_zoom(out, (x / patch_size[0], y / patch_size[1]), order=0)
152
                prediction[ind] = pred_slc
153

154
        # ── Save .nii.gz ─────────────────────────────────────────────────────
155
        save_nii(image,      os.path.join(case_save_dir, "image.nii.gz"), dtype=np.float32)
156
        save_nii(label,      os.path.join(case_save_dir, "gt.nii.gz"),    dtype=np.uint8)
157
        save_nii(prediction, os.path.join(case_save_dir, "pred.nii.gz"),  dtype=np.uint8)
158

159
        # ── Compute 4 metrics per foreground class ────────────────────────────
160
        metric_i = []
161
        for c in range(1, num_classes):
162
            pred_c = (prediction == c)
163
            gt_c   = (label == c)
164
            if pred_c.sum() > 0 and gt_c.sum() > 0:
165
                dice = metric.binary.dc(pred_c, gt_c)
166
                jc   = metric.binary.jc(pred_c, gt_c)
167
                hd95 = metric.binary.hd95(pred_c, gt_c)
168
                asd  = metric.binary.asd(pred_c, gt_c)
169
            elif pred_c.sum() == 0 and gt_c.sum() == 0:
170
                dice, jc, hd95, asd = 1.0, 1.0, 0.0, 0.0
171
            else:
172
                dice, jc, hd95, asd = 0.0, 0.0, 100.0, 100.0
173
            metric_i.append((dice, jc, hd95, asd))
174

175
        if FLAGS.detail:
176
            for c_idx, (dice, jc, hd95, asd) in enumerate(metric_i, start=1):
177
                print(f"  {case_name}  class {c_idx}: "
178
                      f"Dice={dice:.4f}  Jaccard={jc:.4f}  "
179
                      f"HD95={hd95:.2f}  ASD={asd:.2f}")
180

181
        all_metrics.append(metric_i)   # list of 3 tuples
182

183
    # ── Aggregate ─────────────────────────────────────────────────────────────
184
    all_metrics = np.array(all_metrics)   # (N, 3, 4)
185
    mean_metrics = all_metrics.mean(axis=0)   # (3, 4)
186

187
    class_names = ['RV (class 1)', 'Myo (class 2)', 'LV (class 3)']
188
    print("\n" + "=" * 70)
189
    print(f"{'Class':<20} {'Dice':>8} {'Jaccard':>10} {'HD95':>8} {'ASD':>8}")
190
    print("-" * 70)
191
    for name, m in zip(class_names, mean_metrics):
192
        print(f"{name:<20} {m[0]:>8.4f} {m[1]:>10.4f} {m[2]:>8.2f} {m[3]:>8.2f}")
193
    print("=" * 70)
194
    mean_dice    = mean_metrics[:, 0].mean()
195
    mean_jaccard = mean_metrics[:, 1].mean()
196
    mean_hd95    = mean_metrics[:, 2].mean()
197
    mean_asd     = mean_metrics[:, 3].mean()
198
    print(f"{'Mean':<20} {mean_dice:>8.4f} {mean_jaccard:>10.4f} "
199
          f"{mean_hd95:>8.2f} {mean_asd:>8.2f}")
200

201
    return mean_metrics
202

203

204
if __name__ == '__main__':
205
    metric_result = test_calculate_metric()
206
    print("\nFinal mean metrics (Dice / Jaccard / HD95 / ASD):")
207
    print(metric_result)
208

209
# ── Example usage ─────────────────────────────────────────────────────────────
210
# python ACDC_CPAM_test.py --model unet --labelnum 3 --stage_name self_train --gpu 0
211
# python ACDC_CPAM_test.py --model unet --labelnum 7 --stage_name self_train --gpu 0 --detail 0