From 6644b3c99d142fbc260980d359c733ca1ed226c4 Mon Sep 17 00:00:00 2001
From: Jakub Both <jakub.both@uib.no>
Date: Tue, 23 Jul 2024 20:38:19 +0200
Subject: [PATCH] ENH: Add general loading routine for corrections

---
 src/darsia/__init__.py                   |  1 +
 src/darsia/corrections/readcorrection.py | 41 ++++++++++++++++++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 src/darsia/corrections/readcorrection.py

diff --git a/src/darsia/__init__.py b/src/darsia/__init__.py
index 20edc17b..d2eeeb1a 100644
--- a/src/darsia/__init__.py
+++ b/src/darsia/__init__.py
@@ -57,6 +57,7 @@
 from darsia.corrections.color.colorcorrection import *
 from darsia.corrections.color.relativecolorcorrection import *
 from darsia.corrections.color.experimentalcolorcorrection import *
+from darsia.corrections.readcorrection import *
 from darsia.image.coordinatetransformation import *  # Requires affine correction
 from darsia.signals.models.basemodel import *
 from darsia.signals.models.combinedmodel import *
diff --git a/src/darsia/corrections/readcorrection.py b/src/darsia/corrections/readcorrection.py
new file mode 100644
index 00000000..f3d0b995
--- /dev/null
+++ b/src/darsia/corrections/readcorrection.py
@@ -0,0 +1,41 @@
+"""Module containing common reading functionality for corrections."""
+
+from pathlib import Path
+
+import numpy as np
+
+from darsia import (
+    ColorCorrection,
+    CurvatureCorrection,
+    DriftCorrection,
+    IlluminationCorrection,
+    TypeCorrection,
+)
+
+AnyCorrection = (
+    TypeCorrection
+    | DriftCorrection
+    | CurvatureCorrection
+    | IlluminationCorrection
+    | ColorCorrection
+)
+
+
+def read_correction(path: Path) -> AnyCorrection:
+    """General function to read corrections from a file.
+
+    Args:
+        path (Path): path to npz file
+
+    Returns:
+        BaseCorrection: correction object
+
+    """
+
+    # Read class name from npz file
+    class_name = np.load(path, allow_pickle=True)["class_name"].item()
+
+    # Load correction from file
+    correction = eval(class_name)()
+    correction.load(path)
+    return correction