restruture so that discard slices is again in review_mu_map module

mu_map/data/datasets.py

@@ -7,33 +7,10 @@ import numpy as np
 from torch.utils.data import Dataset
 
 from mu_map.data.remove_bed import DEFAULT_BED_CONTOURS_FILENAME, load_contours
+from mu_map.data.review_mu_map import discard_slices
 
 
-HEADER_DISC_FIRST = "discard_first"
-HEADER_DISC_LAST = "discard_last"
-
-
-def discard_slices(row, μ_map):
-    """
-    Discard slices based on the flags in the row of th according table.
-    The row is expected to contain the flags 'discard_first' and 'discard_last'.
-
-    :param row: the row of meta configuration file of a dataset
-    :param μ_map: the μ_map
-    :return: the μ_map with according slices removed
-    """
-    _res = μ_map
-
-    if row[HEADER_DISC_FIRST]:
-        _res  = _res[1:]
-
-    if row[HEADER_DISC_LAST]:
-        _res = _res[:-1]
-
-    return _res
-
-
-def align_images(image_1: np.ndarray, image_2: np.ndarray):
+def align_images(image_1: np.ndarray, image_2: np.ndarray) -> np.ndarray:
     """
     Align one image to another on the first axis (z-axis).
     It is assumed that the second image has less slices than the first.
@@ -66,7 +43,7 @@ class MuMapDataset(Dataset):
         csv_file: str = "meta.csv",
         images_dir: str = "images",
         bed_contours_file: Optional[str] = DEFAULT_BED_CONTOURS_FILENAME,
-        discard_μ_map_slices: bool = True,
+        discard_mu_map_slices: bool = True,
     ):
         super().__init__()
 
@@ -74,14 +51,18 @@ class MuMapDataset(Dataset):
         self.dir_images = os.path.join(dataset_dir, images_dir)
         self.csv_file = os.path.join(dataset_dir, csv_file)
 
-        self.bed_contours_file = os.path.join(dataset_dir, bed_contours_file) if bed_contours_file else None 
-        self.bed_contours = load_contours(self.bed_contours_file) if bed_contours_file else None
+        self.bed_contours_file = (
+            os.path.join(dataset_dir, bed_contours_file) if bed_contours_file else None
+        )
+        self.bed_contours = (
+            load_contours(self.bed_contours_file) if bed_contours_file else None
+        )
 
         # read CSV file and from that access DICOM files
         self.table = pd.read_csv(self.csv_file)
         self.table["id"] = self.table["id"].apply(int)
 
-        self.discard_μ_map_slices = discard_μ_map_slices
+        self.discard_mu_map_slices = discard_mu_map_slices
 
     def __getitem__(self, index: int):
         row = self.table.iloc[index]
@@ -92,7 +73,7 @@ class MuMapDataset(Dataset):
         recon = pydicom.dcmread(recon_file).pixel_array
         mu_map = pydicom.dcmread(mu_map_file).pixel_array
 
-        if self.discard_μ_map_slices:
+        if self.discard_mu_map_slices:
             mu_map = discard_slices(row, mu_map)
 
         if self.bed_contours:
@@ -112,7 +93,6 @@ __all__ = [MuMapDataset.__name__]
 
 if __name__ == "__main__":
     dataset = MuMapDataset("data/tmp")
-    print(f"Images: {len(dataset)}")
 
     import cv2 as cv
 

mu_map/data/review_mu_map.py

-import argparse
-
-import cv2 as cv
 import numpy as np
+import pandas as pd
+
+
+HEADER_DISC_FIRST = "discard_first"
+HEADER_DISC_LAST = "discard_last"
+
+
+def discard_slices(row: pd.Series, μ_map: np.ndarray) -> np.ndarray:
+    """
+    Discard slices based on the flags in the row of th according table.
+    The row is expected to contain the flags 'discard_first' and 'discard_last'.
 
-from mu_map.data.datasets import MuMapDataset, HEADER_DISC_FIRST, HEADER_DISC_LAST
-from mu_map.util import to_grayscale, COLOR_WHITE, COLOR_BLACK
+    :param row: the row of meta configuration file of a dataset
+    :param μ_map: the μ_map
+    :return: the μ_map with according slices removed
+    """
+    _res = μ_map
 
+    if row[HEADER_DISC_FIRST]:
+        _res = _res[1:]
+
+    if row[HEADER_DISC_LAST]:
+        _res = _res[:-1]
+
+    return _res
 
 
 if __name__ == "__main__":
+    import argparse
+
+    import cv2 as cv
+
+    from mu_map.data.datasets import MuMapDataset
+    from mu_map.util import to_grayscale, COLOR_WHITE, COLOR_BLACK
+
     parser = argparse.ArgumentParser(
         description="review all μ-maps in a dataset for broken slices at the start or end",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -65,21 +90,41 @@ if __name__ == "__main__":
         slice_first = cv.resize(slice_first, (512, 512), interpolation=cv.INTER_AREA)
         slice_last = cv.resize(slice_last, (512, 512), interpolation=cv.INTER_AREA)
 
-        cv.putText(slice_first, "First", (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_BLACK, 3)
-        cv.putText(slice_first, "First", (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_WHITE, 1)
+        cv.putText(
+            slice_first, "First", (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_BLACK, 3
+        )
+        cv.putText(
+            slice_first, "First", (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_WHITE, 1
+        )
 
-        cv.putText(slice_last, "Last", (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_BLACK, 3)
-        cv.putText(slice_last, "Last", (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_WHITE, 1)
+        cv.putText(
+            slice_last, "Last", (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_BLACK, 3
+        )
+        cv.putText(
+            slice_last, "Last", (0, 30), cv.FONT_HERSHEY_SIMPLEX, 1, COLOR_WHITE, 1
+        )
 
         space = np.full((slice_first.shape[0], 10), 239, np.uint8)
 
         to_show = np.hstack((slice_first, space, slice_last))
         textposition = (0, to_show.shape[0] - 12)
         cv.putText(
-            to_show, controls, textposition, cv.FONT_HERSHEY_SIMPLEX, 0.75, COLOR_BLACK, 3
+            to_show,
+            controls,
+            textposition,
+            cv.FONT_HERSHEY_SIMPLEX,
+            0.75,
+            COLOR_BLACK,
+            3,
         )
         cv.putText(
-            to_show, controls, textposition, cv.FONT_HERSHEY_SIMPLEX, 0.75, COLOR_WHITE, 1
+            to_show,
+            controls,
+            textposition,
+            cv.FONT_HERSHEY_SIMPLEX,
+            0.75,
+            COLOR_WHITE,
+            1,
         )
 
         cv.imshow(wname, to_show)