update dataset and transform

taigw · taigw · commit 4590b3045d92 · 2025-08-05T10:42:56.000+08:00
fix code for loading h5 images
update transforms
diff --git a/pymic/io/nifty_dataset.py b/pymic/io/nifty_dataset.py
@@ -41,8 +41,9 @@ class NiftyDataset(Dataset):
     :param transform:  (list) List of transforms to be applied on a sample.
         The built-in transforms can listed in :mod:`pymic.transform.trans_dict`.
     """
-    def __init__(self, root_dir, csv_file, modal_num = 1, image_dim = 3, allow_missing_modal = False,
-            with_label = True, transform=None, task = TaskType.SEGMENTATION):
+    def __init__(self, root_dir, csv_file, modal_num = 1, image_dim = 3, 
+            allow_missing_modal = False, label_key = "label",
+            transform=None, task = TaskType.SEGMENTATION):
         self.root_dir   = root_dir
         if(csv_file is not None):
             self.csv_items  = pd.read_csv(csv_file)
@@ -56,10 +57,11 @@ def __init__(self, root_dir, csv_file, modal_num = 1, image_dim = 3, allow_missi
         self.modal_num  = modal_num
         self.image_dim  = image_dim
         self.allow_emtpy= allow_missing_modal
-        self.with_label = with_label
+        self.label_key  = label_key
         self.transform  = transform
         self.task       = task
         self.h5files    = False
+        self.with_label = True
         assert self.task in  [TaskType.SEGMENTATION, TaskType.RECONSTRUCTION]
        
         # check if the files are h5 images, and if the labels are provided.
@@ -69,11 +71,11 @@ def __init__(self, root_dir, csv_file, modal_num = 1, image_dim = 3, allow_missi
             self.h5files = True
             temp_full_name = "{0:}/{1:}".format(self.root_dir, temp_name)
             h5f = h5py.File(temp_full_name, 'r')
-            if('label' not in h5f):
+            if(self.label_key not in h5f):
                 self.with_label = False
         else:
             csv_keys = list(self.csv_items.keys())
-            if('label' not in csv_keys):
+            if(self.label_key not in csv_keys):
                 self.with_label = False
             
             self.image_weight_idx = None
@@ -84,7 +86,7 @@ def __init__(self, root_dir, csv_file, modal_num = 1, image_dim = 3, allow_missi
                 self.pixel_weight_idx = csv_keys.index('pixel_weight')
         if(not self.with_label):
             logging.warning("`label` section is not found in the csv file {0:}".format(
-                csv_file) + "or the corresponding h5 file." + 
+                csv_file) + " or the corresponding h5 file." + 
                 "\n -- This is only allowed for self-supervised learning" + 
                 "\n -- when `SelfSuperviseLabel` is used in the transform, or when" + 
                 "\n -- loading the unlabeled data for preprocessing.")
@@ -94,7 +96,7 @@ def __len__(self):
 
     def __getlabel__(self, idx):
         csv_keys = list(self.csv_items.keys())        
-        label_idx  = csv_keys.index('label')
+        label_idx  = csv_keys.index(self.label_key)
         label_name = self.csv_items.iloc[idx, label_idx]
         label_name_full = "{0:}/{1:}".format(self.root_dir, label_name)
         label = load_image_as_nd_array(label_name_full)['data_array']
@@ -139,8 +141,8 @@ def __getitem__(self, idx):
             img = check_and_expand_dim(h5f['image'][:], self.image_dim)
             sample = {'image':img}
             if(self.with_label):
-                lab = check_and_expand_dim(h5f['label'][:], self.image_dim)
-                sample['label'] = lab
+                lab = check_and_expand_dim(h5f[self.label_key][:], self.image_dim)
+                sample['label'] = np.asarray(lab, np.float32)
             sample['names'] = [sample_name]
         else:            
             for i in range (self.modal_num):
diff --git a/pymic/net_run/agent_seg.py b/pymic/net_run/agent_seg.py
@@ -72,19 +72,22 @@ def get_stage_dataset_from_config(self, stage):
         allow_miss = self.config['dataset'].get('allow_missing_modal', False)
         stage_dir  = self.config['dataset'].get('train_dir', None)
         stage_dim  = self.config['dataset'].get('train_dim', 3)
+        stage_lab_key = self.config['dataset'].get('train_label_key', 'label')
         if(stage == 'valid'): # and "valid_dir" in self.config['dataset']):
             stage_dir = self.config['dataset'].get('valid_dir', stage_dir)
             stage_dim = self.config['dataset'].get('valid_dim', stage_dim)
+            stage_lab_key = self.config['dataset'].get('valid_label_key', 'label')
         if(stage == 'test'): # and "test_dir" in self.config['dataset']):
             stage_dir = self.config['dataset'].get('test_dir', stage_dir)
             stage_dim = self.config['dataset'].get('test_dim', stage_dim)
+            stage_lab_key = self.config['dataset'].get('test_label_key', 'label')
         logging.info("Creating dataset for {0:}".format(stage))
         dataset  = NiftyDataset(root_dir  = stage_dir,
                                 csv_file  = csv_file,
                                 modal_num = modal_num,
                                 image_dim = stage_dim,
                                 allow_missing_modal = allow_miss,
-                                with_label= with_label,
+                                label_key = stage_lab_key,
                                 transform = data_transform, 
                                 task = self.task_type)
         return dataset
diff --git a/pymic/net_run/weak_sup/wsl_dmsps.py b/pymic/net_run/weak_sup/wsl_dmsps.py
@@ -5,11 +5,13 @@
 import random
 import time
 import torch
+from PIL import Image
 from pymic.loss.seg.util import get_soft_label
 from pymic.loss.seg.util import reshape_prediction_and_ground_truth
 from pymic.loss.seg.util import get_classwise_dice
 from pymic.loss.seg.dice import DiceLoss
 from pymic.loss.seg.ce   import CrossEntropyLoss
+# from torch.nn.modules.loss import CrossEntropyLoss as TorchCELoss
 from pymic.net_run.weak_sup import WSLSegAgent
 from pymic.util.ramps import get_rampup_ratio
 
@@ -33,11 +35,11 @@ class WSLDMSPS(WSLSegAgent):
     """
     def __init__(self, config, stage = 'train'):
         net_type = config['network']['net_type']
