Build PSF from images

src/vstarstack/library/calibration/psf.py

@@ -0,0 +1,37 @@
+#
+# Copyright (c) 2024 Vladislav Tsendrovskii
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+# See the GNU General Public License for more details.
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+#
+
+import numpy as np
+import vstarstack.library.data
+import vstarstack.library.merge.simple_add
+import vstarstack.library.image_process.remove_sky
+import vstarstack.library.sky_model.gradient
+import vstarstack.library.common
+
+def prepare_psf(images : vstarstack.library.common.IImageSource, threshold : float) -> vstarstack.library.data.DataFrame:
+    """Prepare PSF from star image"""
+    psf = vstarstack.library.merge.simple_add.simple_add(images)
+    psf = vstarstack.library.image_process.remove_sky.remove_sky_with_model(psf, vstarstack.library.sky_model.gradient.model)
+    for channel in list(psf.get_channels()):
+        if psf.get_channel_option(channel, "weight"):
+            psf.remove_channel(channel)
+            continue
+        image, opts = psf.get_channel(channel)
+        image -= np.amax(image) * threshold
+        image = np.clip(image, 0, None)
+        sumv = np.sum(image)
+        if sumv > 0:
+            image = image / sumv
+        psf.replace_channel(image, channel, **opts)
+    return psf

src/vstarstack/library/data.py

