diff --git a/mu_map/file/interfile.py b/mu_map/file/interfile.py index 85cef0ad767225bc3a2e84d31d4a2215caccd051..a916dda966588edde0f8148a4ea269d03c2eab6f 100644 --- a/mu_map/file/interfile.py +++ b/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)