-        if net_type not in ['UNet2D_DualBranch', 'UNet3D_DualBranch']:
-            raise ValueError("""For WSL_DMPLS, a dual branch network is expected. \
-                It only supports UNet2D_DualBranch and UNet3D_DualBranch currently.""")
+        # if net_type not in ['UNet2D_DualBranch', 'UNet3D_DualBranch']:
+        #     raise ValueError("""For WSL_DMPLS, a dual branch network is expected. \
+        #         It only supports UNet2D_DualBranch and UNet3D_DualBranch currently.""")
         super(WSLDMSPS, self).__init__(config, stage)
-
+  
     def training(self):
         class_num   = self.config['network']['class_num']
         iter_valid  = self.config['training']['iter_valid']
@@ -49,10 +51,12 @@ def training(self):
         if (pseudo_loss_type not in ('dice_loss', 'ce_loss')):
             raise ValueError("""For pseudo supervision loss, only dice_loss and ce_loss \
                 are supported.""")
+        pseudo_loss_func = CrossEntropyLoss() if pseudo_loss_type == 'ce_loss' else DiceLoss()
         train_loss, train_loss_sup, train_loss_reg = 0, 0, 0
         train_dice_list = []
         data_time, gpu_time, loss_time, back_time = 0, 0, 0, 0
         self.net.train()
+        # ce_loss  = CrossEntropyLoss()
         for it in range(iter_valid):
             t0 = time.time()
             try:
