Fixed read_dicom_auto

Former-commit-id: a9a297e
bhklab · Jun 24, 2022 · 290fe0a · 290fe0a
1 parent 67f7af1
commit 290fe0a
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 40 deletions.
diff --git a/README.md b/README.md
@@ -52,13 +52,17 @@ These google collab notebooks will introduce the main functionalities of med-ima
 Med-Imagetools takes two step approch to turn messy medical raw dataset to ML ready dataset.  
 1. ***Autopipeline***: Crawls the raw dataset, forms a network and performs graph query, based on the user defined modalities. The relevant DICOMS, get processed and saved as nrrds
     ```
-    python imgtools/autopipeline.py\
-     [INPUT DIRECTORY] \
-     [OUTPUT DIRECTORY] \
-     --modalities [str: CT,RTSTRUCT,PT] \
-     --spacing [Tuple: (int,int,int)]\
-     --n_jobs [int]\
-     --visualize [bool: True/False]\
+    autopipeline\
+      [INPUT DIRECTORY] \
+      [OUTPUT DIRECTORY] \
+      --modalities [str: CT,RTSTRUCT,PT] \
+      --spacing [Tuple: (int,int,int)]\
+      --n_jobs [int]\
+      --visualize [flag]\
+      --nnunet [flag]\
+      --train_size [float]\
+      --random_state [int]\
+      --roi_yaml_path [str]
     ```
 2. ***class Dataset***: This class converts processed nrrds to torchio subjects, which can be easily converted to torch dataset
     ```

diff --git a/imgtools/io/loaders.py b/imgtools/io/loaders.py
@@ -83,40 +83,42 @@ def read_dicom_pet(path,series=None):
 def read_dicom_auto(path, series=None):
     if path is None:
         return None
-    # print(path, "asdlfkjqeroigjodfklsjg")
     dcms = glob.glob(pathlib.Path(path, "*.dcm").as_posix())
-    meta = dcmread(dcms[0])
-    modality = meta.Modality
-    all_modality_metadata = all_modalities_metadata(meta)
-    if modality == 'CT' or modality == 'MR':
-        dicom_series = read_dicom_series(path,series)#, modality=modality)
-        if modality == 'CT':
-            dicom_series.metadata.update(ct_metadata(meta))
-            dicom_series.metadata.update(all_modality_metadata)
-        else:
-            dicom_series.metadata.update(mr_metadata(meta))
-            dicom_series.metadata.update(all_modality_metadata)
-        return dicom_series
-    elif modality == 'PT':
-        pet = read_dicom_pet(path,series)
-        pet.metadata.update(pet_metadata(meta))
-        pet.metadata.update(all_modality_metadata)
-        return pet
-    elif modality == 'RTSTRUCT':
-        rtstruct = read_dicom_rtstruct(dcms[0])
-        rtstruct.metadata.update(rtstruct_metadata(meta))
-        rtstruct.metadata.update(all_modality_metadata)
-        return rtstruct
-    elif modality == 'RTDOSE':
-        rtdose = read_dicom_rtdose(path)
-        rtdose.metadata.update(all_modality_metadata)
-        return rtdose
-    else:
-        if len(dcms)==1:
-            raise NotImplementedError
+    for dcm in dcms:
+        meta = dcmread(dcm)
+        if meta.SeriesInstanceUID != series:
+            continue
+        modality = meta.Modality
+        all_modality_metadata = all_modalities_metadata(meta)
+        if modality == 'CT' or modality == 'MR':
+            dicom_series = read_dicom_series(path,series)#, modality=modality)
+            if modality == 'CT':
+                dicom_series.metadata.update(ct_metadata(meta))
+                dicom_series.metadata.update(all_modality_metadata)
+            else:
+                dicom_series.metadata.update(mr_metadata(meta))
+                dicom_series.metadata.update(all_modality_metadata)
+            return dicom_series
+        elif modality == 'PT':
+            pet = read_dicom_pet(path,series)
+            pet.metadata.update(pet_metadata(meta))
+            pet.metadata.update(all_modality_metadata)
+            return pet
+        elif modality == 'RTSTRUCT':
+            rtstruct = read_dicom_rtstruct(dcm)
+            rtstruct.metadata.update(rtstruct_metadata(meta))
+            rtstruct.metadata.update(all_modality_metadata)
+            return rtstruct
+        elif modality == 'RTDOSE':
+            rtdose = read_dicom_rtdose(path)
+            rtdose.metadata.update(all_modality_metadata)
+            return rtdose
         else:
-            print("There were no dicoms in this path.")
-            return None
+            if len(dcms)==1:
+                raise NotImplementedError
+            else:
+                print("There were no dicoms in this path.")
+                return None
 
 def read_segmentation(path):
     # TODO read seg.nrrd

diff --git a/imgtools/utils/args.py b/imgtools/utils/args.py
@@ -13,7 +13,7 @@ def parser():
     parser.add_argument("--modalities", type=str, default="CT",
                         help="List of desired modalities. Type as string for ex: RTSTRUCT,CT,RTDOSE")
 
-    parser.add_argument("--visualize", type=bool, default=False,
+    parser.add_argument("--visualize", default=False, action="store_true",
                         help="Whether to visualize the data graph")
 
     parser.add_argument("--spacing", nargs=3, type=float, default=(1., 1., 0.),