diff --git a/mu_map/data/datasets.py b/mu_map/data/datasets.py
index 0f0a321a8de0e30da93b830563ba4ffb789084ce..05e9e6c63ea6af1b1b41f2bc3847ff6d23c354ba 100644
--- a/mu_map/data/datasets.py
+++ b/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
diff --git a/mu_map/data/review_mu_map.py b/mu_map/data/review_mu_map.py
index 720382713feeed5524427a78ebab3891b8ba0fbb..21e66eabdf36fefb7c02fd746c99a12ff5ba3947 100644
--- a/mu_map/data/review_mu_map.py
+++ b/mu_map/data/review_mu_map.py
@@ -1,14 +1,39 @@
-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)