@@ -66,7 +70,6 @@ def training(self):
             y      = self.convert_tensor_type(data['label_prob'])  
                          
             inputs, y = inputs.to(self.device), y.to(self.device)
-            
             # zero the parameter gradients
             self.optimizer.zero_grad()
                 
@@ -78,23 +81,26 @@ def training(self):
             loss_sup2 = self.get_loss_value(data, outputs2, y) 
             loss_sup  = 0.5 * (loss_sup1 + loss_sup2)
 
-            # get pseudo label with dynamical mix
+            # torch_ce_loss  = TorchCELoss(ignore_index=class_num)
+            # torch_ce_loss2 = TorchCELoss()
+            # loss_ce1 = torch_ce_loss(outputs1, label[:].long())
+            # loss_ce2 = torch_ce_loss(outputs2, label[:].long())
+            # loss_sup = 0.5 * (loss_ce1 + loss_ce2)
+
+            # get pseudo label with dynamic mixture
             outputs_soft1 = torch.softmax(outputs1, dim=1)
             outputs_soft2 = torch.softmax(outputs2, dim=1)
-            beta = random.random()
-            pseudo_lab = beta*outputs_soft1.detach() + (1.0-beta)*outputs_soft2.detach()
-            # pseudo_lab = torch.argmax(pseudo_lab, dim = 1, keepdim = True)
-            # pseudo_lab = get_soft_label(pseudo_lab, class_num, self.tensor_type)
-            
-            # calculate the pseudo label supervision loss
-            loss_calculator = DiceLoss() if pseudo_loss_type == 'dice_loss' else CrossEntropyLoss()
-            loss_dict1 = {"prediction":outputs1, 'ground_truth':pseudo_lab}
-            loss_dict2 = {"prediction":outputs2, 'ground_truth':pseudo_lab}
-            loss_reg   = 0.5 * (loss_calculator(loss_dict1) + loss_calculator(loss_dict2))
+            alpha = random.random() 
+            soft_pseudo_label = alpha * outputs_soft1.detach() + (1.0-alpha) * outputs_soft2.detach()
+            # loss_reg = 0.5*(torch_ce_loss2(outputs_soft1, soft_pseudo_label) +torch_ce_loss2(outputs_soft2, soft_pseudo_label) )
             
+            loss_dict1 = {"prediction":outputs_soft1, 'ground_truth':soft_pseudo_label}
+            loss_dict2 = {"prediction":outputs_soft2, 'ground_truth':soft_pseudo_label}
+            loss_reg   = 0.5 * (pseudo_loss_func(loss_dict1) + pseudo_loss_func(loss_dict2))
+
             rampup_ratio = get_rampup_ratio(self.glob_it, rampup_start, rampup_end, "sigmoid")
-            regular_w = wsl_cfg.get('regularize_w', 0.1) * rampup_ratio
-            loss = loss_sup + regular_w*loss_reg
+            regular_w    = wsl_cfg.get('regularize_w', 8.0) * rampup_ratio
+            loss         = loss_sup + regular_w*loss_reg
             t3 = time.time()
             loss.backward()
             t4 = time.time()
@@ -127,5 +133,4 @@ def training(self):
             'avg_fg_dice':train_avg_dice,     'class_dice': train_cls_dice,
             'data_time': data_time, 'forward_time':gpu_time, 
             'loss_time':loss_time, 'backward_time':back_time}
-        return train_scalers
-        
+        return train_scalers
diff --git a/pymic/transform/rotate.py b/pymic/transform/rotate.py
@@ -19,19 +19,13 @@ class RandomRotate(AbstractTransform):
 
     :param `RandomRotate_angle_range_d`: (list/tuple or None) 
         Rotation angle (degree) range along depth axis (x-y plane), e.g., (-90, 90).
-        The length of the list/tuple can be larger than 2, when `RandomRotate_discrete_mode` is True.
         If None, no rotation along this axis. 
     :param `RandomRotate_angle_range_h`: (list/tuple or None) 
         Rotation angle (degree) range along height axis (x-z plane), e.g., (-90, 90).
