relax rotation angle extraction process

Author: KedoKudo

src/imars3d/backend/dataio/data.py

@@ -125,6 +125,9 @@ class load_data(param.ParameterizedFunction):
 
         Currently, we are using a forgiving reader to load the image where a corrupted file
         will not block reading other data.
+
+        The rotation angles are extracted from the filenames if possible, otherwise from the
+        metadata embedded in the tiff files. If both failed, the angle will be set to None.
     """
 
     #
@@ -544,7 +547,7 @@ def _get_filelist_by_dir(
 def _extract_rotation_angles(
     filelist: List[str],
     metadata_idx: int = 65039,
-) -> np.ndarray:
+) -> Optional[np.ndarray]:
     """
     Extract rotation angles in degrees from filename or metadata.
 
@@ -558,40 +561,106 @@ def _extract_rotation_angles(
     Returns
     -------
         rotation_angles
+            Array of rotation angles if successfully extracted, None otherwise.
     """
     # sanity check
-    if filelist == []:
+    if not filelist:
         logger.error("filelist is [].")
         raise ValueError("filelist cannot be empty list.")
 
-    # extract rotation angles from file names
+    # process one file at a time
+    rotation_angles = []
+    for filename in filelist:
+        file_ext = Path(filename).suffix.lower()
+        angle = None
+        if file_ext == ".tiff":
+            # first, let's try to extract the angle from the filename
+            angle = extract_rotation_angle_from_filename(filename)
+            if angle is None:
+                # if failed, try to extract from metadata
+                angle = extract_rotation_angle_from_tiff_metadata(filename, metadata_idx)
+            if angle is None:
+                # if failed, log a warning and move on
+                logger.warning(f"Failed to extract rotation angle from {filename}.")
+        elif file_ext in (".tif", ".fits"):
+            # for tif and fits, we can only extract from filename as the metadata is not reliable
+            angle = extract_rotation_angle_from_filename(filename)
+            if angle is None:
+                # if failed, log a warning and move on
+                logger.warning(f"Failed to extract rotation angle from {filename}.")
+        else:
+            # if the file type is not supported, raise value error
+            logger.error(f"Unsupported file type: {file_ext}")
+            raise ValueError("Unsupported file type.")
+
+        rotation_angles.append(angle)
+
+    # this means we have a list of None
+    if all(angle is None for angle in rotation_angles):
+        logger.warning("Failed to extract any rotation angles.")
+        return None
+
+    # warn users if some angles are missing
+    if any(angle is None for angle in rotation_angles):
+        logger.warning("Some rotation angles are missing. You will see nan in the rotation angles array.")
+
+    return np.array(rotation_angles, dtype=float)
+
+
+def extract_rotation_angle_from_filename(filename: str) -> Optional[float]:
+    """
+    Extract rotation angle in degrees from filename.
+
+    Parameters
+    ----------
+    filename:
+        Filename to extract rotation angle from.
+
+    Returns
+    -------
+        rotation_angle
+            Rotation angle in degrees if successfully extracted, None otherwise.
+    """
+    # extract rotation angle from file names
     # Note
     # ----
     #   For the following file
     #       20191030_ironman_small_0070_300_440_0520.tif(f)
+    #       20191030_ironman_small_0070_300_440_0520.fits
     #   the rotation angle is 300.44 degrees
-    # If all given filenames follows the pattern, we will use the angles from
-    # filenames. Otherwise, we will use the angles from metadata.
-    regex = r"\d{8}_\S*_\d{4}_(?P<deg>\d{3})_(?P<dec>\d{3})_\d*\.tif{1,2}"
-    matches = [re.match(regex, Path(f).name) for f in filelist]
-    if all(matches):
-        logger.info("Using rotation angles from filenames.")
-        rotation_angles = np.array([float(".".join(m.groups())) for m in matches])
+    regex = r"\d{8}_\S*_\d{4}_(?P<deg>\d{3})_(?P<dec>\d{3})_\d*\.(?:tiff?|fits)"
+    match = re.match(regex, Path(filename).name)
+    if match:
+        rotation_angle = float(".".join(match.groups()))
     else:
-        # extract rotation angles from metadata
-        file_exts = set(Path(f).suffix.lower() for f in filelist)
-        if not file_exts.issubset({".tiff", ".tif"}):
-            logger.error("Only .tiff and .tif files are supported.")
-            raise ValueError("Rotation angle from metadata is only supported for .tiff and .tif files.")
+        rotation_angle = None
+    return rotation_angle
+
+
+def extract_rotation_angle_from_tiff_metadata(filename: str, metadata_idx: int = 65039) -> Optional[float]:
+    """
+    Extract rotation angle in degrees from metadata of a tiff file.
+
+    Parameters
+    ----------
+    filename:
+        Filename to extract rotation angle from.
+    metadata_idx:
+        Index of metadata to extract rotation angle from, default is 65039.
+
+    Returns
+    -------
+        rotation_angle
+            Rotation angle in degrees if successfully extracted, None otherwise.
+    """
+    try:
         # -- read metadata
         # img = tifffile.TiffFile("test_with_metadata_0.tiff")
         # img.pages[0].tags[65039].value
         # >> 'RotationActual:0.579840'
-        rotation_angles = np.array(
-            [float(tifffile.TiffFile(f).pages[0].tags[metadata_idx].value.split(":")[-1]) for f in filelist],
-            dtype="float",
-        )
-    return rotation_angles
+        return float(tifffile.TiffFile(filename).pages[0].tags[metadata_idx].value.split(":")[-1])
+    except Exception:
+        return None
 
 
 def _save_data(filename: Path, data: np.ndarray, rot_angles: np.ndarray = None) -> None:

tests/unit/backend/dataio/test_data.py

@@ -15,6 +15,8 @@
     load_data,
     save_checkpoint,
     save_data,
+    extract_rotation_angle_from_filename,
+    extract_rotation_angle_from_tiff_metadata,
 )
 
 
@@ -181,6 +183,31 @@ def test_extract_rotation_angles(data_fixture):
     rst = _extract_rotation_angles([metadata_tiff] * 3)
     ref = np.array([0.1, 0.1, 0.1])
     np.testing.assert_array_almost_equal(rst, ref)
+    # case_2: mixed file types
+    rst = _extract_rotation_angles([good_tiff, metadata_tiff, generic_tiff, generic_fits])
+    ref = np.array([10.02, 0.1, np.nan, np.nan])
+    np.testing.assert_array_equal(rst, ref)
+    # case_3: all files without extractable angles
+    rst = _extract_rotation_angles([generic_tiff, generic_fits])
+    assert rst is None
+
+
+def test_extract_rotation_angle_from_filename():
+    # Test cases for extract_rotation_angle_from_filename
+    assert extract_rotation_angle_from_filename("20191030_sample_0070_300_440_0520.tiff") == 300.44
+    assert extract_rotation_angle_from_filename("20191030_sample_0071_301_441_0521.tif") == 301.441
+    assert extract_rotation_angle_from_filename("20191030_sample_0072_302_442_0522.fits") == 302.442
+    assert extract_rotation_angle_from_filename("generic_file.tiff") is None
+
+
+def test_extract_rotation_angle_from_tiff_metadata(tmpdir):
+    # Create a TIFF file with rotation angle in metadata
+    data = np.ones((3, 3))
+    filename = str(tmpdir / "metadata.tiff")
+    tifffile.imwrite(filename, data, extratags=[(65039, "s", 0, "RotationActual:0.5", True)])
+
+    assert extract_rotation_angle_from_tiff_metadata(filename) == 0.5
+    assert extract_rotation_angle_from_tiff_metadata("non_existent_file.tiff") is None
 
 
 @pytest.fixture(scope="module")