Merge pull request #113 from AFM-SPM/SylviaWhittle/77-fileformat-stp

Add file format STP

AFMReader/stp.py

@@ -0,0 +1,107 @@
+"""For decoding and loading .stp AFM file format into Python Numpy arrays."""
+
+from __future__ import annotations
+from pathlib import Path
+import re
+
+import numpy as np
+
+from AFMReader.logging import logger
+from AFMReader.io import read_double
+
+logger.enable(__package__)
+
+
+# pylint: disable=too-many-locals
+def load_stp(  # noqa: C901 (ignore too complex)
+    file_path: Path | str, header_encoding: str = "latin-1"
+) -> tuple[np.ndarray, float]:
+    """
+    Load image from STP files.
+
+    Parameters
+    ----------
+    file_path : Path | str
+        Path to the .stp file.
+    header_encoding : str
+        Encoding to use for the header of the file. Default is ''latin-1''.
+
+    Returns
+    -------
+    tuple[np.ndarray, float]
+        A tuple containing the image and its pixel to nanometre scaling value.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the file is not found.
+    ValueError
+        If any of the required information are not found in the header.
+    NotImplementedError
+        If the image is non-square.
+    """
+    logger.info(f"Loading image from : {file_path}")
+    file_path = Path(file_path)
+    filename = file_path.stem
+    try:
+        with Path.open(file_path, "rb") as open_file:  # pylint: disable=unspecified-encoding
+            # grab the beggining message, assume that it's within the first 150 bytes
+            beginning_message = str(open_file.read(150))
+            # find the header size in the beginning message
+            header_size_match = re.search(r"Image header size: (\d+)", beginning_message)
+            if header_size_match is None:
+                raise ValueError(f"[{filename}] : 'Image header size' not found in image raw bytes.")
+            header_size = int(header_size_match.group(1))
+
+            # Return to start of file
+            open_file.seek(0)
+            # Read the header
+            header = open_file.read(header_size)
+
+            # decode the header bytes
+            header_decoded = header.decode(header_encoding)
+
+            # find num rows
+            rows_match = re.search(r"Number of rows: (\d+)", header_decoded)
+            if rows_match is None:
+                raise ValueError(f"[{filename}] : 'rows' not found in file header.")
+            rows = int(rows_match.group(1))
+            cols_match = re.search(r"Number of columns: (\d+)", header_decoded)
+            if cols_match is None:
+                raise ValueError(f"[{filename}] : 'cols' not found in file header.")
+            cols = int(cols_match.group(1))
+            x_real_size_match = re.search(r"X Amplitude: (\d+) nm", header_decoded)
+            if x_real_size_match is None:
+                raise ValueError(f"[{filename}] : 'X Amplitude' not found in file header.")
+            x_real_size = float(x_real_size_match.group(1))
+            y_real_size_match = re.search(r"Y Amplitude: (\d+) nm", header_decoded)
+            if y_real_size_match is None:
+                raise ValueError(f"[{filename}] : 'Y Amplitude' not found in file header.")
+            y_real_size = float(y_real_size_match.group(1))
+            if x_real_size != y_real_size:
+                raise NotImplementedError(
+                    f"[{filename}] : X scan size (nm) does not equal Y scan size (nm) ({x_real_size}, {y_real_size})"
+                    "we don't currently support non-square images."
+                )
+
+            # Calculate pixel to nm scaling
+            pixel_to_nm_scaling = x_real_size / cols
+
+            # Read R x C matrix of doubles
+            image_list = []
+            for _ in range(rows):
+                row = []
+                for _ in range(cols):
+                    row.append(read_double(open_file))
+                image_list.append(row)
+            image = np.array(image_list)
+
+    except FileNotFoundError as e:
+        logger.error(f"[{filename}] : File not found : {file_path}")
+        raise e
+    except Exception as e:
+        logger.error(f"[{filename}] : {e}")
+        raise e
+
+    logger.info(f"[{filename}] : Extracted image.")
+    return (image, pixel_to_nm_scaling)

README.md

@@ -35,6 +35,7 @@ Supported file formats
 | `.jpk`      | [Bruker](https://www.bruker.com/) |
 | `.topostats`| [TopoStats](https://github.com/AFM-SPM/TopoStats)  |
 | `.gwy`      | [Gwydion](<http://gwyddion.net>) |
+| `.stp`      | [WSXM AFM software files](http://www.wsxm.eu) |
 
 Support for the following additional formats is planned. Some of these are already supported in TopoStats and are
 awaiting refactoring to move their functionality into AFMReader these are denoted in bold below.
@@ -123,6 +124,17 @@ from AFMReader.jpk import load_jpk
 image, pixel_to_nanometre_scaling_factor = load_jpk(file_path="./my_jpk_file.jpk", channel="height_trace")
 ```
 
+### .stp
+
+You can open `.stp` files using the `load_stp` function. Just pass in the path
+to the file you want to use.
+
+```python
+from AFMReader.stp import load_stp
+
+image, pixel_to_nanometre_scaling_factor = load_stp(file_path="./my_stp_file.stp")
+```
+
 ## Contributing
 
 Bug reports and feature requests are welcome. Please search for existing issues, if none relating to your bug/feature

examples/example_01.ipynb

@@ -257,6 +257,47 @@
     "image, pixel_to_nm_scaling, metadata = load_topostats(file_path=FILE)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Plot the image\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.imshow(image, cmap=\"afmhot\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# STP Files"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Import the load_stp function from AFMReader\n",
+    "from AFMReader.stp import load_stp"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load the STP file as an image and pixel to nm scaling factor\n",
+    "FILE = \"../tests/resources/sample_0.stp\"\n",
+    "image, pixel_to_nm_scaling = load_stp(file_path=FILE)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -273,7 +314,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "new",
+   "display_name": "afmreader",
    "language": "python",
    "name": "python3"
   },
@@ -287,7 +328,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.9.19"
   }
  },
  "nbformat": 4,

tests/test_stp.py

@@ -0,0 +1,39 @@
+"""Test the loading of .stp files."""
+
+from pathlib import Path
+import pytest
+
+import numpy as np
+
+from AFMReader.stp import load_stp
+
+BASE_DIR = Path.cwd()
+RESOURCES = BASE_DIR / "tests" / "resources"
+
+
+@pytest.mark.parametrize(
+    (
+        "file_name",
+        "expected_pixel_to_nm_scaling",
+        "expected_image_shape",
+        "expected_image_dtype",
+        "expected_image_sum",
+    ),
+    [pytest.param("sample_0.stp", 0.9765625, (512, 512), float, -15070620.440757688)],
+)
+def test_load_stp(
+    file_name: str,
+    expected_pixel_to_nm_scaling: float,
+    expected_image_shape: tuple[int, int],
+    expected_image_dtype: type,
+    expected_image_sum: float,
+) -> None:
+    """Test the normal operation of loading a .stp file."""
+    file_path = RESOURCES / file_name
+    result_image, result_pixel_to_nm_scaling = load_stp(file_path=file_path)
+
+    assert result_pixel_to_nm_scaling == expected_pixel_to_nm_scaling
+    assert isinstance(result_image, np.ndarray)
+    assert result_image.shape == expected_image_shape
+    assert result_image.dtype == expected_image_dtype
+    assert result_image.sum() == expected_image_sum