diff --git a/src/vstarstack/library/calibration/psf.py b/src/vstarstack/library/calibration/psf.py new file mode 100644 index 0000000..93015af --- /dev/null +++ b/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 . +# + +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 diff --git a/src/vstarstack/library/data.py b/src/vstarstack/library/data.py index 8488a61..f5648c4 100644 --- a/src/vstarstack/library/data.py +++ b/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]: diff --git a/src/vstarstack/tool/calibration.py b/src/vstarstack/tool/calibration.py index 29dd199..f96f895 100644 --- a/src/vstarstack/tool/calibration.py +++ b/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%]") } diff --git a/src/vstarstack/tool/image_processing/deconvolution.py b/src/vstarstack/tool/image_processing/deconvolution.py index 6250042..8ead948 100644 --- a/src/vstarstack/tool/image_processing/deconvolution.py +++ b/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/ "), + "wiener": (lambda project, argv : _process(project, argv, "Wiener"), "Wiener deconvolution", "inputs/ psf.zip outputs/ "), +} diff --git a/src/vstarstack/tool/image_processing/fixes.py b/src/vstarstack/tool/image_processing/fixes.py index cadd0ea..0373b06 100644 --- a/src/vstarstack/tool/image_processing/fixes.py +++ b/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"), }