-        The length of the list/tuple can be larger than 2, when `RandomRotate_discrete_mode` is True.
         If None, no rotation along this axis. Only used for 3D images. 
     :param `RandomRotate_angle_range_w`: (list/tuple or None) 
         Rotation angle (degree) range along width axis (y-z plane), e.g., (-90, 90).
-        The length of the list/tuple can be larger than 2, when `RandomRotate_discrete_mode` is True.
         If None, no rotation along this axis. Only used for 3D images. 
-    :param `RandomRotate_discrete_mode`: (optional, bool) Whether the rotate angles
-        are discrete values in rangle range. For example, if you only want to rotate 
-        the images with a fixed set of angles like (90, 180, 270), then set discrete_mode mode as True.
     :param `RandomRotate_probability`: (optional, float) 
         The probability of applying RandomRotate. Default is 0.5.
     :param `RandomRotate_inverse`: (optional, bool) 
@@ -42,11 +36,8 @@ def __init__(self, params):
         self.angle_range_d  = params['RandomRotate_angle_range_d'.lower()]
         self.angle_range_h  = params.get('RandomRotate_angle_range_h'.lower(), None)
         self.angle_range_w  = params.get('RandomRotate_angle_range_w'.lower(), None)
-        self.discrete_mode  = params.get('RandomRotate_discrete_mode'.lower(), False)
         self.prob = params.get('RandomRotate_probability'.lower(), 0.5)
         self.inverse = params.get('RandomRotate_inverse'.lower(), True)
