From 5b6d875804136212d0cd7ecae324179bc1aa9506 Mon Sep 17 00:00:00 2001
From: Roy Hoitink <L.D.Hoitink@uu.nl>
Date: Wed, 20 Mar 2024 09:43:07 +0100
Subject: [PATCH] Improve loading of h5file, automatically detect format

---
 src/simulatedmicroscopy/image.py | 48 ++++++++++++++++++++++++--------
 tests/test_image.py              | 21 +++++++++++++-
 2 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/src/simulatedmicroscopy/image.py b/src/simulatedmicroscopy/image.py
index e4401b6..5858323 100644
--- a/src/simulatedmicroscopy/image.py
+++ b/src/simulatedmicroscopy/image.py
@@ -162,7 +162,7 @@ def save_h5file(
 
     @classmethod
     def load_h5file(cls, filename: str) -> type[Image]:
-        """Load data from h5 file (custom format)
+        """Load data from h5 file (custom format or HuygensImage)
 
         Parameters
         ----------
@@ -176,19 +176,40 @@ def load_h5file(cls, filename: str) -> type[Image]:
         """
 
         with h5py.File(filename, "r") as f:
-            image = f["Image"][()]
-            pixel_sizes = [
-                float(f[f"Metadata/DimensionScale{dim.upper()}"][()])
-                for dim in list("zyx")
-            ]
-            if "PixelCoordinates" in f["Metadata"]:
-                pixel_coordinates = f["Metadata/PixelCoordinates"][()]
+            root_elements = list(f.keys())
+            if "Image" in root_elements and "Metadata" in root_elements:
+                return cls._load_h5file_custom(f)
             else:
-                pixel_coordinates = None
+                return HuygensImage(filename)
+    
+    @staticmethod
+    def _load_h5file_custom(f: h5py.File) -> type[Image]:
+        """Load data from h5 file (custom format)
+
+        Parameters
+        ----------
+        f : h5py.File
+            File to load from
+
+        Returns
+        -------
+        Image
+            Resulting image with correct pixel sizes
+        """
+
+        image = f["Image"][()]
+        pixel_sizes = [
+            float(f[f"Metadata/DimensionScale{dim.upper()}"][()])
+            for dim in list("zyx")
+        ]
+        if "PixelCoordinates" in f["Metadata"]:
+            pixel_coordinates = f["Metadata/PixelCoordinates"][()]
+        else:
+            pixel_coordinates = None
 
-            metadata = dict(f["Metadata"].attrs)
+        metadata = dict(f["Metadata"].attrs)
 
-        im = cls(image=image, pixel_sizes=pixel_sizes, metadata=metadata)
+        im = Image(image=image, pixel_sizes=pixel_sizes, metadata=metadata)
         if pixel_coordinates is not None:
             im.pixel_coordinates = pixel_coordinates
         return im
@@ -571,10 +592,15 @@ def __init__(self, filename: str | Path) -> None:
                 root_element = root_elements[0]
             else:
                 # find root element with ImageData
+                root_element = None
                 for re in root_elements:
                     if "ImageData" in f[re].keys():
                         root_element = re
                         break
+                if root_element is None:
+                    raise ValueError(
+                        "No ImageData found in the file, cannot load image"
+                    )
             image = np.squeeze(f[root_element + "/ImageData/Image"][()])
             pixel_sizes = [
                 float(f[root_element + f"/ImageData/DimensionScale{dim}"][()])
diff --git a/tests/test_image.py b/tests/test_image.py
index 6c429be..c971da5 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from simulatedmicroscopy import HuygensImage, Image
+from simulatedmicroscopy import HuygensImage, Image, HuygensPSF
 
 
 def create_demo_image():
@@ -40,6 +40,25 @@ def test_huygens_notexisting(tmp_path):
     with pytest.raises(FileNotFoundError):
         HuygensImage(tmp_path / "thisfiledoesnotexist.h5")
 
+def test_huygens_loading(tmp_path):
+    import h5py
+    im = create_demo_image()
+    pixel_sizes = im.get_pixel_sizes()
+    filename = tmp_path / "testfile.hdf5"
+    with h5py.File(filename, "w") as f:
+        root = f.create_group("testfile")
+        root.create_dataset("ImageData/Image", data=im.image)
+        [root.create_dataset(f"ImageData/DimensionScale{dim}", data=pixel_sizes[i]) for i,dim in enumerate(list("ZYX"))]
+    
+    
+    im_loaded = im.load_h5file(filename)
+    im_loaded_huygens_im = HuygensImage(filename)
+    im_loaded_huygens_psf = HuygensPSF(filename)
+
+
+    assert im_loaded == im
+    assert im_loaded_huygens_im == im
+    assert im_loaded_huygens_psf == im
 
 @pytest.mark.parametrize(
     "unit,conversionfactor",