Merge pull request #372 from IMSY-DKFZ/T_371_seg_many_wvs

T 371 segmentation loader with multiple wavelengths
IMSY-DKFZ · Aug 15, 2024 · c93b0f6 · c93b0f6
2 parents 85463ef + 88f5627
commit c93b0f6
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 12 deletions.
diff --git a/simpa/core/simulation_modules/volume_creation_module/segmentation_based_adapter.py b/simpa/core/simulation_modules/volume_creation_module/segmentation_based_adapter.py
@@ -4,8 +4,6 @@
 
 from simpa.core.simulation_modules.volume_creation_module import VolumeCreationAdapterBase
 from simpa.utils import Tags
-from simpa.utils.constants import property_tags
-from simpa.io_handling import save_hdf5
 import numpy as np
 import torch
 
@@ -44,21 +42,19 @@ def create_simulation_volume(self) -> dict:
 
         for seg_class in segmentation_classes:
             class_properties = class_mapping[seg_class].get_properties_for_wavelength(self.global_settings, wavelength)
-            for prop_tag in property_tags:
-                if isinstance(class_properties[prop_tag], (int, float)) or class_properties[prop_tag] == None:  # scalar
-                    assigned_prop = class_properties[prop_tag]
+            for volume_key in volumes.keys():
+                if isinstance(class_properties[volume_key], (int, float)) or class_properties[volume_key] == None:  # scalar
+                    assigned_prop = class_properties[volume_key]
                     if assigned_prop is None:
                         assigned_prop = torch.nan
-                    volumes[prop_tag][segmentation_volume == seg_class] = assigned_prop
-                elif len(torch.Tensor.size(class_properties[prop_tag])) == 3:  # 3D map
-                    assigned_prop = class_properties[prop_tag][torch.tensor(segmentation_volume == seg_class)]
+                    volumes[volume_key][segmentation_volume == seg_class] = assigned_prop
+                elif len(torch.Tensor.size(class_properties[volume_key])) == 3:  # 3D map
+                    assigned_prop = class_properties[volume_key][torch.tensor(segmentation_volume == seg_class)]
                     assigned_prop[assigned_prop is None] = torch.nan
-                    volumes[prop_tag][torch.tensor(segmentation_volume == seg_class)] = assigned_prop
+                    volumes[volume_key][torch.tensor(segmentation_volume == seg_class)] = assigned_prop
                 else:
                     raise AssertionError("Properties need to either be a scalar or a 3D map.")
 
-        save_hdf5(self.global_settings, self.global_settings[Tags.SIMPA_OUTPUT_PATH], "/settings/")
-
         # convert volumes back to CPU
         for key in volumes.keys():
             volumes[key] = volumes[key].numpy().astype(np.float64, copy=False)

diff --git a/simpa_examples/segmentation_loader.py b/simpa_examples/segmentation_loader.py
@@ -66,7 +66,7 @@ def segmentation_class_mapping():
     settings[Tags.SIMULATION_PATH] = path_manager.get_hdf5_file_save_path()
     settings[Tags.VOLUME_NAME] = "SegmentationTest"
     settings[Tags.RANDOM_SEED] = 1234
-    settings[Tags.WAVELENGTHS] = [700]
+    settings[Tags.WAVELENGTHS] = [700, 800]
     settings[Tags.SPACING_MM] = spacing
     settings[Tags.DIM_VOLUME_X_MM] = segmentation_volume_mask.shape[0] * spacing
     settings[Tags.DIM_VOLUME_Y_MM] = segmentation_volume_mask.shape[1] * spacing