-        if(len(self.angle_range_d) > 2):
-            assert(self.discrete_mode)
 
     def __apply_transformation(self, image, transform_param_list, order = 1):
         """
@@ -61,38 +52,21 @@ def __apply_transformation(self, image, transform_param_list, order = 1):
         return image
 
     def __call__(self, sample):
-        # if(random.random() > self.prob):
-        #     sample['RandomRotate_triggered'] = False
-        #     return sample
-        # else:
-        #     sample['RandomRotate_triggered'] = True
         image = sample['image']
         input_shape = image.shape
         input_dim = len(input_shape) - 1
         
         transform_param_list = []
         if(self.angle_range_d is not None):
-            if(self.discrete_mode):
-                idx = random.randint(0, len(self.angle_range_d) - 1)
-                angle_d = self.angle_range_d[idx]
-            else:
-                angle_d = np.random.uniform(self.angle_range_d[0], self.angle_range_d[1])
+            angle_d = np.random.uniform(self.angle_range_d[0], self.angle_range_d[1])
             transform_param_list.append([angle_d, (-1, -2)])
         if(input_dim == 3):
             if(self.angle_range_h is not None):
-                if(self.discrete_mode):
-                    idx = random.randint(0, len(self.angle_range_h) - 1)
-                    angle_h = self.angle_range_h[idx]
-                else:
-                    angle_h = np.random.uniform(self.angle_range_h[0], self.angle_range_h[1])
-                    transform_param_list.append([angle_h, (-1, -3)])
+                angle_h = np.random.uniform(self.angle_range_h[0], self.angle_range_h[1])
+                transform_param_list.append([angle_h, (-1, -3)])
             if(self.angle_range_w is not None):
-                if(self.discrete_mode):
-                    idx = random.randint(0, len(self.angle_range_w) - 1)
-                    angle_w = self.angle_range_w[idx]
-                else:
-                    angle_w = np.random.uniform(self.angle_range_w[0], self.angle_range_w[1])
-                    transform_param_list.append([angle_w, (-2, -3)])
+                angle_w = np.random.uniform(self.angle_range_w[0], self.angle_range_w[1])
+                transform_param_list.append([angle_w, (-2, -3)])
         assert(len(transform_param_list) > 0)
         # select a random transform from the possible list rather than 
         # use a combination for higher efficiency
@@ -123,4 +97,49 @@ def  inverse_transform_for_prediction(self, sample):
             transform_param_list[i][0] = - transform_param_list[i][0]
         sample['predict'] = self.__apply_transformation(sample['predict'] , 
                                 transform_param_list, 1)
+        return sample
+
+class RandomRot90(AbstractTransform):
+    """
+    Random rotate an image in x-y plane with angles in [90, 180, 270].
+
+    The arguments should be written in the `params` dictionary, and it has the
+    following fields:
+
+    :param `RandomRot90_probability`: (optional, float) 
+        The probability of applying RandomRot90. Default is 0.75.
+    :param `RandomRot90_inverse`: (optional, bool) 
+        Is inverse transform needed for inference. Default is `True`.
+    """
+    def __init__(self, params): 
+        super(RandomRot90, self).__init__(params)
+        self.prob = params.get('RandomRot90_probability'.lower(), 0.75)
+        self.inverse = params.get('RandomRot90_inverse'.lower(), True)
+
+    def __call__(self, sample):
+        if(random.random() > self.prob):
+            sample['RandomRot90_triggered'] = False
+            sample['RandomRot90_Param'] = 0
+            return sample
+        else:
+            sample['RandomRot90_triggered'] = True
+        image = sample['image']      
+        rote_k = random.randint(1, 3)
+        sample['RandomRot90_Param'] = rote_k
+        image_t = np.rot90(image, rote_k, (-2, -1))
+        sample['image'] = image_t
+        if('label' in sample and \
+        self.task in [TaskType.SEGMENTATION, TaskType.RECONSTRUCTION]):
+            sample['label'] = np.rot90(sample['label'], rote_k, (-2, -1))
+        if('pixel_weight' in sample and \
+        self.task in [TaskType.SEGMENTATION, TaskType.RECONSTRUCTION]):
+            sample['pixel_weight'] = np.rot90(sample['pixel_weight'], rote_k, (-2, -1))
+        return sample
+
+    def inverse_transform_for_prediction(self, sample):
+        if(not sample['RandomRot90_triggered']):
+            return sample
+        rote_k = sample['RandomRot90_Param']
+        rote_i = 4 - rote_k
+        sample['predict'] = np.rot90(sample['predict'], rote_i, (-2, -1))
         return sample
diff --git a/pymic/transform/trans_dict.py b/pymic/transform/trans_dict.py
@@ -25,6 +25,7 @@
     'RandomRescale': RandomRescale,
     'RandomFlip': RandomFlip,
     'RandomRotate': RandomRotate,
+    'RandomRot90': RandomRot90,
     'ReduceLabelDim': ReduceLabelDim,
     'Rescale': Rescale,
     'SelfSuperviseLabel': SelfSuperviseLabel,
@@ -43,6 +44,7 @@
 from pymic.transform.normalize import *
 from pymic.transform.crop import *
 from pymic.transform.crop4dino import Crop4Dino
+from pymic.transform.crop4voco import Crop4VoCo
 from pymic.transform.crop4vox2vec import Crop4Vox2Vec
 from pymic.transform.crop4vf import Crop4VolumeFusion, VolumeFusion, VolumeFusionShuffle
 from pymic.transform.label_convert import *  
@@ -57,6 +59,7 @@
     'CropHumanRegion': CropHumanRegion,
     'CenterCrop': CenterCrop,
     'Crop4Dino': Crop4Dino,
+    'Crop4VoCo': Crop4VoCo,
     'Crop4Vox2Vec': Crop4Vox2Vec,
     'Crop4VolumeFusion': Crop4VolumeFusion,
     'GrayscaleToRGB': GrayscaleToRGB,
@@ -83,6 +86,7 @@
     'RandomTranspose': RandomTranspose,
     'RandomFlip': RandomFlip,
     'RandomRotate': RandomRotate,
+    'RandomRot90': RandomRot90,
     'ReduceLabelDim': ReduceLabelDim,
     'Rescale': Rescale,
     'Resample': Resample,
