from monai.data import write_nifti
from monai.utils import set_determinism
from monai.transforms import (
    AddChanneld,
    LoadImaged,
    SpatialPadd,
    RandCropByPosNegLabeld
)

import os
import glob
**set_determinism(0)** ##not effect
data_dir = "G:\Debug\\fineseg\Task08_HepaticVessel"

train_images = sorted(
    glob.glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
train_labels = sorted(
    glob.glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
data_dicts = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images, train_labels)
]
train_data_dicts, val_data_dicts = data_dicts[:-9], data_dicts[-9:]


predict_root="G:\Debug\\fineseg\predict"


loader = LoadImaged(keys=("image", "label"))
data_dict = loader(train_data_dicts[0])
print(f"input:, {train_data_dicts[0]}")
print(f"image shape: {data_dict['image'].shape}")
print(f"label shape: {data_dict['label'].shape}")
print(f"image pixdim:\n{data_dict['image_meta_dict']['pixdim']}")

add_channel = AddChanneld(keys=["image", "label"])
data_dict = add_channel(data_dict)

spatial_pad=SpatialPadd(spatial_size = [512,512,160],method = 'end'
                        ,keys = ["image", "label"])
data_dict = spatial_pad(data_dict)
crop = RandCropByPosNegLabeld(num_samples = 200,keys = ["image", "label"],pos = 1,neg = 0 ,spatial_size = [160,160,160],label_key = "label")
data_dict = crop(data_dict)

for i in range(0,10):
     write_nifti(data =data_dict[i]['image'].squeeze(0), file_name = os.path.join(predict_root, str(i)+".nii.gz"))
     print(data_dict[i]['image'].sum()) # debug image data
     write_nifti(data = data_dict[i]['label'].squeeze(0), file_name = os.path.join(predict_root, str(i) + "label.nii.gz"))