Pyto tomoloading

pyproject.toml

@@ -30,11 +30,11 @@ dependencies = [
     "torch",
     "imageio",
     "scipy",
-    "mrcfile",
     "typer[all]",
     "scikit-image",
     "simpleitk",
-    "pandas"
+    "pandas",
+    "future"
 ]
 
 

src/membrain_seg/segmentation/dataloading/data_utils.py

@@ -1,13 +1,14 @@
 import csv
 import os
-from typing import Any, Callable, Dict, Tuple, Union
+from typing import Any, Callable, Tuple, Union
 
-import mrcfile
 import numpy as np
 import SimpleITK as sitk
 from skimage.util import img_as_float32
 from torch import Tensor, device
 
+from .tomo_io import ImageIO
+
 
 def make_directory_if_not_exists(path: str):
     """
@@ -109,13 +110,15 @@ def load_data_for_inference(data_path: str, transforms: Callable, device: device
         dimension added, and is moved to the appropriate device.
 
     """
-    tomogram = load_tomogram(data_path, normalize_data=True)
+    tomogram, mrc = load_tomogram(
+        data_path, normalize_data=True, return_mrc_object=True
+    )
     tomogram = np.expand_dims(tomogram, 0)
 
     new_data = transforms(tomogram)
     new_data = new_data.unsqueeze(0)  # Add batch dimension
     new_data = new_data.to(device)