@@ -157,6 +157,8 @@ def get_channel(self, channel : str) -> Tuple[np.ndarray, dict]:
         Returns:
             tuple(image, options) - ndarray with pixel data and channel options
         """
+        if channel not in self.channels:
+            return None, None
         return self.channels[channel]["data"], self.channels[channel]["options"]
 
     def get_channels(self) -> List[str]:

src/vstarstack/tool/calibration.py

@@ -15,6 +15,7 @@
 import multiprocessing as mp
 import json
 
+import vstarstack.library.calibration.psf
 import vstarstack.tool.common
 import vstarstack.tool.cfg
 import vstarstack.tool.usage
@@ -232,6 +233,20 @@ def _process_approximate_flat(project : vstarstack.tool.cfg.Project,
     df = vstarstack.library.calibration.flat.approximate_flat_image(df)
     df.store(flat_fname)
 
+def _process_build_psf(project : vstarstack.tool.cfg.Project,
+                       argv : list[str]):
+    input_path = argv[0]
+    psf_fname = argv[1]
+    if len(argv) >= 3:
+        threshold = float(argv[2]) / 100.0
+    else:
+        threshold = 0
+    files = vstarstack.tool.common.listfiles(input_path, ".zip")
+    psfs = [item[1] for item in files]
+    src = vstarstack.library.common.FilesImageSource(psfs)
+    psf = vstarstack.library.calibration.psf.prepare_psf(src, threshold)
+    psf.store(psf_fname)
+
 commands = {
     "flatten": (_process_flatten,
                 "Flatten image",
@@ -253,5 +268,8 @@ def _process_approximate_flat(project : vstarstack.tool.cfg.Project,
                            "flats/ flat.zip"),
     "approximate-flat" : (_process_approximate_flat,
                           "Approximate flat with polynomic function",
-                          "original_flat.zip result_flat.zip")
+                          "original_flat.zip result_flat.zip"),
+    "build-psf" : (_process_build_psf,
+                   "Create point spread function for deconvolution",
+                   "star_images/ psf.zip [threshold%]")
 }

src/vstarstack/tool/image_processing/deconvolution.py

@@ -21,57 +21,61 @@
 import vstarstack.tool.cfg
 import vstarstack.library.data
 
-def deconvolution(df : vstarstack.library.data.DataFrame, strength : int):
-    psf = np.zeros((3, 3))
-
-    psf[1, 1] = 1
-
-    psf[0, 1] = 1/strength
-    psf[2, 1] = 1/strength
-    psf[1, 0] = 1/strength
-    psf[1, 2] = 1/strength
-
-    psf[0, 0] = 1/strength/1.41
-    psf[0, 2] = 1/strength/1.41
-    psf[2, 0] = 1/strength/1.41
-    psf[2, 0] = 1/strength/1.41
-
-    psf = psf / np.sum(psf)
-
+def deconvolution(df : vstarstack.library.data.DataFrame, psf_df : vstarstack.library.data.DataFrame, argv : list, method : str):
     for channel in df.get_channels():
         image, opts = df.get_channel(channel)
         if not opts["brightness"]:
             continue
+        psf,_ = psf_df.get_channel(channel)
+        if psf is None:
+            continue
         norm = np.amax(image)
-        deconvolved_RL = skimage.restoration.richardson_lucy(image/norm, psf, num_iter=100)*norm
-        deconvolved_RL[np.where(np.isnan(deconvolved_RL))] = 0
-        df.replace_channel(deconvolved_RL, channel)
+        image = image / norm
+        if method == "RL":
+            num_steps = int(argv[0])
+            deconvolved = skimage.restoration.richardson_lucy(image, psf, num_iter=num_steps)
+        elif method == "Wiener":
+            balance = float(argv[0])
+            deconvolved = skimage.restoration.wiener(image, psf, balance)
+        else:
+            raise Exception(f"Unknown method {method}")
+        deconvolved[np.where(np.isnan(deconvolved))] = 0
+        deconvolved = deconvolved * norm
+        df.replace_channel(deconvolved, channel, **opts)
     return df
 
-def _process_file(input : str, output : str, strength : int):
+def _process_file(input : str, psf : vstarstack.library.data.DataFrame, output : str, argv : list, method : str):
     df = vstarstack.library.data.DataFrame.load(input)
-    deconvolution(df, strength)
+    deconvolution(df, psf, argv, method)
+    vstarstack.tool.common.check_dir_exists(output)
     df.store(output)
 
-def _process_dir(inputs : str, outputs : str, strength : int):
+def _process_single_file(input : str, psf_fname : str, output : str, argv : list, method : str):
+    psf = vstarstack.library.data.DataFrame.load(psf_fname)
+    _process_file(input, psf, output, argv, method)
+
+def _process_dir(inputs : str, psf_fname : str, outputs : str, argv : list, method : str):
     files = vstarstack.tool.common.listfiles(inputs, ".zip")
+    psf = vstarstack.library.data.DataFrame.load(psf_fname)
     with mp.Pool(vstarstack.tool.cfg.nthreads) as pool:
-        pool.starmap(_process_file, [(fname,
+        pool.starmap(_process_file, [(fname, psf,
                                       os.path.join(outputs, name + ".zip"),
-                                      strength) for name, fname in files])
+                                      argv,
+                                      method) for name, fname in files])
 
-def run(project : vstarstack.tool.cfg.Project, argv : list[str]):
+def _process(project : vstarstack.tool.cfg.Project, argv : list[str], method : str):
     """Deconvolution"""
-    if len(argv) >= 3:
-        inputs = argv[0]
-        strength = int(argv[1])
-        outputs = argv[2]
-    else:
-        inputs = project.config.paths.light.npy
-        outputs = project.config.paths.light.npy
-        strength = int(argv[0])
+    inputs = argv[0]
+    psf_fname = argv[1]
+    outputs = argv[2]
+    argv = argv[3:]
 
     if os.path.isdir(inputs):
-        _process_dir(inputs, outputs, strength)
+        _process_dir(inputs, psf_fname, outputs, argv, method)
     else:
-        _process_file(inputs, outputs, strength)
+        _process_single_file(inputs, psf_fname, outputs, argv, method)
+
+commands = {
+    "rl": (lambda project, argv : _process(project, argv, "RL"),  "Richardson-Lucy deconvolution", "inputs/ psf.zip outputs/ <num_steps>"),
+    "wiener": (lambda project, argv : _process(project, argv, "Wiener"),  "Wiener deconvolution", "inputs/ psf.zip outputs/ <balance>"),
+}

src/vstarstack/tool/image_processing/fixes.py

@@ -50,6 +50,6 @@ def copy(project: vstarstack.tool.cfg.Project, argv: list):
     "border": (vstarstack.tool.image_processing.border.run,     "remove border"),
     "normalize": (vstarstack.tool.image_processing.normalize.run,  "normalize to weight"),
     "blur": (vstarstack.tool.image_processing.blur.run,  "gaussian blur"),
-    "deconvolution": (vstarstack.tool.image_processing.deconvolution.run,  "RL deconvolution"),
+    "deconvolution": (vstarstack.tool.image_processing.deconvolution.commands,  "deconvolution"),
     "select-sharp" : (vstarstack.tool.image_processing.drop_unsharp.commands, "select sharp images"),
 }