156 lines
5.0 KiB
Python
156 lines
5.0 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import argparse
|
|
import os.path as osp
|
|
|
|
import nibabel as nib
|
|
import numpy as np
|
|
from mmengine.utils import mkdir_or_exist
|
|
from PIL import Image
|
|
|
|
|
|
def read_files_from_txt(txt_path):
|
|
with open(txt_path) as f:
|
|
files = f.readlines()
|
|
files = [file.strip() for file in files]
|
|
return files
|
|
|
|
|
|
def read_nii_file(nii_path):
|
|
img = nib.load(nii_path).get_fdata()
|
|
return img
|
|
|
|
|
|
def split_3d_image(img):
|
|
c, _, _ = img.shape
|
|
res = []
|
|
for i in range(c):
|
|
res.append(img[i, :, :])
|
|
return res
|
|
|
|
|
|
def label_mapping(label):
|
|
"""Label mapping from TransUNet paper setting. It only has 9 classes, which
|
|
are 'background', 'aorta', 'gallbladder', 'left_kidney', 'right_kidney',
|
|
'liver', 'pancreas', 'spleen', 'stomach', respectively. Other foreground
|
|
classes in original dataset are all set to background.
|
|
|
|
More details could be found here: https://arxiv.org/abs/2102.04306
|
|
"""
|
|
maped_label = np.zeros_like(label)
|
|
maped_label[label == 8] = 1
|
|
maped_label[label == 4] = 2
|
|
maped_label[label == 3] = 3
|
|
maped_label[label == 2] = 4
|
|
maped_label[label == 6] = 5
|
|
maped_label[label == 11] = 6
|
|
maped_label[label == 1] = 7
|
|
maped_label[label == 7] = 8
|
|
return maped_label
|
|
|
|
|
|
def pares_args():
|
|
parser = argparse.ArgumentParser(
|
|
description='Convert synapse dataset to mmsegmentation format')
|
|
parser.add_argument(
|
|
'--dataset-path', type=str, help='synapse dataset path.')
|
|
parser.add_argument(
|
|
'--save-path',
|
|
default='data/synapse',
|
|
type=str,
|
|
help='save path of the dataset.')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def main():
|
|
args = pares_args()
|
|
dataset_path = args.dataset_path
|
|
save_path = args.save_path
|
|
|
|
if not osp.exists(dataset_path):
|
|
raise ValueError('The dataset path does not exist. '
|
|
'Please enter a correct dataset path.')
|
|
if not osp.exists(osp.join(dataset_path, 'img')) \
|
|
or not osp.exists(osp.join(dataset_path, 'label')):
|
|
raise FileNotFoundError('The dataset structure is incorrect. '
|
|
'Please check your dataset.')
|
|
|
|
train_id = read_files_from_txt(osp.join(dataset_path, 'train.txt'))
|
|
train_id = [idx[3:7] for idx in train_id]
|
|
|
|
test_id = read_files_from_txt(osp.join(dataset_path, 'val.txt'))
|
|
test_id = [idx[3:7] for idx in test_id]
|
|
|
|
mkdir_or_exist(osp.join(save_path, 'img_dir/train'))
|
|
mkdir_or_exist(osp.join(save_path, 'img_dir/val'))
|
|
mkdir_or_exist(osp.join(save_path, 'ann_dir/train'))
|
|
mkdir_or_exist(osp.join(save_path, 'ann_dir/val'))
|
|
|
|
# It follows data preparation pipeline from here:
|
|
# https://github.com/Beckschen/TransUNet/tree/main/datasets
|
|
for i, idx in enumerate(train_id):
|
|
img_3d = read_nii_file(
|
|
osp.join(dataset_path, 'img', 'img' + idx + '.nii.gz'))
|
|
label_3d = read_nii_file(
|
|
osp.join(dataset_path, 'label', 'label' + idx + '.nii.gz'))
|
|
|
|
img_3d = np.clip(img_3d, -125, 275)
|
|
img_3d = (img_3d + 125) / 400
|
|
img_3d *= 255
|
|
img_3d = np.transpose(img_3d, [2, 0, 1])
|
|
img_3d = np.flip(img_3d, 2)
|
|
|
|
label_3d = np.transpose(label_3d, [2, 0, 1])
|
|
label_3d = np.flip(label_3d, 2)
|
|
label_3d = label_mapping(label_3d)
|
|
|
|
for c in range(img_3d.shape[0]):
|
|
img = img_3d[c]
|
|
label = label_3d[c]
|
|
|
|
img = Image.fromarray(img).convert('RGB')
|
|
label = Image.fromarray(label).convert('L')
|
|
img.save(
|
|
osp.join(
|
|
save_path, 'img_dir/train', 'case' + idx.zfill(4) +
|
|
'_slice' + str(c).zfill(3) + '.jpg'))
|
|
label.save(
|
|
osp.join(
|
|
save_path, 'ann_dir/train', 'case' + idx.zfill(4) +
|
|
'_slice' + str(c).zfill(3) + '.png'))
|
|
|
|
for i, idx in enumerate(test_id):
|
|
img_3d = read_nii_file(
|
|
osp.join(dataset_path, 'img', 'img' + idx + '.nii.gz'))
|
|
label_3d = read_nii_file(
|
|
osp.join(dataset_path, 'label', 'label' + idx + '.nii.gz'))
|
|
|
|
img_3d = np.clip(img_3d, -125, 275)
|
|
img_3d = (img_3d + 125) / 400
|
|
img_3d *= 255
|
|
img_3d = np.transpose(img_3d, [2, 0, 1])
|
|
img_3d = np.flip(img_3d, 2)
|
|
|
|
label_3d = np.transpose(label_3d, [2, 0, 1])
|
|
label_3d = np.flip(label_3d, 2)
|
|
label_3d = label_mapping(label_3d)
|
|
|
|
for c in range(img_3d.shape[0]):
|
|
img = img_3d[c]
|
|
label = label_3d[c]
|
|
|
|
img = Image.fromarray(img).convert('RGB')
|
|
label = Image.fromarray(label).convert('L')
|
|
img.save(
|
|
osp.join(
|
|
save_path, 'img_dir/val', 'case' + idx.zfill(4) +
|
|
'_slice' + str(c).zfill(3) + '.jpg'))
|
|
label.save(
|
|
osp.join(
|
|
save_path, 'ann_dir/val', 'case' + idx.zfill(4) +
|
|
'_slice' + str(c).zfill(3) + '.png'))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|