def save_mask(self, pred):
if self.test_idx == 0:
data_path = get_path(self.args)
self.test_imgs, _ = get_test_fnames(self.args, data_path)
fname = os.path.basename(self.test_imgs[self.test_idx]).replace(
"_x", "")
np.save(os.path.join(self.save_dir, fname), pred, allow_pickle=False)
self.test_idx += 1
def setup(self, stage=None):
imgs = load_data(self.data_path, "*_x.npy")
lbls = load_data(self.data_path, "*_y.npy")
self.test_imgs, self.kwargs["meta"] = get_test_fnames(
self.args, self.data_path, self.kwargs["meta"])
if self.args.exec_mode != "predict" or self.args.benchmark:
train_idx, val_idx = list(self.kfold.split(imgs))[self.args.fold]
self.train_imgs = get_split(imgs, train_idx)
self.train_lbls = get_split(lbls, train_idx)
self.val_imgs = get_split(imgs, val_idx)
self.val_lbls = get_split(lbls, val_idx)
if is_main_process():
ntrain, nval = len(self.train_imgs), len(self.val_imgs)
print(f"Number of examples: Train {ntrain} - Val {nval}")
elif is_main_process():
print(f"Number of test examples: {len(self.test_imgs)}")