Adding support for .jpk-qi-data and .bin files

.gitignore

@@ -3,6 +3,9 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+AFMReader/data/*
+AFMReader/notebooks/*
+
 # C extensions
 *.so
 

AFMReader/asd.py

@@ -288,6 +288,38 @@ def load_asd(file_path: str | Path, channel: str):
         return frames, pixel_to_nanometre_scaling_factor, header_dict
 
 
+def get_asd_channels(file_path: Path):
+    """
+    Get the channels available in given .asd file.
+
+    Parameters
+    ----------
+    file_path : Path
+        Path to the .asd file.
+
+    Returns
+    -------
+    list
+        List of channels available in the .asd file.
+    """
+    with Path.open(file_path, "rb", encoding=None) as open_file:  # pylint: disable=unspecified-encoding
+        file_version = read_file_version(open_file)
+
+        if file_version == 0:
+            header_dict = read_header_file_version_0(open_file)
+
+        elif file_version == 1:
+            header_dict = read_header_file_version_1(open_file)
+
+        elif file_version == 2:
+            header_dict = read_header_file_version_2(open_file)
+        else:
+            raise ValueError(
+                f"File version {file_version} unknown. Please add support if you know how to decode this file version."
+            )
+        return [header_dict["channel1"], header_dict["channel2"]]
+
+
 def read_file_version(open_file: BinaryIO) -> int:
     """
     Read the file version from an open asd file. File versions are 0, 1 and 2.

AFMReader/general_loader.py

@@ -1,16 +1,17 @@
 """Switchboard for input files."""
 
 from pathlib import Path
+from typing import Any
 
 import numpy.typing as npt
 
-from AFMReader import asd, gwy, h5_jpk, ibw, jpk, spm, stp, top, topostats
+from AFMReader import asd, gwy, h5_jpk, ibw, jpk, raw_bin, spm, stp, top, topostats, jpk_qi
 from AFMReader.logging import logger
 
 logger.enable(__package__)
 
 
-# pylint: disable=too-few-public-methods
+# pylint: disable=too-few-public-methods,too-many-branches,too-many-statements,fixme
 class LoadFile:
     """
     Class to handle the general loading of an AFM file.
@@ -21,9 +22,11 @@ class LoadFile:
             Path to the AFM image.
         channel : str
             Channel to extract from the AFM image.
+        kwargs : dict, optional
+            Additional keyword arguments to pass to the specific loaders.
     """
 
-    def __init__(self, filepath: str | Path, channel: str):
+    def __init__(self, filepath: str | Path, channel: str, kwargs: dict | None = None):
         """
         Initialise the general LoadFile class with a filepath and channel.
 
@@ -33,39 +36,82 @@ def __init__(self, filepath: str | Path, channel: str):
             Path to the AFM image.
         channel : str
             Channel to extract from the AFM image.
+        kwargs : dict, optional
+            Additional keyword arguments to pass to the specific loaders.
         """
         self.filepath = Path(filepath)
         self.channel = channel
         self.suffix = self.filepath.suffix
+        self.loaded_curves = False
+        self.kwargs = kwargs if kwargs else {}
 
-    def load(self) -> tuple[npt.NDArray | str, float | None]:  # noqa: C901
+        # Store heavy loaded data in a dict to avoid having to reload it
+        self.cached_data: dict[str, Any] = {}
+
+    def load(  # noqa: C901
+        self, channel: str | None = None, kwargs: dict | None = None
+    ) -> tuple[npt.NDArray | str, float | None] | tuple[npt.NDArray | str, float | None, Any]:
         """
         Generally loads a file type that can be handled by AFMReader.
 
+        Parameters
+        ----------
+        channel : str, optional
+            Overriding channel to extract from the AFM image.
+        kwargs : dict, optional
+            Additional keyword arguments to pass to the specific loaders.
+
         Returns
         -------
         tuple
             The image data (stack if ''.asd'' or ''.h5-jpk'') and the pixel to nanometre scaling ratio.
+            If curve data is found, also return the curve data (a large dict of all the curves).
 
         Raises
         ------
         ValueError
             Where the channel is not found, returned as a tuple of "error message" and "None" so that this can be
             propagated to Napari without outright failing.
         """
+        if channel:
+            self.channel = channel
+        if kwargs:
+            self.kwargs = kwargs
         try:
             if self.suffix == ".asd":
                 image, pixel_to_nanometre_scaling_factor, _ = asd.load_asd(self.filepath, self.channel)
             elif self.suffix == ".gwy":
                 image, pixel_to_nanometre_scaling_factor = gwy.load_gwy(self.filepath, self.channel)
             elif self.suffix == ".ibw":
                 image, pixel_to_nanometre_scaling_factor = ibw.load_ibw(self.filepath, self.channel)
-            elif self.suffix == ".jpk":
+            elif self.suffix in [".jpk", ".jpk-qi-image"]:
                 image, pixel_to_nanometre_scaling_factor = jpk.load_jpk(self.filepath, self.channel)
             elif self.suffix == ".spm":
                 image, pixel_to_nanometre_scaling_factor = spm.load_spm(self.filepath, self.channel)
             elif self.suffix == ".h5-jpk":
-                image, pixel_to_nanometre_scaling_factor, _ = h5_jpk.load_h5jpk(self.filepath, self.channel)
+                h5_returned = h5_jpk.load_h5jpk(self.filepath, self.channel, load_curves=not self.loaded_curves)
+                if len(h5_returned) == 3:
+                    image, pixel_to_nanometre_scaling_factor, _ = h5_returned  # type: ignore[misc]
+                elif len(h5_returned) == 4:
+                    image, pixel_to_nanometre_scaling_factor, _, curve_data = h5_returned  # type: ignore[misc]
+                    self.loaded_curves = True
+                    return image, pixel_to_nanometre_scaling_factor, curve_data
+                else:
+                    logger.error(f"Loading h5-jpk file returned unexpected number of items: {len(h5_returned)}")
+            elif self.suffix == ".jpk-qi-data":
+                if "jpk_qi_loader" not in self.cached_data:
+                    self.cached_data["jpk_qi_loader"] = jpk_qi.jpk_qi_loader(
+                        filepath=self.filepath, channel=self.channel, **self.kwargs
+                    )
+                jpk_qi_returned = self.cached_data["jpk_qi_loader"].load(channel=self.channel, **self.kwargs)
+                if len(jpk_qi_returned) == 2:
+                    image, pixel_to_nanometre_scaling_factor = jpk_qi_returned
+                elif len(jpk_qi_returned) == 3:
+                    image, pixel_to_nanometre_scaling_factor, curve_data = jpk_qi_returned
+                    self.loaded_curves = True
+                    return image, pixel_to_nanometre_scaling_factor, curve_data
+                else:
+                    logger.error(f"Loading h5-jpk file returned unexpected number of items: {len(jpk_qi_returned)}")
             elif self.suffix == ".stp":
                 image, pixel_to_nanometre_scaling_factor = stp.load_stp(self.filepath)
             elif self.suffix == ".top":
@@ -82,13 +128,48 @@ def load(self) -> tuple[npt.NDArray | str, float | None]:  # noqa: C901
                         f"'{self.channel}' not in available image keys: "
                         f"{[im for im in image_keys if im in topostats_keys]}"
                     ) from exc
+            elif self.suffix == ".bin":
+                image, pixel_to_nanometre_scaling_factor = raw_bin.load_bin(self.filepath, **self.kwargs)
             else:
                 raise ValueError(f"File type '{self.suffix}' is not currently handled by AFMReader.")
 
             return image, pixel_to_nanometre_scaling_factor
 
         except ValueError as e:
             logger.error(f"{e}")
-            return (e, None)  # cheeky return of an image, px2nm-like tuple object to propagate error message to Napari
+            raise e
 
-    # scope for a "check what channels are available" function similar to above.
+    def get_available_channels(self):  # noqa: C901
+        """
+        Get the available channels for the file type.
+
+        Returns
+        -------
+        list
+            List of available channels.
+        """
+        if self.suffix == ".asd":
+            available_channels = asd.get_asd_channels(self.filepath)
+        elif self.suffix == ".gwy":
+            available_channels = gwy.get_gwy_channels(self.filepath)
+        elif self.suffix == ".ibw":
+            available_channels = ibw.get_ibw_channels(self.filepath)
+        elif self.suffix in [".jpk", ".jpk-qi-image"]:
+            available_channels = jpk.get_jpk_channels(self.filepath)
+        elif self.suffix == ".spm":
+            available_channels = spm.get_spm_channels(self.filepath)
+        elif self.suffix == ".h5-jpk":
+            available_channels = h5_jpk.get_h5jpk_channels(self.filepath)
+        elif self.suffix == ".jpk-qi-data":
+            if "jpk_qi_loader" not in self.cached_data:
+                self.cached_data["jpk_qi_loader"] = jpk_qi.jpk_qi_loader(filepath=self.filepath, **self.kwargs)
+            available_channels = self.cached_data["jpk_qi_loader"].get_available_channels()
+        elif self.suffix == ".topostats":
+            available_channels = ["image", "image_original"]
+        elif self.suffix == ".bin":
+            available_channels = raw_bin.get_bin_channels()
+        elif self.suffix in [".stp", ".top"]:
+            return []
+        else:
+            raise ValueError(f"File type '{self.suffix}' is not currently handled by AFMReader.")
+        return available_channels

AFMReader/gwy.py

@@ -10,6 +10,32 @@
 from AFMReader.io import read_char, read_double, read_null_terminated_string, read_uint32
 
 
+def get_gwy_channels(file_path):
+    """
+    Extract a list of available channels and their corresponding dictionary key ids from the `.gwy` file.
+
+    Parameters
+    ----------
+    file_path : Path or str
+        Path to the .gwy file.
+
+    Returns
+    -------
+    list
+        List of available channels.
+    """
+    image_data_dict: dict[Any, Any] = {}
+    with Path.open(file_path, "rb") as open_file:  # pylint: disable=unspecified-encoding
+        # Read header
+        header = open_file.read(4)
+        logger.debug(f"Gwy file header: {header.decode}")
+
+        gwy_read_object(open_file, data_dict=image_data_dict)
+    channel_ids = gwy_get_channels(gwy_file_structure=image_data_dict)
+
+    return list(channel_ids)
+
+
 def load_gwy(file_path: Path | str, channel: str) -> tuple[np.ndarray[Any, np.float64], float]:
     """
     Extract image and pixel to nm scaling from the .gwy file.