Series selection v1 (#119)

* series selection and update output paths

* save dicom metadata

* Autoformat code

---------

Co-authored-by: louisblankemeier <[email protected]>

bin/C2C

@@ -24,7 +24,7 @@ from comp2comp.muscle_adipose_tissue import (
 )
 from comp2comp.spine import spine
 from comp2comp.utils import orientation
-from comp2comp.utils.process import process_2d, process_3d
+from comp2comp.utils.process import process_3d
 
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 

comp2comp/io/io.py

@@ -8,6 +8,7 @@
 
 # import dicom2nifti
 import dosma as dm
+import pydicom
 import SimpleITK as sitk
 
 from comp2comp.inference_class_base import InferenceClass
@@ -90,14 +91,15 @@ def __call__(self, inference_pipeline):
 
         # if self.input_path is a folder
         if self.input_path.is_dir():
-            dicom_series_to_nifti(
+            ds = dicom_series_to_nifti(
                 self.input_path,
                 output_file=os.path.join(
                     segmentations_output_dir, "converted_dcm.nii.gz"
                 ),
                 reorient_nifti=False,
             )
             inference_pipeline.dicom_series_path = str(self.input_path)
+            inference_pipeline.dicom_ds = ds
         elif str(self.input_path).endswith(".nii"):
             shutil.copy(
                 self.input_path,
@@ -115,6 +117,14 @@ def __call__(self, inference_pipeline):
 def dicom_series_to_nifti(input_path, output_file, reorient_nifti):
     reader = sitk.ImageSeriesReader()
     dicom_names = reader.GetGDCMSeriesFileNames(str(input_path))
+    print("Reading dicoms...")
+    ds = pydicom.filereader.dcmread(dicom_names[0])
+    image_type_list = list(ds.ImageType)
+    if any("gsi" in s.lower() for s in image_type_list):
+        raise ValueError("GSI Image Type detected")
     reader.SetFileNames(dicom_names)
     image = reader.Execute()
+    if image.GetDirection() != (1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0):
+        raise ValueError("Image orientation is not axial")
     sitk.WriteImage(image, output_file)
+    return ds

comp2comp/io/io_utils.py

@@ -1,8 +1,11 @@
 """
 @author: louisblankemeier
 """
+import csv
 import os
 
+import pydicom
+
 
 def find_dicom_files(input_path):
     dicom_series = []
@@ -47,7 +50,7 @@ def get_dicom_or_nifti_paths_and_num(path):
     if path.endswith(".nii") or path.endswith(".nii.gz"):
         return [(path, 1)]
     dicom_nifti_paths = []
-    for root, _, files in os.walk(path):
+    for root, dirs, files in os.walk(path):
         if len(files) > 0:
             if all(file.endswith(".dcm") or file.endswith(".dicom") for file in files):
                 dicom_nifti_paths.append((root, len(files)))
@@ -58,3 +61,17 @@ def get_dicom_or_nifti_paths_and_num(path):
                         dicom_nifti_paths.append((os.path.join(root, file), num_slices))
 
     return dicom_nifti_paths
+
+
+def write_dicom_metadata_to_csv(ds, csv_filename):
+    with open(csv_filename, "w", newline="") as csvfile:
+        csvwriter = csv.writer(csvfile)
+        csvwriter.writerow(["Tag", "Keyword", "Value"])
+
+        for element in ds:
+            tag = element.tag
+            keyword = pydicom.datadict.keyword_for_tag(tag)
+            if keyword == "PixelData":
+                continue
+            value = str(element.value)
+            csvwriter.writerow([tag, keyword, value])

comp2comp/muscle_adipose_tissue/muscle_adipose_tissue.py

@@ -414,7 +414,7 @@ def __call__(self, inference_pipeline, results):
 
     def save_results(self, results):
         """Save results to a CSV file."""
-        categories = self.model_type.categories
+        self.model_type.categories
         df = pd.DataFrame(
             columns=[
                 "Level",

comp2comp/spine/spine.py

@@ -22,6 +22,7 @@
 )
 
 from comp2comp.inference_class_base import InferenceClass
+from comp2comp.io import io_utils
 from comp2comp.models.models import Models
 from comp2comp.spine import spine_utils
 from comp2comp.visualization.dicom import to_dicom
@@ -250,13 +251,14 @@ def __call__(self, inference_pipeline):
         # Compute ROIs
         inference_pipeline.spine_model_type = self.spine_model_type
 
-        (spine_hus, rois, centroids_3d) = spine_utils.compute_rois(
+        (spine_hus, rois, segmentation_hus, centroids_3d) = spine_utils.compute_rois(
             inference_pipeline.segmentation,
             inference_pipeline.medical_volume,
             self.spine_model_type,
         )
 
         inference_pipeline.spine_hus = spine_hus
+        inference_pipeline.segmentation_hus = segmentation_hus
         inference_pipeline.rois = rois
         inference_pipeline.centroids_3d = centroids_3d
 
@@ -272,20 +274,30 @@ def __init__(self):
     def __call__(self, inference_pipeline):
         """Save metrics to a CSV file."""
         self.spine_hus = inference_pipeline.spine_hus
+        self.seg_hus = inference_pipeline.segmentation_hus
         self.output_dir = inference_pipeline.output_dir
         self.csv_output_dir = os.path.join(self.output_dir, "metrics")
         if not os.path.exists(self.csv_output_dir):
             os.makedirs(self.csv_output_dir, exist_ok=True)
         self.save_results()
+        if hasattr(inference_pipeline, "dicom_ds"):
+            if not os.path.exists(os.path.join(self.output_dir, "dicom_metadata.csv")):
+                io_utils.write_dicom_metadata_to_csv(
+                    inference_pipeline.dicom_ds,
+                    os.path.join(self.output_dir, "dicom_metadata.csv"),
+                )
+
         return {}
 
     def save_results(self):
         """Save results to a CSV file."""
-        df = pd.DataFrame(columns=["Level", "ROI HU"])
+        df = pd.DataFrame(columns=["Level", "ROI HU", "Seg HU"])
         for i, level in enumerate(self.spine_hus):
             hu = self.spine_hus[level]
-            row = [level, hu]
+            seg_hu = self.seg_hus[level]
+            row = [level, hu, seg_hu]
             df.loc[i] = row
+        df = df.iloc[::-1]
         df.to_csv(os.path.join(self.csv_output_dir, "spine_metrics.csv"), index=False)
 
 
@@ -327,6 +339,7 @@ def __call__(self, inference_pipeline):
             list(inference_pipeline.centroids_3d.values()),
             output_path,
             spine_hus=inference_pipeline.spine_hus,
+            seg_hus=inference_pipeline.segmentation_hus,
             model_type=spine_model_type,
             pixel_spacing=inference_pipeline.pixel_spacing_list,
             format=self.format,

comp2comp/spine/spine_utils.py

@@ -5,19 +5,17 @@
 import logging
 import math
 import os
-from glob import glob
 from typing import Dict, List
 
 import cv2
 import nibabel as nib
 import numpy as np
-from pydicom.filereader import dcmread
 from scipy.ndimage import zoom
 
 from comp2comp.spine import spine_visualization
 
 
-def find_spine_dicoms(centroids: Dict): #, path: str, levels):
+def find_spine_dicoms(centroids: Dict):  # , path: str, levels):
     """Find the dicom files corresponding to the spine T12 - L5 levels."""
 
     vertical_positions = []
@@ -268,15 +266,19 @@ def compute_rois(seg, img, spine_model_type):
     rois = {}
     spine_hus = {}
     centroids_3d = {}
+    segmentation_hus = {}
     for i, level in enumerate(slices):
         slice = slices[level]
         center_of_mass = compute_center_of_mass(slice)
         centroid = np.array([centroids[level], center_of_mass[1], center_of_mass[0]])
         roi = roi_from_mask(img, centroid)
-        spine_hus[level] = mean_img_mask(img.get_fdata(), roi, i)
+        image_numpy = img.get_fdata()
+        spine_hus[level] = mean_img_mask(image_numpy, roi, i)
         rois[level] = roi
+        mask = (seg_np == spine_model_type.categories[level]).astype(int)
+        segmentation_hus[level] = mean_img_mask(image_numpy, mask, i)
         centroids_3d[level] = centroid
-    return (spine_hus, rois, centroids_3d)
+    return (spine_hus, rois, segmentation_hus, centroids_3d)
 
 
 def keep_two_largest_connected_components(mask: Dict):
@@ -358,6 +360,7 @@ def visualize_coronal_sagittal_spine(
     centroids_3d: np.ndarray,
     output_dir: str,
     spine_hus=None,
+    seg_hus=None,
     model_type=None,
     pixel_spacing=None,
     format="png",
@@ -453,6 +456,7 @@ def visualize_coronal_sagittal_spine(
         output_dir,
         sagittal_name,
         spine_hus=spine_hus,
+        seg_hus=seg_hus,
         model_type=model_type,
         pixel_spacing=pixel_spacing,
     )
@@ -462,6 +466,7 @@ def visualize_coronal_sagittal_spine(
         output_dir,
         coronal_name,
         spine_hus=spine_hus,
+        seg_hus=seg_hus,
         model_type=model_type,
         pixel_spacing=pixel_spacing,
     )

comp2comp/spine/spine_visualization.py

@@ -18,6 +18,7 @@ def spine_binary_segmentation_overlay(
     file_name: str,
     figure_text_key=None,
     spine_hus=None,
+    seg_hus=None,
     spine=True,
     model_type=None,
     pixel_spacing=None,
@@ -66,7 +67,7 @@ def spine_binary_segmentation_overlay(
     _ROI_COLOR = np.array([1.000, 0.340, 0.200])
 
     _SPINE_TEXT_OFFSET_FROM_TOP = 10.0
-    _SPINE_TEXT_OFFSET_FROM_RIGHT = 63.0
+    _SPINE_TEXT_OFFSET_FROM_RIGHT = 40.0
     _SPINE_TEXT_VERTICAL_SPACING = 14.0
 
     img_in = np.clip(img_in, -300, 1800)
@@ -108,18 +109,60 @@ def spine_binary_segmentation_overlay(
             area_threshold=0,
         )
 
+    vis.draw_text(
+        text="ROI",
+        position=(
+            mask.shape[1] - _SPINE_TEXT_OFFSET_FROM_RIGHT - 35,
+            _SPINE_TEXT_OFFSET_FROM_TOP,
+        ),
+        color=[1, 1, 1],
+        font_size=9,
+        horizontal_alignment="center",
+    )
+
+    vis.draw_text(
+        text="Seg",
+        position=(
+            mask.shape[1] - _SPINE_TEXT_OFFSET_FROM_RIGHT,
+            _SPINE_TEXT_OFFSET_FROM_TOP,
+        ),
+        color=[1, 1, 1],
+        font_size=9,
+        horizontal_alignment="center",
+    )
+
     # draw text and lines
     for i, level in enumerate(levels):
         vis.draw_text(
-            text=f"{level}: {round(float(spine_hus[level]))}",
+            text=f"{level}:",
             position=(
-                mask.shape[1] - _SPINE_TEXT_OFFSET_FROM_RIGHT,
-                _SPINE_TEXT_VERTICAL_SPACING * i + _SPINE_TEXT_OFFSET_FROM_TOP,
+                mask.shape[1] - _SPINE_TEXT_OFFSET_FROM_RIGHT - 80,
+                _SPINE_TEXT_VERTICAL_SPACING * (i + 1) + _SPINE_TEXT_OFFSET_FROM_TOP,
             ),
             color=_COLORS[label_map[level]],
             font_size=9,
             horizontal_alignment="left",
         )
+        vis.draw_text(
+            text=f"{round(float(spine_hus[level]))}",
+            position=(
+                mask.shape[1] - _SPINE_TEXT_OFFSET_FROM_RIGHT - 35,
+                _SPINE_TEXT_VERTICAL_SPACING * (i + 1) + _SPINE_TEXT_OFFSET_FROM_TOP,
+            ),
+            color=_COLORS[label_map[level]],
+            font_size=9,
+            horizontal_alignment="center",
+        )
+        vis.draw_text(
+            text=f"{round(float(seg_hus[level]))}",
+            position=(
+                mask.shape[1] - _SPINE_TEXT_OFFSET_FROM_RIGHT,
+                _SPINE_TEXT_VERTICAL_SPACING * (i + 1) + _SPINE_TEXT_OFFSET_FROM_TOP,
+            ),
+            color=_COLORS[label_map[level]],
+            font_size=9,
+            horizontal_alignment="center",
+        )
 
         """
         vis.draw_line(

comp2comp/utils/process.py

@@ -3,14 +3,14 @@
 """
 
 import os
+import shutil
 import sys
 import traceback
 from datetime import datetime
 from pathlib import Path
 from time import time
-import shutil
 
-from comp2comp.io.io_utils import get_dicom_or_nifti_paths_and_num
+from comp2comp.io import io_utils
 
 
 def process_2d(args, pipeline_builder):
@@ -49,7 +49,7 @@ def process_3d(args, pipeline_builder):
         date_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
         output_path = os.path.join(output_path, date_time)
 
-    for path, num in get_dicom_or_nifti_paths_and_num(args.input_path):
+    for path, num in io_utils.get_dicom_or_nifti_paths_and_num(args.input_path):
         try:
             st = time()
 
@@ -87,7 +87,10 @@ def process_3d(args, pipeline_builder):
                 output_dir = Path(
                     os.path.join(
                         output_path,
-                        Path(os.path.basename(os.path.normpath(path))),
+                        Path(os.path.basename(os.path.normpath(args.input_path))),
+                        os.path.relpath(
+                            os.path.normpath(path), os.path.normpath(args.input_path)
+                        ),
                     )
                 )
 
@@ -98,7 +101,6 @@ def process_3d(args, pipeline_builder):
 
             pipeline(output_dir=output_dir, model_dir=model_dir)
 
-
             if not args.save_segmentations:
                 # remove the segmentations folder
                 segmentations_dir = os.path.join(output_dir, "segmentations")
@@ -110,4 +112,6 @@ def process_3d(args, pipeline_builder):
         except Exception:
             print(f"ERROR PROCESSING {path}\n")
             traceback.print_exc()
+            if os.path.exists(output_dir):
+                shutil.rmtree(output_dir)
             continue