refactor interfile module

mu_map/file/interfile.py

@@ -4,34 +4,6 @@ from typing import Dict, Tuple
 
 import numpy as np
 
-Interfile = Tuple[Dict[str, str], np.ndarray]
-
-@dataclass
-class _InterfileKeys:
-    """
-    Data class defining keys for an Interfile header.
-    """
-    placeholder: str = "{_}"
-
-    _dim: str = f"matrix size [{placeholder}]"
-    _spacing: str = f"scaling factor (mm/pixel) [{placeholder}]"
-
-    data_file: str = "name of data file"
-    n_projections: str = "number of projections"
-    bytes_per_pixel: str = "number of bytes per pixel"
-    number_format: str = "number format"
-
-    patient_orientation: str = "patient orientation"
-    patient_rotation: str = "supine"
-
-    def dim(self, index: int) -> str:
-        return self._dim.replace(self.placeholder, str(index))
-
-    def spacing(self, index: int) -> str:
-        return self._spacing.replace(self.placeholder, str(index))
-
-InterfileKeys = _InterfileKeys()
-
 
 """
 A template of an INTERFILE header.
@@ -70,6 +42,54 @@ END OF INTERFILE :=
 """
 
 
+"""
+Default extension for the Interfile header.
+"""
+EXTENSION_INTERFILE_HEADER = ".hv"
+"""
+Default extension for the Interfile image.
+"""
+EXTENSION_INTERFILE_IMAGE = ".v"
+
+
+@dataclass
+class Interfile:
+    """
+    Interfile data type wrapping the header and the image
+    into an object.
+    """
+    header: Dict[str, str]
+    image: np.ndarray
+
+    def __iter__(self):
+        return iter((self.header, self.image))
+
+
+@dataclass
+class _InterfileKeys:
+    """
+    Data class defining keys for an Interfile header.
+    """
+    placeholder: str = "{_}"
+
+    _dim: str = f"matrix size [{placeholder}]"
+    _spacing: str = f"scaling factor (mm/pixel) [{placeholder}]"
+
+    data_file: str = "name of data file"
+    n_projections: str = "number of projections"
+    bytes_per_pixel: str = "number of bytes per pixel"
+    number_format: str = "number format"
+
+    patient_orientation: str = "patient orientation"
+    patient_rotation: str = "supine"
+
+    def dim(self, index: int) -> str:
+        return self._dim.replace(self.placeholder, str(index))
+
+    def spacing(self, index: int) -> str:
+        return self._spacing.replace(self.placeholder, str(index))
+InterfileKeys = _InterfileKeys()
+
 
 def type_by_format(number_format: str, bytes_per_pixel: int) -> type:
     """
@@ -114,7 +134,7 @@ def parse_interfile_header(filename: str) -> Dict[str, str]:
         return parse_interfile_header_str(f.read())
 
 
-def load_interfile(filename: str) -> Tuple[Dict[str, str], np.ndarray]:
+def load_interfile(filename: str) -> Interfile:
     """
     Load an INTERFILE header and its image as a numpy array.
 
@@ -135,7 +155,7 @@ def load_interfile(filename: str) -> Tuple[Dict[str, str], np.ndarray]:
     with open(data_file, mode="rb") as f:
         image = np.frombuffer(f.read(), dtype)
     image = image.reshape((dim_z, dim_y, dim_x))
-    return header, image.copy()
+    return Interfile(header, image.copy())
 
 
 def load_interfile_img(filename: str) -> np.ndarray:
@@ -145,14 +165,16 @@ def load_interfile_img(filename: str) -> np.ndarray:
     :param filename: the filename of the INTERFILE header file
     :return: the image as a numpy array
     """
-    _, image = load_interfile(filename)
-    return image
+    interfile = load_interfile(filename)
+    return interfile.image
 
 
-def write_interfile(filename: str, header: Dict[str, str], image: np.ndarray):
+def write_interfile(filename: str, interfile: Interfile):
     filename = os.path.splitext(filename)[0]
-    filename_data = f"{filename}.v"
-    filename_header = f"{filename}.hv"
+    filename_data = f"{filename}{EXTENSION_INTERFILE_IMAGE}"
+    filename_header = f"{filename}{EXTENSION_INTERFILE_HEADER}"
+
+    header, image = interfile
 
     header[InterfileKeys.data_file] = os.path.basename(filename_data)
     header[InterfileKeys.dim(3)] = str(image.shape[0])
@@ -181,34 +203,34 @@ if __name__ == "__main__":
     args = parser.parse_args()
     args.out = args.interfile if args.out is None else args.out
 
-    header, image = load_interfile(args.interfile)
+    interfile = load_interfile(args.interfile)
 
     if args.cmd == "fill":
         value = args.param
-        image.fill(value)
+        interfile.image.fill(value)
     if args.cmd == "pad":
         target_size = args.param
-        real_size = image.shape[0]
+        real_size = interfile.image.shape[0]
         assert target_size > real_size, f"Cannot pad from a larger size {real_size} to a smaller size {target_size}"
 
         diff = target_size - real_size
         pad_lower = math.ceil(diff / 2)
         pad_upper = math.floor(diff / 2)
-        image = np.pad(image, ((pad_lower, pad_upper), (0, 0), (0, 0)))
+        interfile.image = np.pad(interfile.image, ((pad_lower, pad_upper), (0, 0), (0, 0)))
     if args.cmd == "crop":
         target_size = args.param
-        real_size = image.shape[0]
+        real_size = interfile.image.shape[0]
         assert target_size < real_size, f"Cannot crop from a smaller size {real_size} to a larger size {target_size}"
 
         diff = target_size / 2
         center = real_size // 2
         crop_lower = center - math.floor(diff)
         crop_upper = center + math.ceil(diff)
-        image = image[crop_lower:crop_upper]
+        interfile.image = interfile.image[crop_lower:crop_upper]
     if args.cmd == "flip":
-        image = image[:, :, ::-1]
+        interfile.image = interfile.image[:, :, ::-1]
 
-    write_interfile(args.out, header, image)
+    write_interfile(args.out, interfile)