-    return new_data
+    return new_data, mrc
 
 
 def store_segmented_tomograms(
@@ -124,6 +127,7 @@ def store_segmented_tomograms(
     orig_data_path: str,
     ckpt_token: str,
     store_probabilities: bool = False,
+    mrc_object: ImageIO = None,
 ) -> None:
     """
     Helper function for storing output segmentations.
@@ -144,6 +148,10 @@ def store_segmented_tomograms(
         Checkpoint token.
     store_probabilities : bool, optional
         If True, probabilities are stored before thresholding.
+    mrc_object: ImageIO, optional
+        If given, the mrc_object will be used to retain header information
+        from another tomogram. This way, pixel sizes and other header
+        information is not lost.
     """
     # Create out directory if it doesn't exist yet
     make_directory_if_not_exists(out_folder)
@@ -155,13 +163,13 @@ def store_segmented_tomograms(
         out_file = os.path.join(
             out_folder, os.path.basename(orig_data_path)[:-4] + "_scores.mrc"
         )
-        store_tomogram(out_file, predictions_np)
+        store_tomogram(out_file, predictions_np, mrc=mrc_object)
     predictions_np_thres = predictions.squeeze(0).squeeze(0).cpu().numpy() > 0.0
     out_file_thres = os.path.join(
         out_folder,
         os.path.basename(orig_data_path)[:-4] + "_" + ckpt_token + "_segmented.mrc",
     )
-    store_tomogram(out_file_thres, predictions_np_thres)
+    store_tomogram(out_file_thres, predictions_np_thres, mrc=mrc_object)
     print("MemBrain has finished segmenting your tomogram.")
 
 
@@ -206,8 +214,7 @@ def write_nifti(out_file: str, image: np.ndarray) -> None:
 
 def load_tomogram(
     filename: str,
-    return_pixel_size: bool = False,
-    return_header: bool = False,
+    return_mrc_object: bool = False,
     normalize_data: bool = False,
 ) -> Union[np.ndarray, Tuple[np.ndarray, Any]]:
     """
@@ -219,10 +226,8 @@ def load_tomogram(
     ----------
     filename : str
         File name of the tomogram to load.
-    return_pixel_size : bool, optional
-        If True, return pixel size.
-    return_header : bool, optional
-        If True, return header.
+    return_mrc_object : bool, optional
+        If True, returns mrc object that contains more header information.
     normalize_data : bool, optional
         If True, normalize data.
 
@@ -232,33 +237,19 @@ def load_tomogram(
         Numpy array of the loaded data.
 
     """
-    with mrcfile.open(filename, permissive=True) as mrc:
-        data = np.array(mrc.data)
-        data = np.transpose(data, (2, 1, 0))
-        cella = mrc.header.cella
-        cellb = mrc.header.cellb
-        origin = mrc.header.origin
-        pixel_spacing = np.array([mrc.voxel_size.x, mrc.voxel_size.y, mrc.voxel_size.z])
-        header_dict = {
-            "cella": cella,
-            "cellb": cellb,
-            "origin": origin,
-            "pixel_spacing": pixel_spacing,
-        }
-        if normalize_data:
-            data = img_as_float32(data)
-            data -= np.mean(data)
-            data /= np.std(data)
-        if return_pixel_size:
-            return data, mrc.voxel_size
-        if return_header:
-            return data, header_dict
+    mrc = ImageIO(filename)
+    mrc.read()
+    data = mrc.data
+    if normalize_data:
+        data = img_as_float32(data)
+        data -= np.mean(data)
+        data /= np.std(data)
+    if return_mrc_object:
+        return data, mrc
     return data
 
 
-def store_tomogram(
-    filename: str, tomogram: np.ndarray, header_dict: Dict[str, Any] = None
-) -> None:
+def store_tomogram(filename: str, tomogram: np.ndarray, mrc: ImageIO = None) -> None:
     """
     Store tomogram in specified path.
 
@@ -268,19 +259,15 @@ def store_tomogram(
         Name of the file to store the tomogram.
     tomogram : np.ndarray
         The tomogram data.
-    header_dict : Dict[str, Any], optional
-        Header dictionary to use.
+    mrc : ImageIO, optional
+        Mrc object containing header information.
 
     """
-    if tomogram.dtype != np.int8:
-        tomogram = np.array(tomogram, dtype=np.float32)
-    tomogram = np.transpose(tomogram, (2, 1, 0))
-    with mrcfile.new(filename, overwrite=True) as mrc:
-        mrc.set_data(tomogram)
-        if header_dict is not None:
-            mrc.header.cella = header_dict["cella"]
-            mrc.header.cellb = header_dict["cellb"]
-            mrc.header.origin = header_dict["origin"]
+    if mrc is None:
+        mrc = ImageIO()
+    if tomogram.dtype == bool:
+        tomogram = tomogram.astype("ubyte")
+    mrc.write(filename, tomogram)
 
 
 def normalize_tomogram(tomogram: np.ndarray) -> np.ndarray:

src/membrain_seg/segmentation/dataloading/microscope_db.py

@@ -0,0 +1,173 @@
+#
+# THIS CODE IS COPIED FROM PYTO:
+# https://github.com/vladanl/Pyto/tree/master/pyto/io
+# THIS IS TO AVOID THE INSTALLATION ISSUES AND NO PYPI AVAILABLITY
+# PLASE CITE ACCORDINGLY
+#
+#
+# Lorenz Lamm, July 2023
+
+
+# pixelsize (in pm) at the specimen level for different nominal magnifications
+pixelsize = {}
+
+# physical CCD pixelsize (in um)
+ccd_pixelsize = {}
+
+# number of pixels on CCD
+n_pixels = {}
+
+# nominal magnification for different real magnifications
+nominal_mag = {}
+
+# number of counts per electron
+conversion = {}
+
+
+###################################################
+#
+# Titan 2 with K2 in counting mode recorded with SerialEM
+#
+
+# screen up (nominal) magnification vs. pixel size (at specimen level)
+pixelsize["titan-2_k2-count_sem"] = {}
+
+# counts per electron
+conversion["titan-2_k2-count_sem"] = 15
+
+
+###################################################
+#
+# Polara 2 with K2 in counting mode recorded with SerialEM
+#
+
+# screen up (nominal) magnification vs. pixel size (at specimen level)
+pixelsize["polara-2_k2-count_sem"] = {}
+
+# counts per electron
+conversion["polara-2_k2-count_sem"] = 19
+
+# CCD pixel size
+# ccd_pixelsize['polara-2_01-09'] = 30000
+
+# number of pixels at CCD
+# n_pixels['polara-2_01-09'] = 2048
+
+
+###################################################
+#
+# Polara 1 from 01.07
+#
+
+# screen up (nominal) magnification vs. pixel size (at specimen level)
+pixelsize["polara-1_01-07"] = {
+    18000: 1230,
+    22500: 979,
+    27500: 805,
+    34000: 661,
+    41000: 545,
+    50000: 446,
+    61000: 364,
+}
+
+# counts per electron
+conversion["polara-1_01-07"] = 5.91
+
+# CCD pixel size
+ccd_pixelsize["polara-1_01-07"] = 30000
+
+# number of pixels at CCD
+n_pixels["polara-1_01-07"] = 2048
+
+
+###################################################
+#
+# Polara 1 from 01.09
+#
+
+# screen up (nominal) magnification vs. pixel size (at specimen level)
+pixelsize["polara-1_01-09"] = pixelsize["polara-1_01-07"]
+
+# counts per electron
+conversion["polara-1_01-09"] = 2.3
+
+# CCD pixel size
+ccd_pixelsize["polara-1_01-09"] = 30000
+
+# number of pixels at CCD
+n_pixels["polara-1_01-09"] = 2048
+
+
+###################################################
+#
+# Polara 2 from 01.09
+#
+
+# screen up (nominal) magnification vs. pixel size (at specimen level)
+pixelsize["polara-2_01-09"] = {
+    9300: 1372,
+    13500: 956,
+    18000: 713,
+    22500: 572,
+    27500: 468,
+    34000: 381,
+}
+
+# counts per electron
+conversion["polara-2_01-09"] = 8.1
+
+# CCD pixel size
+ccd_pixelsize["polara-2_01-09"] = 30000
+
+# number of pixels at CCD
+n_pixels["polara-2_01-09"] = 2048
+
+
+###################################################
+#
+# Krios 2, Falcon detector from 09.2011
+#
+
+# screen up (nominal) magnification vs. pixel size [fm] (at specimen level)
+pixelsize["krios-2_falcon_05-2011"] = {18000: 475}
+
+# counts per electron (JP 12.2011)
+conversion["krios-2_falcon_05-2011"] = 134.0
+
+# CCD pixel size (not determined yet)
+ccd_pixelsize["krios-2_falcon_05-2011"] = 1
+
+# number of pixels at CCD
+n_pixels["krios-2_falcon_05-2011"] = 4096
+
+
+###################################################
+#
+# F20 eagle camera
+#
+
+# screen up (nominal) magnification vs. pixel size [fm] (at specimen level)
+pixelsize["f20_eagle"] = {}
+
+# counts per electron
+conversion["f20_eagle"] = 73.0
+
+
+###################################################
+#
+# CM 300
+#
+
+# screen up (nominal) magnification vs. pixel size (at specimen level)
+pixelsize["cm300"] = {13500: 1147, 17500: 821, 23000: 682, 27500: 547}
+
+nominal_mag["cm300"] = {26157: 13500, 36527: 17500, 43974: 23000, 54844: 27500}
+
+# counts per electron
+conversion["cm300"] = 5.5
+
+# CCD pixel size
+ccd_pixelsize["cm300"] = 30000
+
+# number of pixels at CCD
+n_pixels["cm300"] = 2048